Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -1,4 +1,4 @@
-# farmlingua/app/agents/crew_pipeline.py
+# farmlingua/app/agents/crew_pipeline.pymemorysection
 import os
 import sys
 import re
@@ -10,12 +10,13 @@ import numpy as np
 import torch
 import fasttext
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from sentence_transformers import SentenceTransformer
 from app.utils import config
-from app.utils.memory import memory_store
+from app.utils.memory import memory_store  # memory module
 from typing import List
 
+
 hf_cache = "/models/huggingface"
 os.environ["HF_HOME"] = hf_cache
 os.environ["TRANSFORMERS_CACHE"] = hf_cache
@@ -28,11 +29,13 @@ if BASE_DIR not in sys.path:
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
+
 try:
     classifier = joblib.load(config.CLASSIFIER_PATH)
 except Exception:
     classifier = None
 
+
 print(f"Loading expert model ({config.EXPERT_MODEL_NAME})...")
 tokenizer = AutoTokenizer.from_pretrained(config.EXPERT_MODEL_NAME, use_fast=False)
 model = AutoModelForCausalLM.from_pretrained(
@@ -41,8 +44,10 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
+
 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,
@@ -57,56 +62,14 @@ 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)]
 
+#  Translation model
 print(f"Loading translation model ({config.TRANSLATION_MODEL_NAME})...")
-
-NLLB_MODEL = "facebook/nllb-200-distilled-600M"
-print(f"Using model: {NLLB_MODEL}")
-
-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
-    "yor_Latn": "yor_Latn",  # Yoruba
-    "hau_Latn": "hau_Latn",  # Hausa
-    "swh_Latn": "swa_Latn",  # Swahili
-    "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
-}
+translation_pipeline = pipeline(
+    "translation",
+    model=config.TRANSLATION_MODEL_NAME,
+    device=0 if DEVICE == "cuda" else -1,
+    max_new_tokens=400,
+)
 
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
@@ -117,6 +80,7 @@ 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]:
@@ -138,108 +102,16 @@ 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:
-    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)")
+    if not text.strip():
         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 not hasattr(translation_tokenizer, 'lang_code_to_id'):
-        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, 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, trying alternatives...")
-        tgt_code = LANG_CODE_MAP_ALT.get(tgt_lang, "eng_Latn")
-    
-    
-    translation_tokenizer.src_lang = src_code
-    
-   
-    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(
-                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]
-            
-            print(f"  Chunk {i+1}: '{chunk[:30]}...' → '{result[:30]}...'")
-            translated_parts.append(result.strip())
-            
-        except Exception as e:
-            print(f"  Chunk {i+1} error: {e}")
-            translated_parts.append(chunk)
-    
+    for chunk in chunks:
+        res = translation_pipeline(chunk, src_lang=src_lang, tgt_lang=tgt_lang)
+        translated_parts.append(res[0]["translation_text"])
     return " ".join(translated_parts).strip()
 
-
+#  RAG retrieval
 def retrieve_docs(query: str, vs_path: str):
     if not vs_path or not os.path.exists(vs_path):
         return None
@@ -259,6 +131,7 @@ def retrieve_docs(query: str, vs_path: str):
         return "\n\n".join(docs) if docs else None
     return None
 
+
 def get_weather(state_name: str) -> str:
     url = "http://api.weatherapi.com/v1/current.json"
     params = {"key": config.WEATHER_API_KEY, "q": f"{state_name}, Nigeria", "aqi": "no"}
@@ -274,6 +147,7 @@ def get_weather(state_name: str) -> str:
         f"- Wind: {data['current']['wind_kph']} kph"
     )
 
+
 def detect_intent(query: str):
     q_lower = (query or "").lower()
     if any(word in q_lower for word in ["weather", "temperature", "rain", "forecast"]):
@@ -296,6 +170,7 @@ def detect_intent(query: str):
             pass
     return "normal", None
 
+# expert runner
 def run_qwen(messages: List[dict], max_new_tokens: int = 1300) -> str:
     text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
@@ -308,6 +183,7 @@ def run_qwen(messages: List[dict], max_new_tokens: int = 1300) -> str:
     output_ids = generated_ids[0][len(inputs.input_ids[0]):].tolist()
     return tokenizer.decode(output_ids, skip_special_tokens=True).strip()
 
+#  Memory
 MAX_HISTORY_MESSAGES = getattr(config, "MAX_HISTORY_MESSAGES", 30)
 
 def build_messages_from_history(history: List[dict], system_prompt: str) -> List[dict]:
@@ -315,7 +191,11 @@ def build_messages_from_history(history: List[dict], system_prompt: str) -> List
     msgs.extend(history)
     return msgs
 
+
 def strip_markdown(text: str) -> str:
+    """
+    Remove Markdown formatting like **bold**, *italic*, and `inline code`.
+    """
     if not text:
         return ""
     text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)
@@ -324,40 +204,47 @@ def strip_markdown(text: str) -> str:
     text = re.sub(r'^#+\s+', '', text, flags=re.MULTILINE)
     return text
 
+#  Main pipeline
 def run_pipeline(user_query: str, session_id: str = None):
+    """
+    Run FarmLingua pipeline with per-session memory.
+    Each session_id keeps its own history.
+    """
     if session_id is None:
-        session_id = str(uuid.uuid4())
+        session_id = str(uuid.uuid4())  # fallback unique session
 
+    # 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"
         else user_query
     )
-    
-    print(f"Translated to English: {translated_query[:100]}...")
 
     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:]
 
+    
     history.append({"role": "user", "content": translated_query})
 
+    
     system_prompt = (
         "You are FarmLingua, an AI assistant for Nigerian farmers. "
         "Answer directly without repeating the question. "
         "Use clear farmer-friendly English with emojis . "
         "Avoid jargon and irrelevant details. "
         "If asked who built you, say: 'KawaFarm LTD developed me to help farmers.'"
+      
     )
 
+    
     if intent == "weather" and extra:
         weather_text = get_weather(extra)
         history.append({"role": "user", "content": f"Rewrite this weather update simply for farmers:\n{weather_text}"})
@@ -376,22 +263,21 @@ 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)
 
+    # Save assistant reply
     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)
 
+    # Translate back if needed
     final_answer = (
         translate_text(english_answer, src_lang="eng_Latn", tgt_lang=lang_label)
         if lang_label != "eng_Latn"
         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"),
         "answer": final_answer
-    }
+    }