Update app.py

app.py

@@ -1,10 +1,9 @@
-# app.py
 import gradio as gr
 import faiss
 import pickle
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 import os
 print("Files in current directory:", os.listdir())
@@ -19,22 +18,34 @@ chunks = pickle.load(open("chunks.pkl", "rb"))
 metadata = pickle.load(open("metadata.pkl", "rb"))
 
 # -----------------------------
-# Load Hugging Face LLM (CPU-friendly)
+# Load HF‑hosted small LLM
 # -----------------------------
-# Small model for HF Spaces CPU limits
-model_name = "TheBloke/vicuna-3B-1.1-HF"  # You can replace with a smaller model if needed
+model_name = "NousResearch/Nous-Hermes-1.0-GPTQ"
+
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")  # Hugging Face will manage CPU/GPU
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=200)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",  # Works on CPU or GPU
+    torch_dtype="auto"
+)
+
+generator = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=150,
+    do_sample=True,
+    temperature=0.6
+)
 
 print("LLM loaded successfully!")
 
 # -----------------------------
-# Detect query intent
+# Intent detection
 # -----------------------------
 def detect_query(query):
     query = query.lower()
-    
+
     animal = None
     topic = None
 
@@ -69,9 +80,7 @@ def retrieve_context(query):
 
     query_embedding = embed_model.encode([query])
 
-    filtered_embeddings = [index.reconstruct(i) for i in filtered_indices]
-    filtered_embeddings = np.array(filtered_embeddings)
-
+    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]
 
@@ -103,9 +112,12 @@ Question:
 Answer in short and clear sentences.
 """
 
-    # Generate response
-    response = generator(prompt, max_length=200, do_sample=True, temperature=0.5)
-    return response[0]["generated_text"]
+    response = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.6)
+    text = response[0]["generated_text"]
+    # Remove prompt if repeated
+    if prompt.strip() in text:
+        text = text.split(prompt.strip())[-1].strip()
+    return text
 
 # -----------------------------
 # Gradio UI