Update app.py

app.py

@@ -1,96 +1,58 @@
-# 0. Install custom transformers and imports
 import os
-os.system("pip install git+https://github.com/shumingma/transformers.git")
-os.system("pip install python-docx")
-
-import threading
+import faiss
 import torch
-import torch._dynamo
-torch._dynamo.config.suppress_errors = True
-
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-)
+import threading
 import gradio as gr
-import spaces
-from docx import Document
 
-# 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.
-4. Confirm the pickup time before ending the conversation.
-5. Answer questions about ingredients, preparation, etc.
-6. Handle special requests (allergies, modifications) politely.
-7. Provide calorie information if asked.
-Always be polite, helpful, and ensure the customer feels welcomed and cared for!
-"""
+from docx import Document
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MODEL_ID = "microsoft/bitnet-b1.58-2B-4T"
+# === Configuration ===
+MODEL_ID = "microsoft/phi-2"
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SYSTEM_PROMPT = """You are a friendly café assistant. Help customers place orders, check ingredients, and provide warm service."""
 
-# 2. Load tokenizer and model
+# === Load LLM ===
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to(DEVICE)
 
-print(f"Model loaded on device: {model.device}")
+# === Load Embedder ===
+embedder = SentenceTransformer(EMBED_MODEL)
 
-# 3. Load Menu Text from Word document
-def load_menu_text(docx_path):
+# === Load Menu Text ===
+def load_menu(docx_path):
     doc = Document(docx_path)
-    full_text = []
-    for para in doc.paragraphs:
-        if para.text.strip():
-            full_text.append(para.text.strip())
-    return "\n".join(full_text)
+    return [p.text.strip() for p in doc.paragraphs if p.text.strip()]
+
+menu_chunks = load_menu("menu.docx")
+chunk_embeddings = embedder.encode(menu_chunks, convert_to_tensor=True).cpu().numpy()
 
-MENU_TEXT = load_menu_text("menu.docx")
-print(f"Loaded menu text from Word document.")
+# === Build FAISS Index ===
+dimension = chunk_embeddings.shape[1]
+index = faiss.IndexFlatL2(dimension)
+index.add(chunk_embeddings)
 
-# 4. Simple retrieval function (search inside MENU_TEXT)
-def retrieve_context(question, top_k=3):
-    question = question.lower()
-    sentences = MENU_TEXT.split("\n")
-    matches = [s for s in sentences if any(word in s.lower() for word in question.split())]
-    if not matches:
-        return "Sorry, I couldn't find relevant menu information."
-    return "\n\n".join(matches[:top_k])
+# === Retrieval ===
+def retrieve_context_faiss(query, top_k=3):
+    query_vec = embedder.encode([query]).astype("float32")
+    distances, indices = index.search(query_vec, top_k)
+    return "\n".join([menu_chunks[i] for i in indices[0]])
 
-# 5. Chat respond function
-@spaces.GPU
-def respond(
-    message: str,
-    history: list[tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    context = retrieve_context(message)
-    
+# === Generate LLM Response ===
+def generate_response(message, history, system_message, max_tokens, temperature, top_p):
+    context = retrieve_context_faiss(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": f"{message}\n\nRelevant menu info:\n{context}"})
+    for user, bot in history:
+        messages.append({"role": "user", "content": user})
+        messages.append({"role": "assistant", "content": bot})
+    messages.append({"role": "user", "content": f"{message}\n\nRelevant info:\n{context}"})
 
-    prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
 
-    streamer = TextIteratorStreamer(
-        tokenizer, skip_prompt=True, skip_special_tokens=True
-    )
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         **inputs,
         streamer=streamer,
@@ -99,64 +61,31 @@ def respond(
         top_p=top_p,
         do_sample=True,
     )
+
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        yield response
+    output = ""
+    for token in streamer:
+        output += token
+        yield output
 
-# 6. Gradio ChatInterface
+# === UI ===
 demo = gr.ChatInterface(
-    fn=respond,
-    title="Café Eleven Assistant",
-    description="Friendly café assistant based on real menu loaded from Word document!",
+    fn=generate_response,
+    title="Café Eleven RAG Assistant",
+    description="LLM + FAISS powered café chatbot with real-time Word document lookup.",
     examples=[
-        [
-            "What kinds of burgers do you have?",
-            SYSTEM_PROMPT.strip(),
-            512,
-            0.7,
-            0.95,
-        ],
-        [
-            "Do you have gluten-free pastries?",
-            SYSTEM_PROMPT.strip(),
-            512,
-            0.7,
-            0.95,
-        ],
+        ["Do you have vegetarian options?", SYSTEM_PROMPT, 512, 0.7, 0.9],
+        ["What's in the turkey sandwich?", SYSTEM_PROMPT, 512, 0.7, 0.9],
     ],
     additional_inputs=[
-        gr.Textbox(
-            value=SYSTEM_PROMPT.strip(),
-            label="System message"
-        ),
-        gr.Slider(
-            minimum=1,
-            maximum=2048,
-            value=512,
-            step=1,
-            label="Max new tokens"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=4.0,
-            value=0.7,
-            step=0.1,
-            label="Temperature"
-        ),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)"
-        ),
-    ],
+        gr.Textbox(value=SYSTEM_PROMPT, label="System Prompt"),
+        gr.Slider(1, 1024, 512, label="Max Tokens"),
+        gr.Slider(0.1, 2.0, 0.7, label="Temperature"),
+        gr.Slider(0.1, 1.0, 0.9, label="Top-p"),
+    ]
 )
 
-# 7. Launch
 if __name__ == "__main__":
     demo.launch(share=True)