| # Imports | |
| from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
| import fasttext | |
| # Model (Pipeline) class | |
| class TranslateFromAny2XModel: | |
| def __init__(self, nllb_model_path: str, fasttext_model_path: str, target_language="eng_Latn"): | |
| """Initialize the model with paths for NLLB and FastText NLLB LID models.""" | |
| self.nllb_model_path = nllb_model_path | |
| self.fasttext_model_path = fasttext_model_path | |
| self.target_language = target_language | |
| # Load NLLB model and tokenizer | |
| self.model = AutoModelForSeq2SeqLM.from_pretrained(nllb_model_path) | |
| self.tokenizer = AutoTokenizer.from_pretrained(nllb_model_path) | |
| # Load FastText language identification model | |
| self.fasttext_model = fasttext.load_model(fasttext_model_path) | |
| def generate(self, prompt: str) -> str: | |
| """Translates the input prompt to target_language using the NLLB model and source language detection using fastText LID model.""" | |
| self.tokenizer.src_lang = self.fasttext_model.predict(prompt)[0][0].replace("__label__", "") | |
| inputs = self.tokenizer(prompt, return_tensors="pt") | |
| output_tokens = self.model.generate(**inputs, forced_bos_token_id=self.tokenizer.convert_tokens_to_ids(self.target_language))[0] | |
| output = self.tokenizer.decode(output_tokens, skip_special_tokens=True) | |
| return output | |