Update app.py

app.py

@@ -1,18 +1,15 @@
-import torch
-import gradio as gr
-import os
+# ---------- 0. Imports & constants ----------
+import os, torch, gradio as gr
 from pathlib import Path
 from huggingface_hub import login
-from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+
+from llama_index.core import (
+    VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
+)
 from llama_index.core.memory import ChatMemoryBuffer
 from llama_index.llms.huggingface import HuggingFaceLLM
 from llama_index.embeddings.langchain import LangchainEmbedding
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-from llama_index.core import PromptTemplate
-
-
-pdf_files = [str(p) for p in Path(".").glob("*.pdf")]   
-docs = SimpleDirectoryReader(input_files=pdf_files).load_data()
 
 SYSTEM_PROMPT = """
 You are a friendly café assistant for Café Eleven. Your job is to:
@@ -21,7 +18,6 @@ You are a friendly café assistant for Café Eleven. Your job is to:
 3. Ask for pickup time
 4. Suggest add-ons/extras from our menu
 5. Confirm the complete order
-
 Menu items are embedded in the document. Always:
 - Be polite and professional
 - Confirm order details clearly
@@ -29,66 +25,76 @@ Menu items are embedded in the document. Always:
 - Never make up items not in our menu
 """
 
-wrapper_prompt = PromptTemplate(
+WRAPPER_PROMPT = PromptTemplate(
     """<s>[INST] <<SYS>>
 {system_prompt}
 Current conversation:
 {chat_history}
 <</SYS>>
-
 {query_str} [/INST]"""
 )
 
-login(token=os.environ["HF_TOKEN"]) 
-
-llm = HuggingFaceLLM(
-    tokenizer_name="meta-llama/Llama-2-7b-chat-hf",
-    model_name="meta-llama/Llama-2-7b-chat-hf",
-    context_window=3900,
-    max_new_tokens=256,
-    generate_kwargs={"temperature": 0.2, "do_sample": True},
-    device_map="auto",
-    model_kwargs={
-        "torch_dtype": torch.float16,
-        "load_in_4bit": True,
-        "use_auth_token": os.environ["HF_TOKEN"]  
-    },
-    system_prompt=SYSTEM_PROMPT,
-    query_wrapper_prompt=wrapper_prompt,
-)
+login(token=os.environ["HF_TOKEN"])
 
-embed_model = LangchainEmbedding(
-    HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-)
-
-Settings.llm = llm
+# ---------- 1. Pre-load documents & build the vector index (CPU-safe) ----------
+docs         = SimpleDirectoryReader(
+                  input_files=[str(p) for p in Path(".").glob("*.pdf")]
+               ).load_data()
+embed_model  = LangchainEmbedding(
+                  HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+               )
 Settings.embed_model = embed_model
-Settings.chunk_size = 512
+Settings.chunk_size  = 512
 
-# ---------- 3. Build the chat engine ----------
-memory = ChatMemoryBuffer.from_defaults(token_limit=2000)
 index = VectorStoreIndex.from_documents(docs)
-chat_engine = index.as_chat_engine(
-    chat_mode="condense_plus_context",
-    memory=memory,
-    system_prompt=SYSTEM_PROMPT,
-)
 
-# ---------- 4. Gradio UI ----------
+# ---------- 2. Lazy, singleton chat-engine ----------
+_state = {"chat_engine": None}          # filled on first request
+
+def get_chat_engine():
+    if _state["chat_engine"] is None:
+        llm = HuggingFaceLLM(
+            tokenizer_name="meta-llama/Llama-2-7b-chat-hf",
+            model_name   ="meta-llama/Llama-2-7b-chat-hf",
+            context_window=3900,
+            max_new_tokens=256,
+            generate_kwargs={"temperature":0.2, "do_sample":True},
+            device_map   ="auto",          # CUDA now visible
+            model_kwargs ={
+                "torch_dtype": torch.float16,
+                "load_in_4bit": True,
+                "use_auth_token": os.environ["HF_TOKEN"]
+            },
+            system_prompt       = SYSTEM_PROMPT,
+            query_wrapper_prompt= WRAPPER_PROMPT,
+        )
+        Settings.llm = llm
+
+        memory = ChatMemoryBuffer.from_defaults(token_limit=2000)
+        _state["chat_engine"] = index.as_chat_engine(
+            chat_mode="condense_plus_context",
+            memory     =memory,
+            system_prompt=SYSTEM_PROMPT,
+        )
+    return _state["chat_engine"]
+
+# ---------- 3. Gradio UI ----------
 with gr.Blocks(title="Café Eleven Chat") as demo:
     gr.Markdown("## ☕ Café Eleven Ordering Assistant  \nType *quit* to end the chat.")
     chatbot = gr.Chatbot()
-    user_txt = gr.Textbox(show_label=False, placeholder="Hi, I’d like a latte…")
-    clear = gr.Button("Clear")
+    user_in = gr.Textbox(show_label=False, placeholder="Hi, I’d like a latte…")
+    clear   = gr.Button("Clear")
 
-    def respond(message, chat_history):
+    def respond(message, history):
         if message.lower().strip() in {"quit", "exit", "done"}:
-            return "Thank you for your order! We'll see you soon.", chat_history
-        response = chat_engine.chat(message).response
-        chat_history.append((message, response))
-        return "", chat_history
+            return "Thank you for your order! We'll see you soon.", history
+
+        engine   = get_chat_engine()        # GPU & model ready
+        reply    = engine.chat(message).response
+        history.append((message, reply))
+        return "", history
 
-    user_txt.submit(respond, [user_txt, chatbot], [user_txt, chatbot])
+    user_in.submit(respond, [user_in, chatbot], [user_in, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":