Update app.py

app.py

@@ -1,94 +1,71 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-from peft import PeftModel
 import torch
-import os
+from unsloth import FastLanguageModel
 
-# Create offload folder (very important!)
-OFFLOAD_DIR = "offload"
-os.makedirs(OFFLOAD_DIR, exist_ok=True)
+# ── Global model (loaded once at startup) ───────────────────────────────
+print("Loading model...")
 
-# Optimal 4-bit quantization config
-quant_config = BitsAndBytesConfig(
+model, tokenizer = FastLanguageModel.from_pretrained(
+    "unsloth/Phi-3-mini-4k-instruct-bnb-4bit",   # very fast pre-quantized base
+    max_seq_length=2048,
+    dtype=None,                                      # auto (bf16/float16)
     load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
 )
 
-# Model paths
-base_model_name = "unsloth/Phi-3-mini-4k-instruct-bnb-4bit"
-lora_model_name = "saadkhi/SQL_Chat_finetuned_model"
-
-print("Loading base model...")
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
-    quantization_config=quant_config,
-    device_map="auto",
-    trust_remote_code=True,
-    offload_folder=OFFLOAD_DIR,          # ← Required fix!
-)
-
-print("Loading LoRA adapter...")
-model = PeftModel.from_pretrained(
-    base_model,
-    lora_model_name,
-    offload_folder=OFFLOAD_DIR,          # ← Required here too!
+# Load your LoRA adapter
+model = FastLanguageModel.for_inference(
+    model.load_adapter("saadkhi/SQL_Chat_finetuned_model")
 )
 
-tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
-
-model.eval()
 print("Model loaded successfully!")
 
-def chat(message, history):
-    # Build conversation in correct Phi-3 format
+# ── Chat function ───────────────────────────────────────────────────────
+def generate_response(message, history):
+    # Build messages list (multi-turn support)
     messages = []
-    for user, assistant in history:
-        messages.append({"role": "user", "content": user})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
+    for user_msg, assistant_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
     messages.append({"role": "user", "content": message})
 
+    # Use the proper chat template (very important for Phi-3)
     inputs = tokenizer.apply_chat_template(
         messages,
         tokenize=True,
         add_generation_prompt=True,
         return_tensors="pt"
-    ).to(model.device)
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Fast generation settings
+    # Generate
     outputs = model.generate(
-        inputs,
-        max_new_tokens=256,
-        temperature=0.7,
-        do_sample=True,
-        top_p=0.9,
-        repetition_penalty=1.1,
+        input_ids=inputs,
+        max_new_tokens=180,               # ← increased but still reasonable
+        temperature=0.0,
+        do_sample=False,                  # greedy = fastest & most deterministic
         use_cache=True,
-        eos_token_id=tokenizer.eos_token_id,
     )
 
-    response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
-    
-    history.append((message, response))
-    return history, ""
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio UI
-with gr.Blocks(title="SQL Chatbot", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# SQL Chat Assistant")
-    gr.Markdown("Fine-tuned Phi-3 Mini (4-bit) for SQL queries. Responses ~3–10s on GPU.")
-    
-    chatbot = gr.Chatbot(height=500)
-    msg = gr.Textbox(
-        label="Your Question",
-        placeholder="e.g., delete duplicate rows from users table based on email",
-        lines=2
-    )
-    clear = gr.Button("Clear")
+    # Clean up output (remove input prompt part)
+    if "<|assistant|>" in response:
+        response = response.split("<|assistant|>")[-1].strip()
+
+    return response
 
-    msg.submit(chat, [msg, chatbot], [chatbot, msg])
-    clear.click(lambda: ([], ""), None, chatbot)
 
-demo.queue(max_size=30)
-demo.launch()
+# ── Gradio UI ───────────────────────────────────────────────────────────
+demo = gr.ChatInterface(
+    fn=generate_response,
+    title="SQL Chat Assistant (Fast Version)",
+    description="Ask SQL related questions • Powered by Phi-3-mini + your fine-tune",
+    examples=[
+        "Write a query to find duplicate emails in users table",
+        "How to delete rows with NULL values in column price?",
+        "Select top 10 most expensive products",
+    ],
+    cache_examples=False,
+)
+
+if __name__ == "__main__":
+    demo.launch()