Update translation.py

translation.py

@@ -1,56 +1,54 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# 1. TEMPORARY FIX: Switched to the smaller MarianMT model to avoid memory crashes.
-# If this works, the NLLB model is too large for the current Space hardware.
-# NLLB MODEL (Large): model_name = "facebook/nllb-200-distilled-600M"
-model_name = "Helsinki-NLP/opus-mt-en-ur"
-_tokenizer = None
-_model = None
-
-# MarianMT models handle the reverse translation (Urdu-English) by using a separate model pair.
-# We will load the reverse model on demand.
-REVERSE_MODEL_NAME = "Helsinki-NLP/opus-mt-ur-en"
-_reverse_tokenizer = None
-_reverse_model = None
-
-
-def _load_translation_resources():
-    """Loads the main EN-UR model resources (Helsinki-NLP/opus-mt-en-ur)."""
-    global _tokenizer, _model
-    if _tokenizer is None or _model is None:
-        _tokenizer = AutoTokenizer.from_pretrained(model_name)
-        _model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    return _tokenizer, _model
-
-def _load_reverse_translation_resources():
-    """Loads the UR-EN model resources (Helsinki-NLP/opus-mt-ur-en)."""
-    global _reverse_tokenizer, _reverse_model
-    if _reverse_tokenizer is None or _reverse_model is None:
-        _reverse_tokenizer = AutoTokenizer.from_pretrained(REVERSE_MODEL_NAME)
-        _reverse_model = AutoModelForSeq2SeqLM.from_pretrained(REVERSE_MODEL_NAME)
-    return _reverse_tokenizer, _reverse_model
+# --- Model Definitions ---
+# Using separate, small MarianMT models to guarantee stability and avoid memory crashes.
+EN_UR_MODEL_NAME = "Helsinki-NLP/opus-mt-en-ur"
+UR_EN_MODEL_NAME = "Helsinki-NLP/opus-mt-ur-en"
 
+# Lazy-loading variables for EN-UR model
+_en_ur_tokenizer = None
+_en_ur_model = None
+
+# Lazy-loading variables for UR-EN model
+_ur_en_tokenizer = None
+_ur_en_model = None
+
+# --- Resource Loading Functions ---
+
+def _load_en_ur_resources():
+    """Loads the English-to-Urdu MarianMT model."""
+    global _en_ur_tokenizer, _en_ur_model
+    if _en_ur_tokenizer is None or _en_ur_model is None:
+        _en_ur_tokenizer = AutoTokenizer.from_pretrained(EN_UR_MODEL_NAME)
+        _en_ur_model = AutoModelForSeq2SeqLM.from_pretrained(EN_UR_MODEL_NAME)
+    return _en_ur_tokenizer, _en_ur_model
+
+def _load_ur_en_resources():
+    """Loads the Urdu-to-English MarianMT model."""
+    global _ur_en_tokenizer, _ur_en_model
+    if _ur_en_tokenizer is None or _ur_en_model is None:
+        _ur_en_tokenizer = AutoTokenizer.from_pretrained(UR_EN_MODEL_NAME)
+        _ur_en_model = AutoModelForSeq2SeqLM.from_pretrained(UR_EN_MODEL_NAME)
+    return _ur_en_tokenizer, _ur_en_model
+
+# --- Translation Functions ---
 
 def translate_to_urdu(text):
-    """Translates English text to Urdu using the Helsinki-NLP/opus-mt-en-ur model."""
-    # MarianMT models often don't require explicit src_lang or forced_bos_token_id 
-    # for single-pair models, but we use the target language code for safety.
-    tokenizer, model = _load_translation_resources()
+    """Translates English text to Urdu."""
+    tokenizer, model = _load_en_ur_resources()
 
     try:
-        # NOTE: MarianMT tokens are often '>>ur<<' for the target language.
-        input_ids = tokenizer(
-            text,
-            # For MarianMT, the source language is implicit in the model name (en-ur)
-            # but we use the target language token to guide generation.
-            text_target=[""] * len(text) if isinstance(text, list) else "",
-            return_tensors='pt'
-        ).input_ids
-
-        # We set the forced_bos_token_id to the target language code 'ur'
+        # MarianMT requires the target language token to start the generation
+        # We use '>>ur<<' as the start token for this model pair.
+        TGT_LANG_TOKEN = '>>ur<<' 
+        
+        input_ids = tokenizer.encode(text, return_tensors='pt')
+        
         generated_tokens = model.generate(
             input_ids,
-            forced_bos_token_id=tokenizer.get_lang_id("ur"),
+            # CRITICAL FIX: Use the specific language token ID for MarianMT
+            decoder_start_token_id=tokenizer.lang_code_to_id[TGT_LANG_TOKEN],
             num_beams=5,
             max_length=128
         )
@@ -58,26 +56,23 @@ def translate_to_urdu(text):
         return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
     except Exception as exc:
-        # Raised as RuntimeError to be caught in app.py and sent as 500
-        raise RuntimeError("Translation to Urdu failed (MarianMT check)") from exc
-
+        raise RuntimeError("Translation to Urdu failed (MarianMT Final)") from exc
 
 def translate_to_english(text):
-    """Translates Urdu text to English using the Helsinki-NLP/opus-mt-ur-en model."""
-    # Note: This loads a separate UR-EN model pair.
-    tokenizer, model = _load_reverse_translation_resources()
+    """Translates Urdu text to English."""
+    tokenizer, model = _load_ur_en_resources()
 
     try:
-        input_ids = tokenizer(
-            text,
-            text_target=[""] * len(text) if isinstance(text, list) else "",
-            return_tensors='pt'
-        ).input_ids
+        # MarianMT requires the target language token to start the generation
+        # We use '>>en<<' as the start token for this reverse model pair.
+        TGT_LANG_TOKEN = '>>en<<'
+        
+        input_ids = tokenizer.encode(text, return_tensors='pt')
 
-        # We set the forced_bos_token_id to the target language code 'en'
         generated_tokens = model.generate(
             input_ids,
-            forced_bos_token_id=tokenizer.get_lang_id("en"),
+            # CRITICAL FIX: Use the specific language token ID for MarianMT
+            decoder_start_token_id=tokenizer.lang_code_to_id[TGT_LANG_TOKEN],
             num_beams=5,
             max_length=128
         )
@@ -85,19 +80,19 @@ def translate_to_english(text):
         return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
 
     except Exception as exc:
-        raise RuntimeError("Translation to English failed (MarianMT check)") from exc
-
+        raise RuntimeError("Translation to English failed (MarianMT Final)") from exc
 
 # --- Example Usage ---
 if __name__ == "__main__":
-    input_text = "The study investigates the correlation between socioeconomic status and academic achievement."
-    translated_text = translate_to_urdu(input_text)
+    # Test English to Urdu
+    input_text_en = "This is a final test of the translation API."
+    translated_text_ur = translate_to_urdu(input_text_en)
     
-    print(f"Original (English): {input_text}")
-    print(f"Translated (Urdu): {translated_text}")
+    print(f"Original (English): {input_text_en}")
+    print(f"Translated (Urdu): {translated_text_ur}")
     
     # Test Urdu to English translation
-    urdu_text = translated_text
-    back_to_english = translate_to_english(urdu_text)
-    print(f"\nOriginal (Urdu): {urdu_text}")
-    print(f"Translated back (English): {back_to_english}")
+    input_text_ur = "یہ ایپلیکیشن کامیابی سے چل رہی ہے۔"
+    translated_text_en = translate_to_english(input_text_ur)
+    print(f"\nOriginal (Urdu): {input_text_ur}")
+    print(f"Translated back (English): {translated_text_en}")