Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -59,13 +59,37 @@ 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")
+
+# Correct language code mapping for NLLB
 LANG_CODE_MAP = {
-    "eng_Latn": "eng_Latn",
-    "ibo_Latn": "ibo_Latn",
-    "yor_Latn": "yor_Latn",
-    "hau_Latn": "hau_Latn",
-    "swh_Latn": "swa_Latn",
-    "amh_Latn": "amh_Ethi",
+    "eng_Latn": "eng_Latn",  # English
+    "ibo_Latn": "ibo_Latn",  # Igbo
+    "yor_Latn": "yor_Latn",  # Yoruba
+    "hau_Latn": "hau_Latn",  # Hausa
+    "swh_Latn": "swa_Latn",  # Swahili
+    "amh_Latn": "amh_Ethi",  # Amharic
 }
 
 SUPPORTED_LANGS = {
@@ -97,96 +121,77 @@ def chunk_text(text: str, max_len: int = 400) -> List[str]:
         chunks.append(current.strip())
     return chunks
 
-def load_translation_model():
-    """Load translation model with proper configuration"""
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(config.TRANSLATION_MODEL_NAME)
-        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("✓ Custom translation model loaded")
-        return tokenizer, model
-    except Exception as e:
-        print(f"✗ Error loading custom model: {e}")
-        print("Loading standard NLLB model as fallback...")
-        tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
-        model = AutoModelForSeq2SeqLM.from_pretrained(
-            "facebook/nllb-200-distilled-600M",
-            torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            device_map="auto" if DEVICE == "cuda" else None
-        )
-        print("✓ Standard NLLB model loaded as fallback")
-        return tokenizer, model
-
-# Load the model
-translation_tokenizer, translation_model = load_translation_model()
-
-def translate_with_nllb(text: str, src_code: str, tgt_code: str, max_chunk_len: int = 400) -> str:
-    """Translate using NLLB model with forced_bos_token_id"""
-    chunks = chunk_text(text, max_len=max_chunk_len)
-    translated_parts = []
+def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
+    print(f"\n[TRANSLATION] {src_lang} → {tgt_lang}")
+    print(f"  Input: {text[:100]}...")
     
-    # Check if tokenizer has lang_code_to_id
-    if hasattr(translation_tokenizer, 'lang_code_to_id'):
-        try:
-            # 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]
-            
-            for i, chunk in enumerate(chunks):
-                try:
-                    inputs = translation_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512)
-                    
-                    if DEVICE == "cuda":
-                        inputs = {k: v.to(translation_model.device) for k, v in inputs.items()}
-                    
-                    generated_tokens = translation_model.generate(
-                        **inputs,
-                        forced_bos_token_id=forced_bos_token_id,
-                        max_new_tokens=400,
-                        num_beams=4,
-                        early_stopping=True
-                    )
-                    
-                    result = translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
-                    translated_parts.append(result.strip())
-                    
-                except Exception as e:
-                    print(f"  Chunk {i+1} error: {e}")
-                    translated_parts.append(chunk)
-        except Exception as e:
-            print(f"  Language code error: {e}")
-            # Fallback to simple translation
-            return translate_simple(text, max_chunk_len)
-    else:
-        # If no lang_code_to_id, try simple translation
-        return translate_simple(text, max_chunk_len)
+    if not text.strip() or src_lang == tgt_lang:
+        print("  No translation needed (same language)")
+        return text
     
-    return " ".join(translated_parts).strip()
+    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
+    
+    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"
+    
+    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"
+    
+    # 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_simple(text: str, max_chunk_len: int = 400) -> str:
-    """Simple translation without language codes"""
+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]
+    
     for i, chunk in enumerate(chunks):
         try:
-            inputs = translation_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512)
+            inputs = translation_tokenizer(
+                chunk, 
+                return_tensors="pt", 
+                truncation=True, 
+                max_length=512
+            )
             
             if DEVICE == "cuda":
                 inputs = {k: v.to(translation_model.device) for k, v in inputs.items()}
             
             generated_tokens = translation_model.generate(
                 **inputs,
+                forced_bos_token_id=forced_bos_token_id,
                 max_new_tokens=400,
                 num_beams=4,
                 early_stopping=True
             )
             
-            result = translation_tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+            result = translation_tokenizer.batch_decode(
+                generated_tokens, 
+                skip_special_tokens=True
+            )[0]
+            
+            print(f"  Chunk {i+1}: '{chunk[:30]}...' → '{result[:30]}...'")
             translated_parts.append(result.strip())
             
         except Exception as e:
@@ -195,26 +200,6 @@ def translate_simple(text: str, max_chunk_len: int = 400) -> str:
     
     return " ".join(translated_parts).strip()
 
-def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
-    print(f"\n[TRANSLATION] {src_lang} → {tgt_lang}")
-    print(f"  Input: {text[:100]}...")
-    
-    if not text.strip() or src_lang == tgt_lang:
-        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}")
-    
-    if src_code != "eng_Latn" and tgt_code != "eng_Latn":
-        print(f"  Two-step translation: {src_code}→eng_Latn→{tgt_code}")
-        to_english = translate_with_nllb(text, src_code, "eng_Latn", max_chunk_len)
-        return translate_with_nllb(to_english, "eng_Latn", tgt_code, max_chunk_len)
-    
-    return translate_with_nllb(text, src_code, tgt_code, max_chunk_len)
-
 def retrieve_docs(query: str, vs_path: str):
     if not vs_path or not os.path.exists(vs_path):
         return None