Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -59,30 +59,36 @@ def detect_language(text: str, top_k: int = 1):
 
 print(f"Loading translation model ({config.TRANSLATION_MODEL_NAME})...")
 
-# SIMPLIFIED: Directly load the NLLB model
-translation_tokenizer = AutoTokenizer.from_pretrained(config.TRANSLATION_MODEL_NAME)
-translation_model = AutoModelForSeq2SeqLM.from_pretrained(
-    config.TRANSLATION_MODEL_NAME,
-    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-    device_map="auto" if DEVICE == "cuda" else None
-)
-
-print(f"✓ Translation model loaded: {config.TRANSLATION_MODEL_NAME}")
-
-# Verify language codes are available
-if hasattr(translation_tokenizer, 'lang_code_to_id'):
-    print(f"Available language codes in tokenizer:")
-    # Show languages we care about
-    target_langs = ["eng_Latn", "ibo_Latn", "yor_Latn", "hau_Latn", "swa_Latn", "amh_Ethi"]
-    for lang in target_langs:
-        if lang in translation_tokenizer.lang_code_to_id:
-            print(f"  ✓ {lang}")
-        else:
-            print(f"  ✗ {lang} (not found)")
-else:
-    print("Warning: Tokenizer doesn't have lang_code_to_id attribute")
+NLLB_MODEL = "facebook/nllb-200-distilled-600M"
+print(f"Using model: {NLLB_MODEL}")
 
-# Correct language code mapping for NLLB
+try:
+    translation_tokenizer = AutoTokenizer.from_pretrained(NLLB_MODEL)
+    translation_model = AutoModelForSeq2SeqLM.from_pretrained(
+        NLLB_MODEL,
+        torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+        device_map="auto" if DEVICE == "cuda" else None
+    )
+    
+    print(f"✓ Translation model loaded successfully")
+    
+    # DEBUG: Check tokenizer properties
+    print(f"Tokenizer type: {type(translation_tokenizer).__name__}")
+    print(f"Has lang_code_to_id: {hasattr(translation_tokenizer, 'lang_code_to_id')}")
+    
+    if hasattr(translation_tokenizer, 'lang_code_to_id'):
+        print(f"Sample language codes: {list(translation_tokenizer.lang_code_to_id.keys())[:10]}")
+    else:
+        
+        from transformers import AutoConfig
+        config_model = AutoConfig.from_pretrained(NLLB_MODEL)
+        print(f"Model config: {config_model}")
+        
+except Exception as e:
+    print(f"✗ Error loading translation model: {e}")
+    raise
+
+# Language code mapping 
 LANG_CODE_MAP = {
     "eng_Latn": "eng_Latn",  # English
     "ibo_Latn": "ibo_Latn",  # Igbo
@@ -92,6 +98,16 @@ LANG_CODE_MAP = {
     "amh_Latn": "amh_Ethi",  # Amharic
 }
 
+# Alternative mapping k
+LANG_CODE_MAP_ALT = {
+    "eng_Latn": "en",  # English
+    "ibo_Latn": "ig",  # Igbo
+    "yor_Latn": "yo",  # Yoruba
+    "hau_Latn": "ha",  # Hausa
+    "swh_Latn": "sw",  # Swahili
+    "amh_Latn": "am",  # Amharic
+}
+
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
     "ibo_Latn": "Igbo",
@@ -129,43 +145,66 @@ def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int =
         print("  No translation needed (same language)")
         return text
     
+  
     src_code = LANG_CODE_MAP.get(src_lang, "eng_Latn")
     tgt_code = LANG_CODE_MAP.get(tgt_lang, "eng_Latn")
     
     print(f"  Using codes: {src_code} → {tgt_code}")
     
-    # Check if codes are available
+    
     if not hasattr(translation_tokenizer, 'lang_code_to_id'):
-        print("  ERROR: Tokenizer doesn't support language codes")
-        return text
+        print("  WARNING: Tokenizer doesn't have lang_code_to_id")
+        print("  Trying alternative method...")
+        
+       
+        src_code_alt = LANG_CODE_MAP_ALT.get(src_lang, "en")
+        tgt_code_alt = LANG_CODE_MAP_ALT.get(tgt_lang, "en")
+        
+        
+        try:
+            from transformers import pipeline
+            translator = pipeline(
+                "translation",
+                model=translation_model,
+                tokenizer=translation_tokenizer,
+                src_lang=src_code_alt,
+                tgt_lang=tgt_code_alt,
+                device=0 if DEVICE == "cuda" else -1,
+                max_length=400
+            )
+            
+            chunks = chunk_text(text, max_len=max_chunk_len)
+            translated_parts = []
+            
+            for chunk in chunks:
+                result = translator(chunk)
+                translated_parts.append(result[0]["translation_text"])
+            
+            return " ".join(translated_parts).strip()
+            
+        except Exception as e:
+            print(f"  Pipeline translation failed: {e}")
+            return text 
+    
     
     if src_code not in translation_tokenizer.lang_code_to_id:
-        print(f"  WARNING: Source code {src_code} not found, using English")
-        src_code = "eng_Latn"
+        print(f"  WARNING: Source code {src_code} not found, trying alternatives...")
+  
+        src_code = LANG_CODE_MAP_ALT.get(src_lang, "eng_Latn")
     
     if tgt_code not in translation_tokenizer.lang_code_to_id:
-        print(f"  WARNING: Target code {tgt_code} not found, using English")
-        tgt_code = "eng_Latn"
+        print(f"  WARNING: Target code {tgt_code} not found, trying alternatives...")
+        tgt_code = LANG_CODE_MAP_ALT.get(tgt_lang, "eng_Latn")
     
-    # Handle non-English to non-English translation
-    if src_code != "eng_Latn" and tgt_code != "eng_Latn":
-        print(f"  Two-step translation: {src_code}→eng_Latn→{tgt_code}")
-        to_english = translate_direct(text, src_code, "eng_Latn", max_chunk_len)
-        return translate_direct(to_english, "eng_Latn", tgt_code, max_chunk_len)
     
-    return translate_direct(text, src_code, tgt_code, max_chunk_len)
-
-def translate_direct(text: str, src_code: str, tgt_code: str, max_chunk_len: int = 400) -> str:
-    """Direct translation using forced_bos_token_id"""
-    chunks = chunk_text(text, max_len=max_chunk_len)
-    translated_parts = []
-    
-    # Set source language
     translation_tokenizer.src_lang = src_code
     
-    # Get forced bos token ID
+   
     forced_bos_token_id = translation_tokenizer.lang_code_to_id[tgt_code]
     
+    chunks = chunk_text(text, max_len=max_chunk_len)
+    translated_parts = []
+    
     for i, chunk in enumerate(chunks):
         try:
             inputs = translation_tokenizer(
@@ -200,6 +239,7 @@ def translate_direct(text: str, src_code: str, tgt_code: str, max_chunk_len: int
     
     return " ".join(translated_parts).strip()
 
+
 def retrieve_docs(query: str, vs_path: str):
     if not vs_path or not os.path.exists(vs_path):
         return None