61 lines
2.5 KiB
Python
61 lines
2.5 KiB
Python
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
import folder_paths
|
||
import os
|
||
import torch
|
||
import gc
|
||
|
||
|
||
models_dir = folder_paths.models_dir
|
||
model_path = os.path.join(models_dir, "TTS")
|
||
|
||
LANGUAGES = ["Arabic", "Bengali", "Czech", "German", "English", "Spanish", "Persian", "French", "Hebrew", "Hindi",
|
||
"Indonesian", "Italian", "Japanese", "Khmer", "Korean", "Lao", "Malay", "Burmese", "Dutch", "Polish",
|
||
"Portuguese", "Russian", "Thai", "Tagalog", "Turkish", "Urdu", "Vietnamese", "中文"]
|
||
|
||
MODEL_CACHE = None
|
||
TOKENIZER = None
|
||
class GemmaxRun:
|
||
def __init__(self):
|
||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||
self.model_name = None
|
||
@classmethod
|
||
def INPUT_TYPES(s):
|
||
return {
|
||
"required": {
|
||
"model":(["GemmaX2-28-2B-v0.1", "GemmaX2-28-9B-v0.1", "GemmaX2-28-2B-4bit", "GemmaX2-28-2B-8bit"],{"default": "GemmaX2-28-2B-4bit"}),
|
||
"source_language": (LANGUAGES, {"default": "English"}),
|
||
"target_language": (LANGUAGES, {"default": "中文"}),
|
||
"text": ("STRING", {"forceInput": True}),
|
||
"max_new_tokens": ("INT", {"default": 200, "min": 1,}),
|
||
"unload_model": ("BOOLEAN", {"default": True}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("STRING",)
|
||
RETURN_NAMES = ("translations",)
|
||
FUNCTION = "translate"
|
||
CATEGORY = "🎤MW/MW-gemmax"
|
||
def translate(self, model, source_language, target_language, text, max_new_tokens, unload_model):
|
||
model_id = model_path + "/" + model
|
||
global MODEL_CACHE, TOKENIZER
|
||
if MODEL_CACHE is None or self.model_name != model_id:
|
||
self.model_name = model_id
|
||
MODEL_CACHE = AutoModelForCausalLM.from_pretrained(model_id).eval().to(self.device)
|
||
TOKENIZER = AutoTokenizer.from_pretrained(model_id)
|
||
|
||
text = "将文本从{}翻译成{}：\n\n{}:{}\n\n{}:".format(source_language, target_language, source_language, text, target_language)
|
||
|
||
inputs = TOKENIZER(text, return_tensors="pt").to(self.device)
|
||
outputs = MODEL_CACHE.generate(**inputs, max_new_tokens=max_new_tokens)
|
||
translations = TOKENIZER.decode(outputs[0], skip_special_tokens=True)
|
||
|
||
translations = translations.split(f"\n\n{target_language}:")[-1].strip('"“”[] ')
|
||
|
||
if unload_model:
|
||
TOKENIZER = None
|
||
MODEL_CACHE = None
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
|
||
return (translations,)
|