Update app.py

app.py

@@ -1,106 +1,76 @@
-# app.py
 import os
-import glob
-import faiss
+import fitz  # PyMuPDF
 import numpy as np
-
-import gradio as gr
-import spaces
-
-from unstructured.partition.pdf import partition_pdf
+import faiss
 from sentence_transformers import SentenceTransformer
-from transformers import RagTokenizer, RagSequenceForGeneration
-
-# ─── Configuration ─────────────────────────────────────────────
-PDF_FOLDER = "meal_plans"
-MODEL_NAME = "facebook/rag-sequence-nq"
-EMBED_MODEL = "all-MiniLM-L6-v2"
-TOP_K = 5
-
-# ─── 1) LOAD + CHUNK ALL PDFs ──────────────────────────────────
-rag_tokenizer = RagTokenizer.from_pretrained(MODEL_NAME)
-texts, sources, pages = [], [], []
-
-for pdf_path in glob.glob(f"{PDF_FOLDER}/*.pdf"):
-    book = os.path.basename(pdf_path)
-    pages_data = partition_pdf(filename=pdf_path)
-    for pg_num, page in enumerate(pages_data, start=1):
-        enc = rag_tokenizer(
-            page.text,
-            max_length=800,
-            truncation=True,
-            return_overflowing_tokens=True,
-            stride=50,
-            return_tensors="pt"
-        )
-        for token_ids in enc["input_ids"]:
-            chunk = rag_tokenizer.decode(token_ids, skip_special_tokens=True)
-            texts.append(chunk)
-            sources.append(book)
-            pages.append(pg_num)
+import gradio as gr
+import spaces  # for ZeroGPU
 
-# ─── 2) EMBED + BUILD FAISS INDEX ─────────────────────────────
-embedder = SentenceTransformer(EMBED_MODEL)
-embeddings = embedder.encode(texts, convert_to_numpy=True)
-dim = embeddings.shape[1]
-index = faiss.IndexFlatL2(dim)
-index.add(embeddings)
+@spaces.GPU
+def query_app(user_input, include_source, verbose):
+    return search_index(user_input, index, documents, include_source, verbose)
 
-# ─── 3) LOAD RAG GENERATOR ────────────────────────────────────
-tokenizer = RagTokenizer.from_pretrained(MODEL_NAME)
-generator = RagSequenceForGeneration.from_pretrained(MODEL_NAME)
+# PDF reader
+def extract_text_from_pdf(folder_path="meal_plans"):
+    documents = []
+    for filename in os.listdir(folder_path):
+        if filename.lower().endswith(".pdf"):
+            path = os.path.join(folder_path, filename)
+            try:
+                doc = fitz.open(path)
+                text = ""
+                for page in doc:
+                    text += page.get_text()
+                documents.append({"text": text, "source": filename})
+            except Exception as e:
+                print(f"Error reading {filename}: {e}")
+    return documents
 
-@spaces.GPU
-def respond(
-    message: str,
-    history: list[tuple[str,str]],
-    goal: str,
-    diet: list[str],
-    meals: int,
-    avoid: str,
-    weeks: str
-):
-    # build prefs string
-    avoid_list = [a.strip() for a in avoid.split(",") if a.strip()]
-    prefs = (
-        f"Goal={goal}; Diet={','.join(diet)}; "
-        f"Meals={meals}/day; Avoid={','.join(avoid_list)}; Duration={weeks}"
-    )
-    # 1) RETRIEVE top-k chunks
-    q_emb = embedder.encode([message], convert_to_numpy=True)
-    D, I  = index.search(q_emb, TOP_K)
-    context = "\n".join(f"[{sources[i]} p{pages[i]}] {texts[i]}" for i in I[0])
+# Index builder
+def create_index(docs):
+    texts = [doc["text"] for doc in docs]
+    embeddings = model.encode(texts)
+    dim = embeddings[0].shape[0]
+    index = faiss.IndexFlatL2(dim)
+    index.add(np.array(embeddings).astype("float32"))
+    return index
 
-    # 2) BUILD PROMPT with guardrail
-    prompt = (
-        "SYSTEM: Only answer using the provided CONTEXT. "
-        "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"
-    )
+# Search logic
+def search_index(query, index, docs, include_source=True, verbose=False, top_k=3):
+    query_vec = model.encode([query])
+    D, I = index.search(np.array(query_vec).astype("float32"), top_k)
+    responses = []
+    for i in I[0]:
+        doc = docs[i]
+        snippet = doc["text"][:750 if verbose else 300].replace("\n", " ").strip()
+        label = f"**📄 {doc['source']}**\n" if include_source else ""
+        responses.append(f"{label}{snippet}...")
+    return "\n\n---\n\n".join(responses)
 
-    # 3) GENERATE
-    inputs  = tokenizer([prompt], return_tensors="pt")
-    outputs = generator.generate(**inputs, num_beams=2, max_new_tokens=200)
-    answer  = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+# Setup
+model = SentenceTransformer("all-MiniLM-L6-v2")
+documents = extract_text_from_pdf("meal_plans")
+index = create_index(documents)
 
-    # update chat history
-    history = history or []
-    history.append((message, answer))
-    return history
+# Gradio UI
+with gr.Blocks(title="Meal Plan Chat Assistant") as demo:
+    gr.Markdown("## 🍽️ Meal Plan Assistant\nChat with your PDF documents in `meal_plans/` folder.")
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.Chatbot()
+            user_input = gr.Textbox(placeholder="Ask something...", show_label=False)
+            send_btn = gr.Button("Ask")
+        with gr.Column(scale=1):
+            include_source = gr.Checkbox(label="Include Source", value=True)
+            verbose = gr.Checkbox(label="Verbose Mode", value=False)
 
-# ─── 4) BUILD UI ────────────────────────────────────────────────
-goal  = gr.Dropdown(["Lose weight","Bulk","Maintain"], label="Goal", value="Lose weight")
-diet  = gr.CheckboxGroup(["Omnivore","Vegetarian","Vegan","Keto","Paleo","Low-Carb"], label="Diet Style")
-meals = gr.Slider(1, 6, step=1, value=3, 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"], label="Plan Length", value="1 week")
+    def user_query(msg, history, source, verbose_mode):
+        answer = query_app(msg, source, verbose_mode)
+        history = history + [(msg, answer)]
+        return history, history
 
-demo = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[goal, diet, meals, avoid, weeks]
-)
+    send_btn.click(user_query, 
+                   inputs=[user_input, chatbot, include_source, verbose], 
+                   outputs=[chatbot, chatbot])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()