translate.py

# 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)