Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -61,7 +61,7 @@ def detect_language(text: str, top_k: int = 1):
 
 print(f"Loading translation model ({config.TRANSLATION_MODEL_NAME})...")
 
-
+# Load tokenizer and model
 translation_tokenizer = AutoTokenizer.from_pretrained(config.TRANSLATION_MODEL_NAME)
 translation_model = AutoModelForSeq2SeqLM.from_pretrained(
     config.TRANSLATION_MODEL_NAME,
@@ -69,18 +69,18 @@ translation_model = AutoModelForSeq2SeqLM.from_pretrained(
     device_map="auto" if DEVICE == "cuda" else None
 )
 
+print(" Translation model loaded successfully")
+
 
 LANG_CODE_MAP = {
     "eng_Latn": "eng_Latn",  # English
     "ibo_Latn": "ibo_Latn",  # Igbo
     "yor_Latn": "yor_Latn",  # Yoruba
     "hau_Latn": "hau_Latn",  # Hausa
-    "swh_Latn": "swh_Latn",  # Swahili 
-    "amh_Latn": "amh_Latn",  # Amharic 
+    "swh_Latn": "swa_Latn",  # Swahili 
+    "amh_Latn": "amh_Ethi",  # Amharic 
 }
 
-
-
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
     "ibo_Latn": "Igbo",
@@ -113,44 +113,67 @@ def chunk_text(text: str, max_len: int = 400) -> List[str]:
 
 def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
     """
-    Translate text using the custom NLLB model directly
+    Translate text using  
+    IMPORTANT: Model expects format "src_lang text" -> "tgt_lang translation"
     """
+    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
     
     # Get language codes
     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 not hasattr(translation_tokenizer, 'lang_code_to_id'):
-        print("Warning: Tokenizer doesn't have lang_code_to_id attribute")
-        print(f"Available tokenizer special tokens: {translation_tokenizer.special_tokens_map}")
+    if src_code != "eng_Latn" and tgt_code != "eng_Latn":
+        print(f"  WARNING: Model wasn't trained on {src_code}→{tgt_code}")
+        print(f"  Will translate {src_code}→eng_Latn→{tgt_code}")
+        
 
-        return translate_text_simple(text, src_lang, tgt_lang, max_chunk_len)
+        to_english = translate_text_single(text, src_code, "eng_Latn", max_chunk_len)
+        return translate_text_single(to_english, "eng_Latn", tgt_code, max_chunk_len)
     
 
-    if src_code not in translation_tokenizer.lang_code_to_id:
-        print(f"Warning: Source language code '{src_code}' not found in tokenizer")
-        src_code = "eng_Latn"  
-    
-    if tgt_code not in translation_tokenizer.lang_code_to_id:
-        print(f"Warning: Target language code '{tgt_code}' not found in tokenizer")
-        tgt_code = "eng_Latn"  
-    
+    return translate_text_single(text, src_code, tgt_code, max_chunk_len)
 
-    translation_tokenizer.src_lang = src_code
-    
+def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len: int = 400) -> str:
+    """
+    Perform single translation step
+    """
 
-    forced_bos_token_id = translation_tokenizer.lang_code_to_id[tgt_code]
+   supported_pairs = [
+        ("eng_Latn", "ibo_Latn"), ("ibo_Latn", "eng_Latn"),
+        ("eng_Latn", "yor_Latn"), ("yor_Latn", "eng_Latn"),
+        ("eng_Latn", "hau_Latn"), ("hau_Latn", "eng_Latn"),
+        ("eng_Latn", "swa_Latn"), ("swa_Latn", "eng_Latn"),  
+        ("eng_Latn", "amh_Ethi"), ("amh_Ethi", "eng_Latn"), 
+    ]
+    
+    if (src_code, tgt_code) not in supported_pairs:
+        print(f"  WARNING: Pair {src_code}→{tgt_code} may not work well")
     
     chunks = chunk_text(text, max_len=max_chunk_len)
     translated_parts = []
     
