revamped code of translate.py

translate.py

@@ -12,18 +12,44 @@
 #     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)
+
 import spaces
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 
 # 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)
+model_name = "facebook/m2m100_418M"  # Facebook's multilingual model
+tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+translation_model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+
 @spaces.GPU
 def translate_text(text, src_lang="auto", tgt_lang="en"):
     """Translate text from any language to English."""
+    # Tokenize the input text
     inputs = tokenizer(text, return_tensors="pt", padding=True)
+    
+    # Set the source language and target language for the model
+    # If source language is auto, the model will automatically detect it
+    tokenizer.src_lang = src_lang if src_lang != "auto" else None
+    tokenizer.tgt_lang = tgt_lang
+
+    # Generate the translation
     translated = translation_model.generate(**inputs)
+    
+    # Decode the translated text
     return tokenizer.decode(translated[0], skip_special_tokens=True)
 
 
+