from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig, AutoModelForCausalLM import google.generativeai as genai import torch, os, sys, ast from utils import standardize_lang from functools import wraps from batching import generate_text, Gemini from logging_config import logger from multiprocessing import Process,Event # root dir sys.path.insert(0, os.path.dirname(os.path.dirname(__file__))) from config import LOCAL_FILES_ONLY, available_langs, curr_models, BATCH_SIZE, device, GEMINI_KEY, MAX_INPUT_TOKENS, MAX_OUTPUT_TOKENS, seq_llm_models, api_llm_models, causal_llm_models ############################## # translation decorator def translate(translation_func): @wraps(translation_func) def wrapper(text, *args, **kwargs): try: if len(text) == 0: return [] return translation_func(text, *args, **kwargs) except Exception as e: logger.error(f"Translation error with the following function: {translation_func.__name__}. Text: {text}\nError: {e}") return wrapper ############################### ############################### def init_GEMINI(models_and_rates = None): if not models_and_rates: ## this is default for free tier models_and_rates = {'gemini-1.5-pro': 2, 'gemini-1.5-flash': 15, 'gemini-1.5-flash-8b': 8, 'gemini-1.0-pro': 15} # order from most pref to least pref models = [Gemini(name, rate) for name, rate in models_and_rates.items()] for model in models: model.start() genai.configure(api_key=GEMINI_KEY) return models def translate_GEMINI(text, models, from_lang, target_lang): safety_settings = { "HARM_CATEGORY_HARASSMENT": "BLOCK_NONE", "HARM_CATEGORY_HATE_SPEECH": "BLOCK_NONE", "HARM_CATEGORY_SEXUALLY_EXPLICIT": "BLOCK_NONE", "HARM_CATEGORY_DANGEROUS_CONTENT": "BLOCK_NONE"} prompt = f"Without any additional remarks, and without any code, translate the following items of the Python list from {from_lang} into {target_lang} and output as a Python list ensuring proper escaping of characters: {text}" for model in models: if model.curr_calls < model.rate: try: response = genai.GenerativeModel(model.name).generate_content(prompt, safety_settings=safety_settings) model.curr_calls += 1 logger.info(repr(model)) logger.info(f'Model Response: {response.text.strip()}') return ast.literal_eval(response.text.strip()) except Exception as e: logger.error(f"Error with model {model.name}. Error: {e}") logger.error("No models available to translate. Please wait for a model to be available.") ############################### # Best model by far. Aya-23-8B. Gemma is relatively good. If I get the time to quantize either gemma or aya those will be good to use. llama3.2 is really good as well. def init_AYA(): model_id = "CohereForAI/aya-23-8B" tokenizer = AutoTokenizer.from_pretrained(model_id, local_files_only=True) model = AutoModelForCausalLM.from_pretrained(model_id, locals_files_only=True, torch_dtype=torch.float16).to(device) model.eval() return (model, tokenizer) ############################## # M2M100 model def init_M2M(): tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY) model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M", local_files_only=LOCAL_FILES_ONLY).to(device) model.eval() return (model, tokenizer) def translate_M2M(text, model, tokenizer, from_lang = 'ch_sim', target_lang = 'en') -> list[str]: model_lang_from = standardize_lang(from_lang)['translation_model_lang'] model_lang_to = standardize_lang(target_lang)['translation_model_lang'] if len(text) == 0: return [] tokenizer.src_lang = model_lang_from generated_translations = generate_text(text, model,tokenizer, batch_size=BATCH_SIZE, max_length=MAX_INPUT_TOKENS, max_new_tokens=MAX_OUTPUT_TOKENS, forced_bos_token_id=tokenizer.get_lang_id(model_lang_to)) return generated_translations ############################### ############################### # Helsinki-NLP model Opus MT # Refer here for all the models https://huggingface.co/Helsinki-NLP def get_OPUS_model(from_lang, target_lang): model_lang_from = standardize_lang(from_lang)['translation_model_lang'] model_lang_to = standardize_lang(target_lang)['translation_model_lang'] return f"Helsinki-NLP/opus-mt-{model_lang_from}-{model_lang_to}" def init_OPUS(from_lang = 'ch_sim', target_lang = 'en'): opus_model = get_OPUS_model(from_lang, target_lang) tokenizer = AutoTokenizer.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY) model = AutoModelForSeq2SeqLM.from_pretrained(opus_model, local_files_only=LOCAL_FILES_ONLY, torch_dtype=torch.float16).to(device) model.eval() return (model, tokenizer) def translate_OPUS(text: list[str], model, tokenizer) -> list[str]: translated_text = generate_text(model,tokenizer, text, batch_size=BATCH_SIZE, device=device, max_length=MAX_INPUT_TOKENS, max_new_tokens=MAX_OUTPUT_TOKENS) return translated_text ############################### def init_Seq_LLM(model_type, **kwargs): # model = 'opus' or 'm2m' if model_type == 'opus': return init_OPUS(**kwargs) elif model_type == 'm2m': return init_M2M() else: raise ValueError(f"Invalid model. Please use {' or '.join(curr_models)}.") def init_API_LLM(model_type, **kwargs): # model = 'gemma' if model_type == 'gemini': return init_GEMINI(**kwargs) else: raise ValueError(f"Invalid model type. Please use {' or '.join(api_llm_models)}.") def init_Causal_LLM(model_type, **kwargs): pass ### @translate def translate_Seq_LLM(text, model_type, # 'opus' or 'm2m' model, tokenizer, **kwargs): if model_type == 'opus': return translate_OPUS(text, model, tokenizer) elif model_type == 'm2m': try: return translate_M2M(text, model, tokenizer, **kwargs) except Exception as e: logger.error(f"Error with M2M model. Error: {e}") # raise ValueError(f"Please provide the correct from_lang and target_lang variables if you are using the M2M model. Use the list from {available_langs}.") else: raise ValueError(f"Invalid model. Please use {' or '.join(curr_models)}.") ### if you want to use any other translation, just define a translate function with input text and output text. # def translate_api(text): #@translate #def translate_Causal_LLM(text, model_type, model) @translate def translate_API_LLM(text: list[str], model_type: str, # 'gemma' models: list, # list of objects of classes defined in batching.py from_lang: str, # suggested to use ISO 639-1 codes target_lang: str # suggested to use ISO 639-1 codes ) -> list[str]: if model_type == 'gemini': from_lang = standardize_lang(from_lang)['translation_model_lang'] target_lang = standardize_lang(target_lang)['translation_model_lang'] return translate_GEMINI(text, models, from_lang, target_lang) else: raise ValueError(f"Invalid model. Please use {' or '.join(api_llm_models)}.") @translate def translate_Causal_LLM(text: list[str], model_type, # aya model, tokenizer, from_lang: str, target_lang: str) -> list[str]: model_lang_from = standardize_lang(from_lang)['translation_model_lang'] model_lang_to = standardize_lang(target_lang)['translation_model_lang'] if len(text) == 0: return [] pass # choose between local Seq2Seq LLM or obtain translations from an API def init_func(model): if model in seq_llm_models: return init_Seq_LLM elif model in api_llm_models: return init_API_LLM elif model in causal_llm_models: return init_Causal_LLM else: raise ValueError("Invalid model category. Please use either 'seq' or 'api'.") def translate_func(model): if model in seq_llm_models: return translate_Seq_LLM elif model in api_llm_models: return translate_API_LLM elif model in causal_llm_models: return translate_Causal_LLM else: raise ValueError("Invalid model category. Please use either 'seq' or 'api'.") ### todo: if cuda is not detected, default to online translation as cpu just won't cut it bro. Parallel process it over multiple websites to make it faster if __name__ == "__main__": models = init_GEMINI() print(translate_API_LLM(['想要借用合成台。', '有什么卖的?', '声音太大会有什么影响吗?', '不怕丢东西吗?', '再见。', '布纳马', '想买什么自己拿。把钱留在旁边就好。', '回顾', '隐藏'], 'gemini', models, from_lang='ch_sim', target_lang='en')) # model, tokenizer = init_M2M() # print(translate_Seq_LLM( ['想要借用合成台。', '有什么卖的?', '声音太大会有什么影响吗?', '不怕丢东西吗?', '再见。', '布纳马', '想买什么自己拿。把钱留在旁边就好。', '回顾', '隐藏'], 'm2m', model, tokenizer, from_lang='ch_sim', target_lang='en'))