Update app.py

app.py

@@ -1,118 +1,95 @@
-import os
+# ---------- 0. Imports & constants ----------
+import os, torch, gradio as gr
 from pathlib import Path
-import streamlit as st
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.chains import ConversationalRetrievalChain
-from langchain_community.llms import HuggingFaceHub
-from langchain.memory import ConversationBufferMemory
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import PyMuPDFLoader
+from huggingface_hub import login
 
-# Initialize session state FIRST (before any other Streamlit operations)
-def initialize_session_state():
-    if 'messages' not in st.session_state:
-        st.session_state.messages = [{
-            "role": "assistant",
-            "content": "Hi! Welcome to Café Eleven. What would you like to order today?"
-        }]
+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
 
-# Main application function
-def main():
-    initialize_session_state()
-    
-    # App configuration
-    st.set_page_config(page_title="Café Eleven", page_icon="☕")
-    st.title("☕ Café Eleven Ordering Assistant")
-    
-    # Display chat messages
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.write(message["content"])
-    
-    # Chat processing functions
-    @st.cache_resource
-    def setup_chain():
-        pdf_files = list(Path(".").glob("*.pdf"))
-        if not pdf_files:
-            st.error("No PDF menus found! Please upload some.")
-            st.stop()
-        
-        # Process all PDFs
-        docs = []
-        for pdf in pdf_files:
-            loader = PyMuPDFLoader(str(pdf))
-            docs.extend(loader.load())
-        
-        # Split text
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,
-            chunk_overlap=200
-        )
-        splits = text_splitter.split_documents(docs)
-        
-        # Create embeddings
-        embeddings = HuggingFaceEmbeddings(
-            model_name="sentence-transformers/all-mpnet-base-v2"
-        )
-        vectorstore = FAISS.from_documents(splits, embeddings)
-        
-        # Setup LLM
-        llm = HuggingFaceHub(
-            repo_id="meta-llama/Llama-2-7b-chat-hf",
-            huggingfacehub_api_token=os.environ.get("HF_TOKEN", ""),
+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/drinks from our menu
+3. Answer questions about ingredients, preparation, etc.
+4. Process special requests (allergies, modifications)
+5. Provide a friendly farewell
+
+Always be polite and helpful!"""
+
+WRAPPER_PROMPT = PromptTemplate(
+    "[INST]<<SYS>>\n" + SYSTEM_PROMPT + "\n<</SYS>>\n\n{query_str} [/INST]"
+)
+
+login(token=os.environ["HF_TOKEN"])
+
+# ---------- 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
+
+index = VectorStoreIndex.from_documents(docs)
+
+# ---------- 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",
             model_kwargs={
-                "temperature": 0.2,
-                "max_length": 256
-            }
+                "torch_dtype": torch.float16,
+                "load_in_4bit": True,
+                "use_auth_token": os.environ["HF_TOKEN"]
+            },
+            system_prompt=SYSTEM_PROMPT,
+            query_wrapper_prompt=WRAPPER_PROMPT,
         )
-        
-        # Create conversation chain
-        return ConversationalRetrievalChain.from_llm(
-            llm=llm,
-            retriever=vectorstore.as_retriever(),
-            memory=ConversationBufferMemory(
-                memory_key="chat_history",
-                return_messages=True
-            )
+        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 ----------
+def respond(message, chat_history):
+    if message.lower().strip() in {"quit", "exit", "done"}:
+        return "Thank you for your order! We'll see you soon.", chat_history
+
+    engine = get_chat_engine()
+    response = engine.chat(message).response
+    chat_history.append((message, response))
+    return "", chat_history
+
+with gr.Blocks(title="Café Eleven Chat") as demo:
+    gr.Markdown("## ☕ Café Eleven Ordering Assistant")
+    gr.Markdown("Type your order or question below. Type 'quit' to end the chat.")
     
-    # Handle chat input
-    if prompt := st.chat_input("What would you like to order?"):
-        # Add user message
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        with st.chat_message("user"):
-            st.write(prompt)
-        
-        # Generate response
-        with st.chat_message("assistant"):
-            with st.spinner("Preparing your order..."):
-                try:
-                    chain = setup_chain()
-                    response = chain({"question": prompt})["answer"]
-                    st.write(response)
-                    st.session_state.messages.append({
-                        "role": "assistant",
-                        "content": response
-                    })
-                except Exception as e:
-                    st.error(f"Error: {str(e)}")
+    chatbot = gr.Chatbot(height=500)
+    msg = gr.Textbox(label="Your message", placeholder="Hi, I'd like a latte...")
+    clear = gr.Button("Clear Chat")
     
-    # PDF upload in sidebar
-    with st.sidebar:
-        st.header("Menu Management")
-        uploaded_files = st.file_uploader(
-            "Upload menu PDFs",
-            type="pdf",
-            accept_multiple_files=True
-        )
-        if uploaded_files:
-            for file in uploaded_files:
-                with open(file.name, "wb") as f:
-                    f.write(file.getbuffer())
-            st.success(f"Uploaded {len(uploaded_files)} new menu(s)")
-            st.cache_resource.clear()
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
 
-# Entry point for Hugging Face Spaces
+# For public sharing
 if __name__ == "__main__":
-    main()
+    demo.launch(share=True)