Update app.py

app.py

@@ -1,31 +1,122 @@
 from llama_cpp import Llama
 import gradio as gr
+import faiss
+import pickle
+import numpy as np
+from sentence_transformers import SentenceTransformer
 
-# IMPORTANT: Use chat_format="qwen" (qwen2 is NOT supported)
+# -----------------------------
+# Load LLM
+# -----------------------------
 model = Llama(
     model_path="qwen2.5-1.5B-q4.gguf",
     n_ctx=4096,
     n_gpu_layers=0,
-    chat_format="qwen",  
+    chat_format="qwen",
 )
 
+# -----------------------------
+# Load RAG
+# -----------------------------
+embed_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
+
+index = faiss.read_index("faiss_index.bin")
+chunks = pickle.load(open("chunks.pkl", "rb"))
+metadata = pickle.load(open("metadata.pkl", "rb"))
+
+# -----------------------------
+# Detect query intent
+# -----------------------------
+def detect_query(query):
+    query = query.lower()
+    
+    animal = None
+    topic = None
+
+    if "goat" in query:
+        animal = "goat"
+    elif "cow" in query:
+        animal = "cow"
+
+    if any(word in query for word in ["feed", "diet", "khilana"]):
+        topic = "feeding"
+    elif any(word in query for word in ["disease", "bimari"]):
+        topic = "disease"
+
+    return animal, topic
+
+# -----------------------------
+# Retrieve context (RAG)
+# -----------------------------
+def retrieve_context(query):
+    animal, topic = detect_query(query)
+
+    filtered_indices = []
+    for i, meta in enumerate(metadata):
+        if animal and meta["animal"] != animal:
+            continue
+        if topic and meta["topic"] != topic:
+            continue
+        filtered_indices.append(i)
+
+    if not filtered_indices:
+        filtered_indices = list(range(len(chunks)))
+
+    query_embedding = embed_model.encode([query])
+
+    filtered_embeddings = [index.reconstruct(i) for i in filtered_indices]
+    filtered_embeddings = np.array(filtered_embeddings)
+
+    distances = np.linalg.norm(filtered_embeddings - query_embedding, axis=1)
+    top_indices = distances.argsort()[:2]
+
+    context = ""
+    for idx in top_indices:
+        real_index = filtered_indices[idx]
+        context += chunks[real_index] + "\n"
+
+    return context
+
+# -----------------------------
+# Chat function (UPDATED)
+# -----------------------------
 def chat(user_input):
+    context = retrieve_context(user_input)
+
+    prompt = f"""
+You are a livestock expert assistant.
+
+Use ONLY the information below to answer.
+If answer is not present, say "I don't know".
+
+Context:
+{context}
+
+Question:
+{user_input}
+
+Answer in short and clear sentences.
+"""
+
     messages = [
-        {"role": "system", "content": "You are a helpful assistant. Answer ONLY the question. Do NOT continue, do NOT ask questions, do NOT add extra text."},
-        {"role": "user", "content": user_input}
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": prompt}
     ]
 
     response = model.create_chat_completion(
         messages=messages,
-        max_tokens=256,
-        temperature=0.7,
+        max_tokens=200,
+        temperature=0.5,
     )
 
     return response["choices"][0]["message"]["content"]
 
+# -----------------------------
+# Gradio UI (UNCHANGED)
+# -----------------------------
 gr.Interface(
     fn=chat,
     inputs="text",
     outputs="text",
-    title="Qwen2.5-1.5B Q4 Chatbot"
+    title="Livestock Chatbot"
 ).launch()