Update app/agents/crew_pipeline.py

app/agents/crew_pipeline.py

@@ -1,4 +1,4 @@
-# farmlingua/app/agents/crew_pipeline.pymemorysection
+# farmlingua/app/agents/crew_pipeline.py
 import os
 import sys
 import re
@@ -10,13 +10,12 @@ import numpy as np
 import torch
 import fasttext
 from huggingface_hub import hf_hub_download
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
 from sentence_transformers import SentenceTransformer
 from app.utils import config
-from app.utils.memory import memory_store  # memory module
+from app.utils.memory import memory_store
 from typing import List
 
-
 hf_cache = "/models/huggingface"
 os.environ["HF_HOME"] = hf_cache
 os.environ["TRANSFORMERS_CACHE"] = hf_cache
@@ -29,13 +28,11 @@ 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(
@@ -44,10 +41,8 @@ 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,
@@ -62,14 +57,16 @@ 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_pipeline = pipeline(
-    "translation",
-    model=config.TRANSLATION_MODEL_NAME,
-    device=0 if DEVICE == "cuda" else -1,
-    max_new_tokens=400,
-)
+
+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",
@@ -80,7 +77,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]:
@@ -101,17 +97,124 @@ def chunk_text(text: str, max_len: int = 400) -> List[str]:
         chunks.append(current.strip())
     return chunks
 
-def translate_text(text: str, src_lang: str, tgt_lang: str, max_chunk_len: int = 400) -> str:
-    if not text.strip():
-        return text
+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 = []
-    for chunk in chunks:
-        res = translation_pipeline(chunk, src_lang=src_lang, tgt_lang=tgt_lang)
-        translated_parts.append(res[0]["translation_text"])
+    
+    # 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)
+    
     return " ".join(translated_parts).strip()
 
-#  RAG retrieval
+def translate_simple(text: str, max_chunk_len: int = 400) -> str:
+    """Simple translation without language codes"""
+    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,
+                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)
+    
+    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
@@ -131,7 +234,6 @@ 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"}
@@ -147,7 +249,6 @@ 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"]):
@@ -170,7 +271,6 @@ 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)
@@ -183,7 +283,6 @@ 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]:
@@ -191,11 +290,7 @@ 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)
@@ -204,47 +299,40 @@ 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())  # fallback unique session
+        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"
         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}"})
@@ -263,19 +351,20 @@ 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"),