Update app.py

app.py

@@ -47,62 +47,71 @@
 
 
 
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel
-from transformers import BitsAndBytesConfig
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
+from unsloth import FastLanguageModel
+import torch
 
-base_model = "unsloth/Phi-3-mini-4k-instruct-bnb-4bit"
-finetuned_model = "saadkhi/SQL_Chat_finetuned_model"
+# Load model once at startup — Unsloth makes it 2.5x faster
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name="unsloth/Phi-3-mini-4k-instruct-bnb-4bit",
+    max_seq_length=4096,
+    dtype=None,  # Auto detect (bfloat16 if supported)
+    load_in_4bit=True,
+)
 
-tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+# Load your fine-tuned LoRA adapter
+model = FastLanguageModel.get_peft_model(
+    model,
+    "saadkhi/SQL_Chat_finetuned_model",  # Your HF repo
+)
 
-bnb_config = BitsAndBytesConfig(load_in_4bit=True)
+# Enable fast inference mode (critical for speed!)
+FastLanguageModel.for_inference(model)
 
-model = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    quantization_config=bnb_config,
-    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
-    device_map="auto",
-    trust_remote_code=True,
-)
-model = PeftModel.from_pretrained(model, finetuned_model)
-model.eval()
+def chat(message, history):
+    # Build proper Phi-3 chat format
+    messages = []
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
+    messages.append({"role": "user", "content": message})
 
-def chat(user_prompt):
-    # Proper Phi-3 chat format
-    messages = [{"role": "user", "content": user_prompt}]
-    
+    # Apply chat template and tokenize
     inputs = tokenizer.apply_chat_template(
         messages,
         tokenize=True,
         add_generation_prompt=True,
         return_tensors="pt"
-    ).to(device)
-    
-    with torch.inference_mode():
-        outputs = model.generate(
-            inputs,
-            max_new_tokens=256,      # Increased a bit for full SQL
-            temperature=0.7,
-            do_sample=True,          # Better for creativity, faster
-            top_p=0.9,
-            repetition_penalty=1.1,
-        )
+    ).to(model.device)
+
+    # Generate fast
+    output = model.generate(
+        input_ids=inputs,
+        max_new_tokens=256,
+        temperature=0.7,
+        do_sample=True,
+        top_p=0.9,
+        use_cache=True,
+        repetition_penalty=1.1,
+    )
+
+    # Decode only the new part
+    response = tokenizer.decode(output[0][inputs.shape[-1]:], skip_special_tokens=True)
+    history.append((message, response))
+    return history, ""
+
+# Clean Gradio Chat Interface
+with gr.Blocks(title="SQL Chatbot", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# SQL Chat Assistant")
+    gr.Markdown("Ask any SQL-related question. Fast responses powered by fine-tuned Phi-3 Mini.")
     
-    # Clean response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    response = response.split("<|assistant|>")[-1].split("<|end|>")[0].strip()
-    
-    return response
+    chatbot = gr.Chatbot(height=500)
+    msg = gr.Textbox(label="Your Message", placeholder="e.g., delete duplicate rows from users table", lines=2)
+    clear = gr.Button("Clear")
 
-iface = gr.ChatInterface(
-    fn=chat,
-    title="Fast SQL Chatbot",
-    description="Ask SQL questions (e.g., 'delete duplicate rows based on email')"
-)
+    msg.submit(chat, [msg, chatbot], [chatbot, msg])
+    clear.click(lambda: ([], ""), None, chatbot)
 
-iface.launch()
+demo.queue(max_size=20)  # Handle multiple users smoothly
+demo.launch()