|
|
|
|
|
|
|
|
from transformers import MarianMTModel, MarianTokenizer
|
|
|
|
|
|
class Translator:
|
|
|
def __init__(self, src_lang="en", tgt_lang="fr"):
|
|
|
model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
|
|
|
self.tokenizer = MarianTokenizer.from_pretrained(model_name)
|
|
|
self.model = MarianMTModel.from_pretrained(model_name)
|
|
|
|
|
|
def translate(self, text):
|
|
|
inputs = self.tokenizer(text, return_tensors="pt", padding=True)
|
|
|
translated = self.model.generate(**inputs)
|
|
|
return self.tokenizer.decode(translated[0], skip_special_tokens=True)
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
translator = Translator(src_lang="en", tgt_lang="es")
|
|
|
text = "Hello, how are you?"
|
|
|
print(f"Original: {text}")
|
|
|
print(f"Translated: {translator.translate(text)}") |
