Update app.py

app.py

@@ -1,63 +1,95 @@
 import os
-import faiss
-import torch
+os.system("pip install git+https://github.com/shumingma/transformers.git")
+os.system("pip install python-docx")
+
 import threading
-import gradio as gr
+import torch
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
 
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+)
+import gradio as gr
+import spaces
 from docx import Document
-from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from spaces import GPU
 
-# === 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."""
 
-# === Load LLM ===
+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!
+"""
+
+MODEL_ID = "microsoft/bitnet-b1.58-2B-4T"
+
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to(DEVICE)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
 
-# === Load Embedder ===
-embedder = SentenceTransformer(EMBED_MODEL)
+print(f"Model loaded on device: {model.device}")
 
-# === Load Menu Text ===
-def load_menu(docx_path):
-    doc = Document(docx_path)
-    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()
+def load_menu_text(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)
 
-# === Build FAISS Index ===
-dimension = chunk_embeddings.shape[1]
-index = faiss.IndexFlatL2(dimension)
-index.add(chunk_embeddings)
+MENU_TEXT = load_menu_text("menu.docx")
+print(f"Loaded menu text from Word document.")
 
-# === 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]])
 
+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])
 
-# === Generate LLM Response ===
-@spaces.GPU  # Only if you're using ZeroGPU
-def generate_response(message, history, system_message, max_tokens, temperature, top_p):
-    context = retrieve_context_faiss(message)
 
+@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)
+    
     messages = [{"role": "system", "content": system_message}]
     for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
-    
-    messages.append({"role": "user", "content": f"{message}\n\nRelevant info:\n{context}"})
-    
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        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}"})
+
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
     inputs = tokenizer(prompt, return_tensors="pt").to(model.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,
@@ -66,32 +98,64 @@ def generate_response(message, history, system_message, max_tokens, temperature,
         top_p=top_p,
         do_sample=True,
     )
-
     thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
     thread.start()
 
-    output = ""
-    for token in streamer:
-        output += token
-        yield output
-        
-print("Inputs received:", message, history, system_message, max_tokens, temperature, top_p)
-# === UI ===
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
+
+
 demo = gr.ChatInterface(
-    fn=generate_response,
-    title="Café Eleven RAG Assistant",
-    description="LLM + FAISS powered café chatbot with real-time Word document lookup.",
+    fn=respond,
+    title="Café Eleven Assistant",
+    description="Friendly café assistant based on real menu loaded from Word document!",
     examples=[
-        ["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],
+        [
+            "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,
+        ],
     ],
     additional_inputs=[
-        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"),
-    ]
+        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)"
+        ),
+    ],
 )
 
+
 if __name__ == "__main__":
     demo.launch(share=True)