Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -43,7 +43,6 @@ model = AutoModelForCausalLM.from_pretrained(
 
 embedder = SentenceTransformer(config.EMBEDDING_MODEL)
 
-# Language detector
 print(f"Loading FastText language identifier ({config.LANG_ID_MODEL_REPO})...")
 lang_model_path = hf_hub_download(
     repo_id=config.LANG_ID_MODEL_REPO,
@@ -58,10 +57,7 @@ def detect_language(text: str, top_k: int = 1):
     labels, probs = lang_identifier.predict(clean_text, k=top_k)
     return [(l.replace("__label__", ""), float(p)) for l, p in zip(labels, probs)]
 
-
 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,16 +65,15 @@ translation_model = AutoModelForSeq2SeqLM.from_pretrained(
     device_map="auto" if DEVICE == "cuda" else None
 )
 
-print(" Translation model loaded successfully")
-
+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": "swa_Latn",  # Swahili 
-    "amh_Latn": "amh_Ethi",  # Amharic 
+    "eng_Latn": "eng_Latn",
+    "ibo_Latn": "ibo_Latn",
+    "yor_Latn": "yor_Latn",
+    "hau_Latn": "hau_Latn",
+    "swh_Latn": "swa_Latn",
+    "amh_Latn": "amh_Ethi",
 }
 
 SUPPORTED_LANGS = {
@@ -90,7 +85,6 @@ SUPPORTED_LANGS = {
     "amh_Latn": "Amharic",
 }
 
-# Text chunking
 _SENTENCE_SPLIT_RE = re.compile(r'(?<=[.!?])\s+')
 
 def chunk_text(text: str, max_len: int = 400) -> List[str]:
@@ -112,10 +106,6 @@ def chunk_text(text: str, max_len: int = 400) -> List[str]:
     return chunks
 
 def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
-    """
-    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]}...")
     
@@ -123,35 +113,26 @@ def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int =
         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 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}")
-        
-
         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)
     
-
     return translate_text_single(text, src_code, tgt_code, max_chunk_len)
 
 def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len: int = 400) -> str:
-    """
-    Perform single translation step
-    """
-
-   supported_pairs = [
+    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"), 
+        ("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:
@@ -164,10 +145,8 @@ def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len
         print(f"  Chunk {i+1}/{len(chunks)}: '{chunk[:50]}...'")
         
         try:
-
             input_text = f"{src_code} {chunk}"
             
-            # Tokenize
             inputs = translation_tokenizer(
                 input_text, 
                 return_tensors="pt", 
@@ -178,7 +157,6 @@ def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len
             if DEVICE == "cuda":
                 inputs = {k: v.to(translation_model.device) for k, v in inputs.items()}
             
-            # Generate translation
             generated_tokens = translation_model.generate(
                 **inputs,
                 max_new_tokens=400,
@@ -186,13 +164,11 @@ def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len
                 early_stopping=True
             )
             
-            # Decode
             result = translation_tokenizer.batch_decode(
                 generated_tokens, 
                 skip_special_tokens=True
             )[0]
             
-
             if result.startswith(tgt_code + " "):
                 result = result[len(tgt_code) + 1:]
             
@@ -201,13 +177,12 @@ def translate_text_single(text: str, src_code: str, tgt_code: str, max_chunk_len
             
         except Exception as e:
             print(f"    ERROR: {e}")
-            translated_parts.append(chunk)  # Return original as fallback
+            translated_parts.append(chunk)
     
     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):
     if not vs_path or not os.path.exists(vs_path):
         return None
@@ -296,14 +271,12 @@ def run_pipeline(user_query: str, session_id: str = None):
     if session_id is None:
         session_id = str(uuid.uuid4())
 
-    # Language detection
     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 = (
         translate_text(user_query, src_lang=lang_label, tgt_lang="eng_Latn")
         if lang_label != "eng_Latn"
@@ -314,7 +287,6 @@ def run_pipeline(user_query: str, session_id: str = None):
 
     intent, extra = detect_intent(translated_query)
 
-    # Load conversation history
     history = memory_store.get_history(session_id) or []
     if len(history) > MAX_HISTORY_MESSAGES:
         history = history[-MAX_HISTORY_MESSAGES:]
@@ -347,13 +319,11 @@ def run_pipeline(user_query: str, session_id: str = None):
         messages_for_qwen = build_messages_from_history(history, system_prompt)
         english_answer = run_qwen(messages_for_qwen, max_new_tokens=700)
 
-
     history.append({"role": "assistant", "content": english_answer})
     if len(history) > MAX_HISTORY_MESSAGES:
         history = history[-MAX_HISTORY_MESSAGES:]
     memory_store.save_history(session_id, history)
 
-
     final_answer = (
         translate_text(english_answer, src_lang="eng_Latn", tgt_lang=lang_label)
         if lang_label != "eng_Latn"