File size: 1,412 Bytes
20d49b2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
# 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
|