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,25 +62,15 @@ 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})...")
-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
+translation_pipeline = pipeline(
+    "translation",
+    model=config.TRANSLATION_MODEL_NAME,
+    device=0 if DEVICE == "cuda" else -1,
+    max_new_tokens=400,
 )
 
-print("Translation model loaded successfully")
-
-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",
-}
-
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
     "ibo_Latn": "Igbo",
@@ -85,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]:
@@ -106,83 +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 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:
-    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 i, chunk in enumerate(chunks):
-        print(f"  Chunk {i+1}/{len(chunks)}: '{chunk[:50]}...'")
-        
-        try:
-            input_text = f"{src_code} {chunk}"
-            
-            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()}
-            
-            generated_tokens = translation_model.generate(
-                **inputs,
-                max_new_tokens=400,
-                num_beams=4,
-                early_stopping=True
-            )
-            
-            result = translation_tokenizer.batch_decode(
-                generated_tokens, 
-                skip_special_tokens=True
-            )[0]
-            
-            if result.startswith(tgt_code + " "):
-                result = result[len(tgt_code) + 1:]
-            
-            print(f"    → '{result[:50]}...'")
-            translated_parts.append(result.strip())
-            
-        except Exception as e:
-            print(f"    ERROR: {e}")
-            translated_parts.append(chunk)
-    
-    final_result = " ".join(translated_parts).strip()
-    print(f"  Final: '{final_result[:100]}...'")
-    return final_result
+    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
@@ -202,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"}
@@ -217,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"]):
@@ -239,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)
@@ -251,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]:
@@ -258,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)
@@ -267,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}"})
@@ -319,20 +263,19 @@ 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"),