Update app.py

app.py

@@ -1,69 +1,152 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
 import gradio as gr
-import pickle
 import faiss
+import pickle
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+
+# -----------------------------
+# Load embedding model (for RAG)
+# -----------------------------
+embed_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
+
+# -----------------------------
+# Load FAISS + data
+# -----------------------------
+index = faiss.read_index("faiss_index.bin")
+chunks = pickle.load(open("chunks.pkl", "rb"))
+metadata = pickle.load(open("metadata.pkl", "rb"))
+
+# -----------------------------
+# Intent detection
+# -----------------------------
+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])
 
-# ----------------------------
-# Model Setup
-# ----------------------------
+    filtered_embeddings = np.array([index.reconstruct(i) for i in filtered_indices])
+    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.strip()
+
+# -----------------------------
+# Load Qwen model (CPU SAFE)
+# -----------------------------
 MODEL_NAME = "Qwen/Qwen3.5-0.8B"
 
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# CPU mode
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    device_map="cpu",  # force CPU
-    torch_dtype=torch.float32  # use float32 for CPU
+    torch_dtype=torch.float32   # CPU safe
 )
 
-# Text generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# ----------------------------
-# Load FAISS Index (optional)
-# ----------------------------
-try:
-    index = faiss.read_index("faiss_index.bin")
-    with open("metadata.pkl", "rb") as f:
-        metadata = pickle.load(f)
-except Exception as e:
-    print("FAISS index or metadata not found:", e)
-    index = None
-    metadata = None
-
-# ----------------------------
-# Chat Function
-# ----------------------------
-def chat_fn(user_input, chat_history=[]):
-    # If retrieval is enabled
-    if index:
-        query_vector = tokenizer(user_input, return_tensors="pt")["input_ids"].float().mean(dim=1).detach().numpy()
-        D, I = index.search(query_vector, k=3)
-        retrieved_texts = [metadata[i] for i in I[0]]
-        context = " ".join(retrieved_texts)
-        prompt = f"Context: {context}\nUser: {user_input}\nAI:"
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=150,
+    do_sample=True,
+    temperature=0.6
+)
+
+print("Model loaded successfully!")
+
+# -----------------------------
+# Chat function (RAG + LLM)
+# -----------------------------
+def chat_fn(user_input, history):
+    if history is None:
+        history = []
+
+    context = retrieve_context(user_input)
+
+    # 🔥 If no context → strict fallback
+    if not context:
+        response = "I don't know."
     else:
-        prompt = f"User: {user_input}\nAI:"
+        # 🔥 If context is small → return directly (FAST RAG)
+        if len(context.split()) < 50:
+            response = context.strip()
+        else:
+            # 🔥 Use LLM with strict instruction
+            prompt = f"""
+You are a livestock expert assistant for goats and cows.
+
+Use ONLY the information below to answer.
+If the answer is not present, say "I don't know".
+
+Context:
+{context}
+
+Question:
+{user_input}
 
-    # Generate response
-    output = generator(prompt, max_length=300, do_sample=True, top_p=0.9, temperature=0.7)
-    response = output[0]["generated_text"].split("AI:")[-1].strip()
+Answer in short and clear sentences.
+"""
 
-    # Update chat history
-    chat_history.append((user_input, response))
-    return chat_history, chat_history
+            output = generator(prompt)
+            text = output[0]["generated_text"]
 
-# ----------------------------
+            # Clean output
+            if prompt.strip() in text:
+                text = text.split(prompt.strip())[-1].strip()
+
+            response = text
+
+    history.append((user_input, response))
+    return history
+
+# -----------------------------
 # Gradio UI
-# ----------------------------
+# -----------------------------
 with gr.Blocks() as demo:
+    gr.Markdown("## 🐐 Livestock Chatbot (RAG + Qwen)")
+
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your Message")
+    msg = gr.Textbox(label="Ask about goats or cows")
     btn = gr.Button("Send")
 
-    btn.click(chat_fn, [msg, chatbot], [chatbot, chatbot])
+    btn.click(chat_fn, [msg, chatbot], chatbot)
 
 demo.launch()