Update app.py

app.py

@@ -1,22 +1,31 @@
 import gradio as gr
 import spaces
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+import faiss
+import numpy as np
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+from transformers import RagTokenizer, RagSequenceForGeneration
 
-# ——— Configuration ———
-MODEL_NAME   = "facebook/rag-sequence-nq"
+# — Config —
 DATASET_NAME = "username/mealplan-chunks"
+INDEX_PATH   = "mealplan.index"
+MODEL_NAME   = "facebook/rag-sequence-nq"
+
+# — Load chunks & FAISS index —
+ds = load_dataset(DATASET_NAME, split="train")
+texts   = ds["text"]
+sources = ds["source"]
+pages   = ds["page"]
+
+# — Embeddings embedder & FAISS —
+embedder         = SentenceTransformer("all-MiniLM-L6-v2")
+chunk_embeddings = embedder.encode(texts, convert_to_numpy=True)
+index            = faiss.read_index(INDEX_PATH)
 
-# ——— Initialize RAG ———
+# — RAG generator —
 tokenizer = RagTokenizer.from_pretrained(MODEL_NAME)
-retriever = RagRetriever.from_pretrained(
-    MODEL_NAME,
-    index_name="exact",
-    use_dummy_dataset=False,
-    dataset_name=DATASET_NAME
-)
-model = RagSequenceForGeneration.from_pretrained(MODEL_NAME, retriever=retriever)
+model     = RagSequenceForGeneration.from_pretrained(MODEL_NAME)
 
-# ——— Core chat callback ———
 @spaces.GPU
 def respond(
     message: str,
@@ -27,38 +36,45 @@ def respond(
     avoid: str,
     weeks: str,
 ):
-    # parse avoidances
+    # Parse preferences
     avoid_list = [a.strip() for a in avoid.split(",") if a.strip()]
-    # build prefs string
     prefs = (
-        f"Goal={goal}; "
-        f"Diet={','.join(diet)}; "
-        f"Meals={meals}/day; "
-        f"Avoid={','.join(avoid_list)}; "
-        f"Duration={weeks}"
+        f"Goal={goal}; Diet={','.join(diet)}; "
+        f"Meals={meals}/day; Avoid={','.join(avoid_list)}; Duration={weeks}"
     )
-    # system guardrail + prefs + question
+
+    # 1) Query embedding & FAISS search
+    q_emb = embedder.encode([message], convert_to_numpy=True)
+    D, I  = index.search(q_emb, 5)  # top-5
+    ctx_chunks = [
+        f"[{sources[i]} p{pages[i]}] {texts[i]}" for i in I[0]
+    ]
+    context = "\n".join(ctx_chunks)
+
+    # 2) Build prompt
     prompt = (
         "SYSTEM: Only answer using the provided CONTEXT. "
-        "If it’s not there, say “I’m sorry, I don’t know.”\n"
+        "If it’s not there, say \"I'm sorry, I don't know.\"\n"
         f"PREFS: {prefs}\n"
+        f"CONTEXT:\n{context}\n"
         f"Q: {message}\n"
     )
-    # generate
+
+    # 3) Generate
     inputs  = tokenizer([prompt], return_tensors="pt")
     outputs = model.generate(**inputs, num_beams=2, max_new_tokens=200)
     answer  = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
 
+    # 4) Update history
     history = history or []
     history.append((message, answer))
     return history
 
-# ——— Build the Chat Interface ———
-# preference controls
+# — Build Gradio chat interface —
 goal  = gr.Dropdown(["Lose weight","Bulk","Maintain"], value="Lose weight", label="Goal")
 diet  = gr.CheckboxGroup(["Omnivore","Vegetarian","Vegan","Keto","Paleo","Low-Carb"], label="Diet Style")
-meals = gr.Slider(1, 6, value=3, step=1, label="Meals per day")
-avoid = gr.Textbox(placeholder="e.g. Gluten, Dairy, Nuts, Eggs, Soy…", label="Avoidances (comma-separated)")
+meals = gr.Slider(1,6,value=3,step=1,label="Meals per day")
+avoid = gr.Textbox(placeholder="e.g. Gluten, Dairy, Nuts...", label="Avoidances (comma-separated)")
 weeks = gr.Dropdown(["1 week","2 weeks","3 weeks","4 weeks"], value="1 week", label="Plan Length")
 
 demo = gr.ChatInterface(