Update app.py

app.py

@@ -3,6 +3,8 @@ import faiss
 import pickle
 import numpy as np
 from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 import os
 
 print("Files in current directory:", os.listdir())
@@ -11,11 +13,33 @@ print("Files in current directory:", os.listdir())
 # Load RAG components
 # -----------------------------
 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"))
 
+# -----------------------------
+# Load Qwen 2.5B Instruct model
+# -----------------------------
+model_name = "Qwen/Qwen2.5-1.5B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=200,
+    do_sample=True,
+    temperature=0.6
+)
+
+print("Qwen model loaded successfully!")
+
 # -----------------------------
 # Intent detection
 # -----------------------------
@@ -42,40 +66,55 @@ def detect_query(query):
 def retrieve_context(query, top_k=2):
     animal, topic = detect_query(query)
 
-    # Filter relevant chunks based on metadata
     filtered_indices = [
         i for i, meta in enumerate(metadata)
         if (not animal or meta["animal"] == animal) and
            (not topic or meta["topic"] == topic)
     ]
 
-    # If no specific filter matches, consider all chunks
     if not filtered_indices:
         filtered_indices = list(range(len(chunks)))
 
-    # Embed query
     query_embedding = embed_model.encode([query])
     filtered_embeddings = np.array([index.reconstruct(i) for i in filtered_indices])
 
-    # Compute distances and get top-k closest chunks
     distances = np.linalg.norm(filtered_embeddings - query_embedding, axis=1)
     top_indices = distances.argsort()[:top_k]
 
-    # Combine top chunks into context
     context = "\n".join(chunks[filtered_indices[idx]] for idx in top_indices)
-
     return context.strip()
 
 # -----------------------------
-# Chat function (RAG only)
+# Chat function (RAG + Qwen)
 # -----------------------------
 def chat(user_input):
     context = retrieve_context(user_input)
     if not context:
         return "I don't know."
-    
-    # Return context with clear formatting
-    return f"Answer from retrieved data:\n\n{context}"
+
+    prompt = f"""
+You are a livestock expert assistant.
+
+Use ONLY the information below to answer the question.
+If the answer is not present, say "I don't know".
+
+Context:
+{context}
+
+Question:
+{user_input}
+
+Answer in full, clear sentences.
+"""
+
+    response = generator(prompt, max_new_tokens=200, do_sample=True, temperature=0.6)
+    text = response[0]["generated_text"]
+
+    # Remove prompt repetition
+    if prompt.strip() in text:
+        text = text.split(prompt.strip())[-1].strip()
+
+    return text
 
 # -----------------------------
 # Gradio UI
@@ -84,7 +123,6 @@ gr.Interface(
     fn=chat,
     inputs=gr.Textbox(lines=2, placeholder="Ask a question about livestock..."),
     outputs=gr.Textbox(),
-    title="Livestock Chatbot (RAG only)",
-    description="This chatbot answers livestock questions using only retrieved data. No AI model is used.",
-    allow_flagging="never"
+    title="Livestock Chatbot (RAG + Qwen)",
+    description="This chatbot answers livestock questions using retrieved data and Qwen Instruct model."
 ).launch()