english-learning-chatbot / translate.py
SandaAbhishekSagar
revamped code of translate.py
7295e60
# from transformers import MarianMTModel, MarianTokenizer
# def translate_text(text, src_lang="es", tgt_lang="en"):
# model_name = f"Helsinki-NLP/opus-mt-{src_lang}-{tgt_lang}"
# tokenizer = MarianTokenizer.from_pretrained(model_name)
# model = MarianMTModel.from_pretrained(model_name)
# inputs = tokenizer(text, return_tensors="pt", padding=True)
# translated = model.generate(**inputs)
# return tokenizer.decode(translated[0], skip_special_tokens=True)
# if __name__ == "__main__":
# input_text = "¿Cómo estás?"
# print("Translated Text:", translate_text(input_text, src_lang="es", tgt_lang="en"))
import spaces
from transformers import MarianMTModel, MarianTokenizer
# Preload the translation model globally
model_name = "Helsinki-NLP/opus-mt-mul-en"
tokenizer = MarianTokenizer.from_pretrained(model_name)
translation_model = MarianMTModel.from_pretrained(model_name)
@spaces.GPU
def translate_text(text, src_lang="auto", tgt_lang="en"):
"""Translate text from any language to English."""
inputs = tokenizer(text, return_tensors="pt", padding=True)
translated = translation_model.generate(**inputs)
return tokenizer.decode(translated[0], skip_special_tokens=True)