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,7 +10,7 @@ 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, NllbTokenizer
 from sentence_transformers import SentenceTransformer
 from app.utils import config
 from app.utils.memory import memory_store  # memory module
@@ -64,26 +64,24 @@ def detect_language(text: str, top_k: int = 1):
 
 #  Translation model
 print(f"Loading translation model ({config.TRANSLATION_MODEL_NAME})...")
-from transformers import AutoModelForSeq2SeqLM 
 
-
-
-translation_tokenizer = AutoTokenizer.from_pretrained(
+translation_tokenizer = NllbTokenizer.from_pretrained(
     config.TRANSLATION_MODEL_NAME,
-    use_fast = True
+    cache_dir=hf_cache
 )
 
 translation_model = AutoModelForSeq2SeqLM.from_pretrained(
-    config.TRANSLATION_MODEL_NAME, 
-    torch_dtype = 'auto',
-    device_map = 'auto' if DEVICE == 'cuda' else None
-    
+    config.TRANSLATION_MODEL_NAME,
+    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+    cache_dir=hf_cache
 )
-if DEVICE == 'cpu':
-    translation_model = translation_model.to('cpu')
 
+if DEVICE == "cuda":
+    translation_model = translation_model.to("cuda")
+else:
+    translation_model = translation_model.to("cpu")
 
-    
+print(f"Translation model loaded on {DEVICE}")
 
 SUPPORTED_LANGS = {
     "eng_Latn": "English",
@@ -116,37 +114,50 @@ 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 NLLB model"""
     if not text.strip():
         return text
+    
+    if src_lang == tgt_lang:
+        return text
+    
     chunks = chunk_text(text, max_len=max_chunk_len)
     translated_parts = []
+    
     for chunk in chunks:
+        
+        translation_tokenizer.src_lang = src_lang
+        
+        # Tokenize
         inputs = translation_tokenizer(
             chunk,
-            return_tensors = 'pt',
-            padding = True,
-            truncation = True,
-            max_length = 400 
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
         ).to(translation_model.device)
-
-        #setting the target language token
+        
+       
         forced_bos_token_id = translation_tokenizer.convert_tokens_to_ids(tgt_lang)
-
+        
+        # Generate translation
         generated_tokens = translation_model.generate(
             **inputs,
-            forced_bos_token_id = forced_bos_token_id,
-            max_new_tokens = 400,
-            num_beams = 5,
-            early_stopping = True
+            forced_bos_token_id=forced_bos_token_id,
+            max_new_tokens=512,
+            num_beams=5,
+            early_stopping=True
         )
-
+        
+        # Decode
         translated_text = translation_tokenizer.batch_decode(
             generated_tokens,
-            skip_special_tokens = True
+            skip_special_tokens=True
         )[0]
+        
         translated_parts.append(translated_text)
-
-    return "".join(translated_parts).strip()   
+    
+    return " ".join(translated_parts).strip()
 
 
 #  RAG retrieval
@@ -242,7 +253,7 @@ 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.
@@ -273,7 +284,7 @@ def run_pipeline(user_query: str, session_id: str = None):
     system_prompt = (
         "You are FarmLingua, an AI assistant for Nigerian farmers. "
         "Answer questions directly and accurately with helpful farming advice. "
-        "Use clear, simple language with occasional emojis 🌾. "
+        "Use clear, simple language with occasional emojis . "
         "Be concise and focus on practical, actionable information. "
         "If asked who built you, say: 'KawaFarm LTD developed me to help farmers.'"
     )