Update app.py

app.py

@@ -1,95 +1,28 @@
-import gradio as gr
-import faiss
+# app.py
+import os
 import pickle
-import numpy as np
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
-import os
-from huggingface_hub import login
-
-# -----------------------------
-# Login with HF_TOKEN (secret)
-# -----------------------------
-HF_TOKEN = os.environ["HF_TOKEN"]  # Must be set in Space secrets
-login(HF_TOKEN)
-
-print("Files in current directory:", os.listdir())
-
-# -----------------------------
-# Load RAG components
-# -----------------------------
-print("Loading embedding model...")
-embed_model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
-
-print("Loading FAISS index and metadata...")
-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)))
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import faiss
 
-    query_embedding = embed_model.encode([query])
-    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]
+# --- Load HF token from Space secrets ---
+HF_TOKEN = os.environ["HF_TOKEN"]
 
-    context = ""
-    for idx in top_indices:
-        real_index = filtered_indices[idx]
-        context += chunks[real_index] + "\n"
+# --- Model configuration ---
+MODEL_NAME = "Qwen/Qwen3.5-0.8B-Base"
 
-    return context.strip()
+print(f"Loading model {MODEL_NAME} on CPU...")
 
-# -----------------------------
-# Load Qwen3.5‑0.8B‑Base
-# -----------------------------
-model_name = "Qwen/Qwen3.5-0.8B-Base"
-print(f"Loading model {model_name} (may take a while on CPU)...")
+# --- Load tokenizer ---
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+# --- Load model (CPU only) ---
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float32,
-    device_map={"": "cpu"},
+    MODEL_NAME,
     use_auth_token=HF_TOKEN
 )
 
+# --- Setup text-generation pipeline ---
 generator = pipeline(
     "text-generation",
     model=model,
@@ -97,45 +30,54 @@ generator = pipeline(
     max_new_tokens=150,
     do_sample=True,
     temperature=0.7,
-    device=-1  # CPU only
+    device=-1  # CPU
 )
 
-print("LLM loaded successfully!")
-
-# -----------------------------
-# Chat function
-# -----------------------------
-def chat(user_input):
-    context = retrieve_context(user_input)
-    if not context:
-        return "I don't know."
-
-    prompt = f"""
-You are a livestock expert assistant for goats and cows.
-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.
-"""
-    response = generator(prompt)
-    text = response[0]["generated_text"]
-    if prompt.strip() in text:
-        text = text.split(prompt.strip())[-1].strip()
-    return text
-
-# -----------------------------
-# Gradio UI
-# -----------------------------
-gr.Interface(
-    fn=chat,
-    inputs="text",
-    outputs="text",
-    title="Livestock Chatbot (RAG + Qwen3.5‑0.8B‑Base)",
-    description="Answers livestock questions using RAG retrieval and the Qwen3.5‑0.8B base model."
-).launch()
+print("Model loaded successfully!")
+
+# --- Load FAISS index and metadata ---
+if os.path.exists("faiss_index.bin") and os.path.exists("metadata.pkl"):
+    print("Loading FAISS index and metadata...")
+    index = faiss.read_index("faiss_index.bin")
+    with open("metadata.pkl", "rb") as f:
+        metadata = pickle.load(f)
+    print("FAISS index loaded.")
+else:
+    print("FAISS index or metadata not found. Make sure you uploaded faiss_index.bin and metadata.pkl")
+    index = None
+    metadata = None
+
+# --- Embeddings model for query ---
+from sentence_transformers import SentenceTransformer
+embed_model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
+
+# --- RAG retrieval function ---
+def retrieve_docs(query, top_k=3):
+    if index is None or metadata is None:
+        return []
+    q_embed = embed_model.encode([query])
+    distances, idxs = index.search(q_embed, top_k)
+    docs = [metadata[i] for i in idxs[0]]
+    return docs
+
+# --- Chatbot function ---
+def chat(query):
+    # Retrieve relevant docs
+    retrieved_docs = retrieve_docs(query)
+    context = "\n".join(retrieved_docs) if retrieved_docs else ""
+    
+    # Combine context with user query
+    prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
+    
+    # Generate response
+    output = generator(prompt)
+    return output[0]["generated_text"]
+
+# --- Example usage ---
+if __name__ == "__main__":
+    while True:
+        query = input("You: ")
+        if query.lower() in ["exit", "quit"]:
+            break
+        answer = chat(query)
+        print("Bot:", answer)