optimized sol, review needed

app.py

@@ -1,19 +1,40 @@
+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"
 
-# Load base + finetuned model
 base_model = "unsloth/Phi-3-mini-4k-instruct-bnb-4bit"
 finetuned_model = "saadkhi/SQL_Chat_finetuned_model"
 
 tokenizer = AutoTokenizer.from_pretrained(base_model)
-model = AutoModelForCausalLM.from_pretrained(base_model)
-model = PeftModel.from_pretrained(model, finetuned_model)
+
+bnb = BitsAndBytesConfig(load_in_4bit=True)
+
+model = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    quantization_config=bnb,
+    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+    device_map="auto"
+)
+
+model = PeftModel.from_pretrained(model, finetuned_model).to(device)
+model.eval()
 
 def chat(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=200)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
+    with torch.inference_mode():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=60,
+            temperature=0.1,
+            do_sample=False
+        )
+
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="SQL Chatbot")
-iface.launch()
+iface.launch()

app_old.txt

@@ -0,0 +1,18 @@
+from gradio_client import Client
+
+# Use the existing Hugging Face Space as the backend
+client = Client("saadkhi/SQL_chatbot_API")
+
+
+def chat(prompt: str) -> str:
+    """Proxy the prompt to the remote Space /chat endpoint."""
+    return client.predict(
+        prompt=prompt,
+        api_name="/chat",
+    )
+
+
+if __name__ == "__main__":
+    # Simple CLI test
+    user_prompt = input("Enter your SQL question: ")
+    print(chat(user_prompt))

requirements.txt

@@ -1,5 +1,6 @@
 gradio
 transformers
 peft
+accelerate
+bitsandbytes
 torch
-bitsandbytes