Update app.py

app.py

@@ -1,13 +1,13 @@
-from transformers import NllbTokenizer, NllbForConditionalGeneration
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
 import gradio as gr
 import torch
 
 # Define the model
-model_name = "facebook/nllb-200-distilled-600M"
+model_name = "facebook/m2m100_418M"
 
 try:
-    tokenizer = NllbTokenizer.from_pretrained(model_name)
-    model = NllbForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
 except Exception as e:
     print(f"Error loading model or tokenizer: {e}")
     exit(1)
@@ -81,7 +81,7 @@ def transliterate_farsi_to_cyrillic(farsi_text):
         "سلام": "Салом", 
         "خداحافظ": "Худоҳафиз", 
         "شب بخیر": "Шаби хайр",
-        "صبح بخیر": "Субҳи хайر", 
+        "صبح بخیر": "Субҳи хайр", 
         "ممنون": "Ташаккур", 
         "خواهش میکنم": "Илтимос",
         "چطور هستی؟": "Чӣ хел ҳастӣ?", 
@@ -92,7 +92,7 @@ def transliterate_farsi_to_cyrillic(farsi_text):
         "تشکر": "Ташаккур",
         "فارسی": "Форсī",
         "اسم من": "Номи ман",
-        "لطفا": "Илтимос",
+        "لطفا": "Илтимوس",
         "کمک": "Кумак",
         "هستی": "ҳастī",
         "هستید": "ҳастед",
@@ -158,9 +158,9 @@ def translate_to_cyrillic_farsi(text):
         return "Error: Please enter text in English (ASCII characters only).", ""
 
     # Try full sentence translation first
-    tokenizer.src_lang = "eng_Latn"  # English in Latin script
+    tokenizer.src_lang = "en"
     encoded_text = tokenizer(text, return_tensors="pt", padding=True).to(device)
-    translated = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id("fas_Arab"))  # Farsi in Arabic script
+    translated = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id("fa"))
     farsi_text = tokenizer.decode(translated[0], skip_special_tokens=True)
 
     # Check if the translation is valid Farsi
@@ -172,9 +172,9 @@ def translate_to_cyrillic_farsi(text):
             if phrase in common_phrases:
                 farsi_translations.append(common_phrases[phrase])
             else:
-                tokenizer.src_lang = "eng_Latn"
+                tokenizer.src_lang = "en"
                 encoded_text = tokenizer(phrase, return_tensors="pt", padding=True).to(device)
-                translated = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id("fas_Arab"))
+                translated = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id("fa"))
                 translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
                 if any(c in "ابتثجحخدذرزسشصضطظعغفقکگلمنوهیءأؤئء،؟" for c in translated_text.replace(" ", "")):
                     farsi_translations.append(translated_text)