Update app.py

app.py

@@ -1,22 +1,35 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    BitsAndBytesConfig
+)
 
 MODEL_NAME = "HuggingFaceH4/zephyr-7b-beta"
-OFFLOAD_DIR = "./offload"  # Folder to store offloaded weights
+OFFLOAD_DIR = "./offload"
+
+# Configure 8-bit quantization with CPU fallback
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0,
+    llm_int8_has_fp16_weight=False,
+    llm_int8_enable_fp32_cpu_offload=True   # ✅ allow CPU fallback
+)
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# Load model with 8-bit quantization and offloading
+# Load model with quantization + offloading
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    device_map="auto",       # automatically places layers on CPU/GPU
-    load_in_8bit=True,       # reduce memory usage
-    offload_folder=OFFLOAD_DIR,
-    torch_dtype="auto"       # chooses FP16 if possible
+    device_map="auto",                # spreads across CPU/GPU automatically
+    quantization_config=bnb_config,
+    offload_folder=OFFLOAD_DIR,       # spill to disk if RAM full
+    dtype="auto"
 )
 
-# Create text-generation pipeline
+# Build pipeline
 generator = pipeline(
     "text-generation",
     model=model,
@@ -30,12 +43,8 @@ Speak in a friendly, school-spirited, and enthusiastic tone.
 Always provide interesting facts about WPI when asked questions, and stay in character as Gompei.
 """
 
-# Chat history
-chat_history = []
-
 def chatbot(message, history):
-    global chat_history
-    # Build context with system prompt and previous conversation
+    # Build context with system prompt + conversation history
     context = system_prompt
     for user, bot in history:
         context += f"\nUser: {user}\nGompei: {bot}"
@@ -50,9 +59,6 @@ def chatbot(message, history):
         top_p=0.9
     )
     reply = response[0]["generated_text"].split("Gompei:")[-1].strip()
-
-    # Update history
-    chat_history.append((message, reply))
     return reply
 
 # Gradio interface
@@ -63,5 +69,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    # Disable SSR for faster build
     demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)