Update app.py

app.py

@@ -1,6 +1,7 @@
-# 0. Install custom transformers and import packages
+# 0. Install custom transformers and imports
 import os
 os.system("pip install git+https://github.com/shumingma/transformers.git")
+os.system("pip install sentence-transformers")
 
 import threading
 import torch
@@ -12,11 +13,15 @@ from transformers import (
     AutoTokenizer,
     TextIteratorStreamer,
 )
+from sentence_transformers import SentenceTransformer
 import gradio as gr
 import spaces
 
-# 1. System prompt for Café Eleven
+from pathlib import Path
+
+# 1. System prompt
 SYSTEM_PROMPT = """
+You are a friendly café assistant for Café Eleven. Your job is to:
 1. Greet the customer warmly.
 2. Help them order food and drinks from our menu.
 3. Ask the customer for their desired pickup time.
@@ -39,7 +44,37 @@ model = AutoModelForCausalLM.from_pretrained(
 
 print(f"Model loaded on device: {model.device}")
 
-# 3. Response generation with streaming
+# 3. Load PDF files and create simple document store
+from PyPDF2 import PdfReader
+
+# Read all PDFs into a list of small chunks
+def load_pdfs(folder_path="."):
+    docs = []
+    for pdf_file in Path(folder_path).glob("*.pdf"):
+        reader = PdfReader(str(pdf_file))
+        for page in reader.pages:
+            text = page.extract_text()
+            if text:
+                for para in text.split("\n\n"):
+                    if len(para.strip()) > 20:  # keep meaningful text
+                        docs.append(para.strip())
+    return docs
+
+document_chunks = load_pdfs(".")
+print(f"Loaded {len(document_chunks)} text chunks from PDFs.")
+
+# 4. Create embeddings
+embedder = SentenceTransformer("all-MiniLM-L6-v2")  # Fast small model
+doc_embeddings = embedder.encode(document_chunks, normalize_embeddings=True)
+
+# 5. Retrieval function
+def retrieve_context(question, top_k=3):
+    question_embedding = embedder.encode(question, normalize_embeddings=True)
+    scores = torch.tensor(doc_embeddings) @ torch.tensor(question_embedding)
+    top_indices = torch.topk(scores, k=min(top_k, len(scores))).indices.tolist()
+    return "\n\n".join([document_chunks[idx] for idx in top_indices])
+
+# 6. Chat respond function
 @spaces.GPU
 def respond(
     message: str,
@@ -49,13 +84,15 @@ def respond(
     temperature: float,
     top_p: float,
 ):
+    context = retrieve_context(message)
+    
     messages = [{"role": "system", "content": system_message}]
     for user_msg, bot_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
         if bot_msg:
             messages.append({"role": "assistant", "content": bot_msg})
-    messages.append({"role": "user", "content": message})
+    messages.append({"role": "user", "content": f"{message}\n\nRelevant menu info:\n{context}"})
 
     prompt = tokenizer.apply_chat_template(
         messages, tokenize=False, add_generation_prompt=True
@@ -81,21 +118,21 @@ def respond(
         response += new_text
         yield response
 
-# 4. Gradio ChatInterface setup
+# 7. Gradio UI
 demo = gr.ChatInterface(
     fn=respond,
     title="Café Eleven Assistant",
-    description="A friendly café chatbot to help you with orders and menu questions!",
+    description="Friendly café assistant with real menu knowledge!",
     examples=[
         [
-            "Can I get a recommendation for breakfast?",
+            "What kinds of burgers do you have?",
             SYSTEM_PROMPT.strip(),
             512,
             0.7,
             0.95,
         ],
         [
-            "Do you have vegan menu options?",
+            "Do you have any gluten-free pastries?",
             SYSTEM_PROMPT.strip(),
             512,
             0.7,
@@ -131,6 +168,6 @@ demo = gr.ChatInterface(
     ],
 )
 
-# 5. Launch normally
+# 8. Launch
 if __name__ == "__main__":
     demo.launch()