-    for chunk in chunks:
+    for i, chunk in enumerate(chunks):
+        print(f"  Chunk {i+1}/{len(chunks)}: '{chunk[:50]}...'")
+        
         try:
 
-            inputs = translation_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512)
+            input_text = f"{src_code} {chunk}"
+            
+            # Tokenize
+            inputs = translation_tokenizer(
+                input_text, 
+                return_tensors="pt", 
+                truncation=True, 
+                max_length=512
+            )
             
             if DEVICE == "cuda":
                 inputs = {k: v.to(translation_model.device) for k, v in inputs.items()}
@@ -158,7 +181,6 @@ def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int =
             # Generate translation
             generated_tokens = translation_model.generate(
                 **inputs,
-                forced_bos_token_id=forced_bos_token_id,
                 max_new_tokens=400,
                 num_beams=4,
                 early_stopping=True
@@ -170,49 +192,20 @@ def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int =
                 skip_special_tokens=True
             )[0]
             
-            translated_parts.append(result)
-            
-        except Exception as e:
-            print(f"Translation error ({src_code}->{tgt_code}): {e}")
 
-            translated_parts.append(chunk)
-    
-    return " ".join(translated_parts).strip()
-
-def translate_text_simple(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
-    """
-    Simple fallback translation function if the main one fails
-    """
-    if not text.strip() or src_lang == tgt_lang:
-        return text
-    
-    chunks = chunk_text(text, max_len=max_chunk_len)
-    translated_parts = []
-    
-    for chunk in chunks:
-        try:
-
-            inputs = translation_tokenizer(chunk, return_tensors="pt", truncation=True, max_length=512)
+            if result.startswith(tgt_code + " "):
+                result = result[len(tgt_code) + 1:]
             
-            if DEVICE == "cuda":
-                inputs = {k: v.to(translation_model.device) for k, v in inputs.items()}
+            print(f"    → '{result[:50]}...'")
+            translated_parts.append(result.strip())
             
-            generated_tokens = translation_model.generate(
-                **inputs,
-                max_new_tokens=400
-            )
-            
-            result = translation_tokenizer.batch_decode(
-                generated_tokens, 
-                skip_special_tokens=True
-            )[0]
-            
-            translated_parts.append(result)
         except Exception as e:
-            print(f"Simple translation error: {e}")
-            translated_parts.append(chunk)
+            print(f"    ERROR: {e}")
+            translated_parts.append(chunk)  # Return original as fallback
     
-    return " ".join(translated_parts).strip()
+    final_result = " ".join(translated_parts).strip()
+    print(f"  Final: '{final_result[:100]}...'")
+    return final_result
 
 # RAG retrieval
 def retrieve_docs(query: str, vs_path: str):
@@ -307,6 +300,8 @@ def run_pipeline(user_query: str, session_id: str = None):
     lang_label, prob = detect_language(user_query, top_k=1)[0]
     if lang_label not in SUPPORTED_LANGS:
         lang_label = "eng_Latn"
+    
+    print(f"Detected language: {SUPPORTED_LANGS.get(lang_label, 'Unknown')}")
 
 
     translated_query = (
@@ -314,6 +309,8 @@ def run_pipeline(user_query: str, session_id: str = None):
         if lang_label != "eng_Latn"
         else user_query
     )
+    
+    print(f"Translated to English: {translated_query[:100]}...")
 
     intent, extra = detect_intent(translated_query)
 
@@ -363,6 +360,9 @@ def run_pipeline(user_query: str, session_id: str = None):
         else english_answer
     )
     final_answer = strip_markdown(final_answer)
+    
+    print(f"Final answer: {final_answer[:100]}...")
+    
     return {
         "session_id": session_id,
         "detected_language": SUPPORTED_LANGS.get(lang_label, "Unknown"),