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,9 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
-
 embedder = SentenceTransformer(config.EMBEDDING_MODEL)
 
-#   language detector
+# 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,15 +58,29 @@ 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_en_to_fr",
-    model=config.TRANSLATION_MODEL_NAME,
-    device=0 if DEVICE == "cuda" else -1,
-    max_new_tokens=400,
+
+
+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
 )
 
+
+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 
+}
+
+
+
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
     "ibo_Latn": "Igbo",
@@ -102,16 +112,109 @@ 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:
-    if not text.strip():
+    """
+    Translate text using the custom NLLB model directly
+    """
+    if not text.strip() or src_lang == tgt_lang:
+        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")
+    
+
+    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}")
+
+        return translate_text_simple(text, src_lang, tgt_lang, 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"  
+    
+
+    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 chunk in 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()}
+            
+            # 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
+            )
+            
+            # Decode
+            result = translation_tokenizer.batch_decode(
+                generated_tokens, 
+                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:
-        res = translation_pipeline(chunk, src_lang=src_lang, tgt_lang=tgt_lang)
-        translated_parts.append(res[0]["translation_text"])
+        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
+            )
+            
+            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)
+    
     return " ".join(translated_parts).strip()
 
-#  RAG retrieval
+# RAG retrieval
 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,20 +299,16 @@ 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"
 
+
     translated_query = (
         translate_text(user_query, src_lang=lang_label, tgt_lang="eng_Latn")
         if lang_label != "eng_Latn"
@@ -226,25 +317,21 @@ def run_pipeline(user_query: str, session_id: str = None):
 
     intent, extra = detect_intent(translated_query)
 
-    #  Load conversation history
+    # 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,13 +350,13 @@ 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"
@@ -280,4 +367,4 @@ def run_pipeline(user_query: str, session_id: str = None):
         "session_id": session_id,
         "detected_language": SUPPORTED_LANGS.get(lang_label, "Unknown"),
         "answer": final_answer
-    }
+    }