Update app.py

app.py

@@ -15,18 +15,6 @@ import os
 # Initialize Botok
 wt = WordTokenizer()
 
-# Load TibetaMind (Tibetan → Chinese)
-tibetamind_tokenizer = AutoTokenizer.from_pretrained("tibetamind/tibetan-chinese")
-tibetamind_model = AutoModelForSeq2SeqLM.from_pretrained("tibetamind/tibetan-chinese")
-
-# MBART for Chinese → English
-translation_tokenizer = MBart50TokenizerFast.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt", use_fast=False
-)
-translation_model = MBartForConditionalGeneration.from_pretrained(
-    "facebook/mbart-large-50-many-to-many-mmt"
-)
-
 HF_TOKEN = os.getenv("HF_TOKEN")
 
 try:
@@ -91,6 +79,11 @@ def translate_with_quota(text, src_lang="bo", tgt_lang="en"):
 tts_tibetan = pipeline("text-to-speech", model="facebook/mms-tts-bod")
 #tts_sanskrit = pipeline("text-to-speech", model="facebook/mms-tts-san")
 
+# Load MBART-50
+tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt", use_fast=False)
+model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+
+
 # Use the slow tokenizer to avoid the bug
 ##translation_tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt",use_fast=False)
 #translation_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")
@@ -161,22 +154,17 @@ def run_task(text, language, task):
             # 1) Segment Tibetan text with Botok
             tokens = [t.text for t in wt.tokenize(text)]
             segmented_text = " ".join(tokens)
-    
-            # 2) Tibetan → Chinese via TibetaMind
-            inputs = tibetamind_tokenizer(segmented_text, return_tensors="pt")
-            outputs = tibetamind_model.generate(**inputs, max_new_tokens=256)
-            chinese_text = tibetamind_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-            # 3) Chinese → English via MBART
-            translation_tokenizer.src_lang = "zh_CN"
-            forced_bos = translation_tokenizer.lang_code_to_id["en_XX"]
-            inputs = translation_tokenizer(chinese_text, return_tensors="pt")
-            outputs = translation_model.generate(
-                **inputs,
-                max_new_tokens=256,
-                forced_bos_token_id=forced_bos
-            )
-            english_text = translation_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            # Set source and target languages
+            tokenizer.src_lang = "bo_CN"
+            forced_bos = tokenizer.lang_code_to_id["en_XX"]
+
+            # Translate
+            inputs = tokenizer(text, return_tensors="pt")
+            outputs = model.generate(**inputs, max_new_tokens=256, forced_bos_token_id=forced_bos)
+            english_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            print(english_text)
     
             return None, None, english_text
             #translated_text = translate_with_quota(text, src_lang="bo", tgt_lang="en")
@@ -195,6 +183,9 @@ def run_task(text, language, task):
     else:
         return None, None, "Unsupported task"
 
+
+
+
     
 # --- Gradio Interface ---
 iface = gr.Interface(