from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM, GPTQConfig, AutoModelForCausalLM import google.generativeai as genai import torch, os, sys, ast, json from utils import standardize_lang from functools import wraps import random import batching from batching import generate_text, Gemini, Groq 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_API_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_API_LLM(from_lang, target_lang): from_lang = standardize_lang(from_lang)['translation_model_lang'] target_lang = standardize_lang(target_lang)['translation_model_lang'] with open('api_models.json', 'r') as f: models_and_rates = json.load(f) models = [] for class_type, class_models in models_and_rates.items(): cls = getattr(batching, class_type) instantiated_objects = [ cls(model, rate) for model, rate in class_models.items()] models.extend(instantiated_objects) for model in models: model.start() model.set_lang(from_lang, target_lang) return models def translate_API_LLM(text, models): random.shuffle(models) for model in models: try: return model.translate(text) except: continue logger.error("All models have failed to translate the text.") ############################### # 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_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_Causal_LLM(text, model_type, model) @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_API_LLM('ja', 'en') print(translate_API_LLM(['こんにちは'], models))