Update app.py

app.py

@@ -1,18 +1,21 @@
-# ---------- 0. Imports & constants ----------
+# 0. Install custom transformers and imports
 import os
+os.system("pip install git+https://github.com/shumingma/transformers.git")
+
+import threading
 import torch
-import gradio as gr
-from pathlib import Path
-from huggingface_hub import login
+import torch._dynamo
+torch._dynamo.config.suppress_errors = True
 
-from llama_index.core import (
-    VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
 )
-from llama_index.core.memory import ChatMemoryBuffer
-from llama_index.llms.huggingface import HuggingFaceLLM
-from llama_index.embeddings.langchain import LangchainEmbedding
-from langchain_huggingface import HuggingFaceEmbeddings
+import gradio as gr
+import spaces
 
+# 1. System prompt (your original one)
 SYSTEM_PROMPT = """
 You are a friendly café assistant for Café Eleven. Your job is to:
 1. Greet the customer warmly
@@ -24,70 +27,114 @@ You are a friendly café assistant for Café Eleven. Your job is to:
 Always be polite and helpful!
 """
 
-WRAPPER_PROMPT = PromptTemplate(
-    "[INST]<<SYS>>\n" + SYSTEM_PROMPT + "\n<</SYS>>\n\n{query_str} [/INST]"
-)
+# 2. Model info
+MODEL_ID = "microsoft/bitnet-b1.58-2B-4T"
 
-# ---------- 1. Login & Load Data ----------
-login(token=os.environ["HF_TOKEN"])
-
-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")
+# 3. Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
 )
-Settings.embed_model = embed_model
-Settings.chunk_size = 512
-
-index = VectorStoreIndex.from_documents(docs)
 
-# ---------- 2. Initialize Chat Engine ----------
-_state = {"chat_engine": None}
+print(f"Model loaded on device: {model.device}")
 
-def get_chat_engine():
-    if _state["chat_engine"] is None:
-        llm = HuggingFaceLLM(
-            tokenizer_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-            model_name="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
-            context_window=2048,
-            max_new_tokens=256,
-            generate_kwargs={"temperature": 0.2, "do_sample": True},
-            device_map="auto",
-            model_kwargs={
-                "use_auth_token": os.environ["HF_TOKEN"]
-                        },
-            system_prompt=SYSTEM_PROMPT,
-            query_wrapper_prompt=WRAPPER_PROMPT,)
-        Settings.llm = llm
+# 4. Respond function
+@spaces.GPU
+def respond(
+    message: str,
+    history: list[tuple[str, str]],
+    system_message: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    """
+    Generate a chat response using streaming with TextIteratorStreamer.
+    """
+    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})
 
-        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"]
+    prompt = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-# ---------- 3. Simple Chat Function ----------
-def chat_with_cafe_eleven(message: str) -> str:
-    if message.lower().strip() in {"quit", "exit", "done"}:
-        return "Thank you for your order! We'll see you soon."
+    streamer = TextIteratorStreamer(
+        tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+    generate_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
-    engine = get_chat_engine()
-    response = engine.chat(message).response
-    return response
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
 
-# ---------- 4. Gradio UI ----------
-iface = gr.Interface(
-    fn=chat_with_cafe_eleven,
-    inputs=gr.Textbox(lines=2, placeholder="Ask about menu items, orders, etc..."),
-    outputs="text",
+# 5. Gradio UI
+demo = gr.ChatInterface(
+    fn=respond,
     title="Café Eleven Assistant",
-    description="A friendly café assistant to help you with orders and questions!"
+    description="A friendly café chatbot to help you with orders and menu questions!",
+    examples=[
+        [
+            "Can I get a recommendation for breakfast?",
+            SYSTEM_PROMPT.strip(),
+            512,
+            0.7,
+            0.95,
+        ],
+        [
+            "Do you have vegan menu options?",
+            SYSTEM_PROMPT.strip(),
+            512,
+            0.7,
+            0.95,
+        ],
+    ],
+    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)"
+        ),
+    ],
 )
 
-# ---------- 5. Launch App ----------
+# 6. Launch
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)