Update app.py

app.py

@@ -1,103 +1,47 @@
-# app.py - Fixed for recent Gradio versions (no allow_flagging)
-
 import torch
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from peft import PeftModel
-
-# ────────────────────────────────────────────────────────────────
-#   Fastest practical configuration
-# ────────────────────────────────────────────────────────────────
-
-BASE_MODEL = "unsloth/Phi-3-mini-4k-instruct-bnb-4bit"
-LORA_PATH  = "saadkhi/SQL_Chat_finetuned_model"
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-MAX_NEW_TOKENS = 180
-TEMPERATURE    = 0.0          # greedy = fastest
-DO_SAMPLE      = False
+MODEL_ID = "saadkhi/SQL_Chat_finetuned_model"
 
-# ────────────────────────────────────────────────────────────────
-#   4-bit quantization (very important for speed)
-# ────────────────────────────────────────────────────────────────
+app = FastAPI(title="SQL Chatbot API")
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit              = True,
-    bnb_4bit_quant_type       = "nf4",
-    bnb_4bit_use_double_quant = True,
-    bnb_4bit_compute_dtype    = torch.bfloat16
-)
-
-print("Loading quantized base model...")
+# Load model once (on startup)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(
-    BASE_MODEL,
-    quantization_config = bnb_config,
-    device_map          = "auto",
-    trust_remote_code   = True,
-    torch_dtype         = torch.bfloat16
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-print("Loading LoRA adapters...")
-model = PeftModel.from_pretrained(model, LORA_PATH)
-
-# Merge LoRA into base model → much faster inference
-model = model.merge_and_unload()
+class QueryRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 256
 
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+class QueryResponse(BaseModel):
+    response: str
 
-model.eval()
-print("Model ready!")
 
-# ────────────────────────────────────────────────────────────────
-def generate_sql(prompt: str):
-    messages = [{"role": "user", "content": prompt}]
-    
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
+@app.post("/generate", response_model=QueryResponse)
+def generate_answer(request: QueryRequest):
+    inputs = tokenizer(
+        request.prompt,
         return_tensors="pt"
     ).to(model.device)
 
-    with torch.inference_mode():
-        outputs = model.generate(
-            input_ids      = inputs,
-            max_new_tokens = MAX_NEW_TOKENS,
-            temperature    = TEMPERATURE,
-            do_sample      = DO_SAMPLE,
-            use_cache      = True,
-            pad_token_id   = tokenizer.eos_token_id,
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=request.max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
         )
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Clean output
-    if "<|assistant|>" in response:
-        response = response.split("<|assistant|>", 1)[-1].strip()
-    response = response.split("<|end|>")[0].strip() if "<|end|>" in response else response
-
-    return response
-
-# ────────────────────────────────────────────────────────────────
-#   Gradio interface - modern style (no allow_flagging)
-# ────────────────────────────────────────────────────────────────
-
-demo = gr.Interface(
-    fn=generate_sql,
-    inputs=gr.Textbox(
-        label="Ask SQL related question",
-        placeholder="Show me all employees with salary > 50000...",
-        lines=3
-    ),
-    outputs=gr.Textbox(label="Generated SQL / Answer"),
-    title="SQL Chatbot - Optimized",
-    description="Phi-3-mini 4bit + LoRA merged",
-    examples=[
-        ["Find duplicate emails in users table"],
-        ["Top 5 highest paid employees"],
-        ["Count orders per customer last month"]
-    ],
-    # flag button is disabled by default in newer versions → no need for allow_flagging
-)
+    output_text = tokenizer.decode(
+        output_ids[0],
+        skip_special_tokens=True
+    )
 
-if __name__ == "__main__":
-    demo.launch()
+    return {"response": output_text}