Update app.py

app.py

@@ -1,83 +1,75 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import gradio as gr
 
-# Model Names
-BASE_MODEL_NAME = "microsoft/phi-2"
-ADAPTER_REPO = "Shriti09/Microsoft-Phi-QLora"
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load tokenizer and model
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
-tokenizer.pad_token = tokenizer.eos_token
+# Base model and adapter paths
+base_model_name = "microsoft/phi-2"  # Pull from HF Hub directly
+adapter_path = "./phi2-qlora-adapter"  # Your uploaded adapter folder in Space repo
 
-print("Loading base model...")
-base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_NAME, device_map="auto", torch_dtype=torch.float16)
+print("🔧 Loading base model...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_name,
+    device_map="auto",
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+)
 
-print("Loading LoRA adapter...")
-model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)
+print("🔧 Loading LoRA adapter...")
+adapter_model = PeftModel.from_pretrained(base_model, adapter_path)
 
-# Merge adapter into the base model
-model = model.merge_and_unload()
-model.eval()
+print("🔗 Merging adapter into base model...")
+merged_model = adapter_model.merge_and_unload()
+merged_model.eval()
 
-# Function to generate responses
-def generate_response(message, chat_history, temperature, top_p, max_tokens):
-    # Combine history with the new message
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+print("✅ Model ready for inference!")
+
+# Chat function with history
+def chat_fn(message, history):
+    # Combine conversation history into one prompt
     full_prompt = ""
-    for user_msg, bot_msg in chat_history:
+    for user_msg, bot_msg in history:
         full_prompt += f"User: {user_msg}\nAI: {bot_msg}\n"
     full_prompt += f"User: {message}\nAI:"
 
-    # Tokenize and generate
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_length=len(inputs["input_ids"][0]) + max_tokens,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    
-    # Decode and extract the AI response
+    # Tokenize inputs
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+
+    with torch.no_grad():
+        outputs = merged_model.generate(
+            **inputs,
+            max_new_tokens=150,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
+        )
+
+    # Decode and return only the AI's latest response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # Only return the new part of the response
     response = response.split("AI:")[-1].strip()
 
-    # Update history
-    chat_history.append((message, response))
-    return chat_history, chat_history
+    # Append to history
+    history.append((message, response))
+    return history, history
 
-# Gradio UI with Blocks
-with gr.Blocks() as demo:
-    gr.Markdown("<h1><center>🤖 Phi-2 QLoRA Chatbot</center></h1>")
-    gr.Markdown("Chat with Microsoft Phi-2 fine-tuned using QLoRA adapters!")
+# Gradio UI
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1>🧠 Phi-2 QLoRA Chatbot</h1>")
 
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Ask me something...", label="Your Message")
-    clear = gr.Button("🗑️ Clear Chat")
+    message = gr.Textbox(label="Your message:")
+    clear = gr.Button("Clear chat")
 
-    # Add sliders for controlling generation behavior
-    with gr.Row():
-        temp_slider = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature")
-        top_p_slider = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="Top-p (nucleus sampling)")
-        max_tokens_slider = gr.Slider(64, 1024, value=256, step=64, label="Max Tokens")
-
-    # State to hold chat history
     state = gr.State([])
 
-    # On send message
-    def on_message(message, history, temperature, top_p, max_tokens):
-        return generate_response(message, history, temperature, top_p, max_tokens)
-
-    # Button actions
-    msg.submit(on_message, 
-               [msg, state, temp_slider, top_p_slider, max_tokens_slider], 
-               [chatbot, state])
-    
-    clear.click(lambda: ([], []), None, [chatbot, state])
+    message.submit(chat_fn, [message, state], [chatbot, state])
+    clear.click(lambda: [], None, chatbot)
+    clear.click(lambda: [], None, state)
 
-# Launch the Gradio app
-demo.launch()
+# Run with queue for multiple users
+demo.queue(concurrency_count=2).launch()