Update app.py

app.py

@@ -6,12 +6,14 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Optimize CPU
+import gradio as gr
+import threading
+
 torch.set_num_threads(1)
 
-app = FastAPI(title="SQL Generator API")
+app = FastAPI()
 
-BASE_MODEL = "distilgpt2"
+BASE_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
 print("Loading model...")
 
@@ -22,12 +24,6 @@ model.eval()
 
 print("Model ready")
 
-# ─────────────────────────
-# Request Schema
-# ─────────────────────────
-class Query(BaseModel):
-    text: str
-
 # ─────────────────────────
 # SQL FILTER
 # ─────────────────────────
@@ -40,9 +36,6 @@ SQL_KEYWORDS = [
 def is_sql_related(text):
     return any(k in text.lower() for k in SQL_KEYWORDS)
 
-# ─────────────────────────
-# Generator
-# ─────────────────────────
 SYSTEM_PROMPT = """
 You are an expert SQL generator.
 Only output SQL query.
@@ -50,10 +43,10 @@ Only output SQL query.
 
 def generate_sql(user_input: str):
     if not user_input.strip():
-        return "Empty input."
+        return "Enter SQL question."
 
     if not is_sql_related(user_input):
-        return "Only SQL-related queries allowed."
+        return "Only SQL/database questions allowed."
 
     prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"
 
@@ -62,9 +55,9 @@ def generate_sql(user_input: str):
     with torch.no_grad():
         output = model.generate(
             **inputs,
-            max_new_tokens=80,
-            temperature=0.2,
-            do_sample=True,
+            max_new_tokens=120,
+            temperature=0.1,
+            do_sample=False,
             pad_token_id=tokenizer.eos_token_id,
         )
 
@@ -74,13 +67,30 @@ def generate_sql(user_input: str):
     return result
 
 # ─────────────────────────
-# Routes
+# FastAPI Routes
 # ─────────────────────────
+class Query(BaseModel):
+    text: str
+
 @app.get("/")
 def root():
-    return {"status": "API is running"}
+    return {"status": "API running"}
 
 @app.post("/generate")
 def generate(query: Query):
-    result = generate_sql(query.text)
-    return {"result": result}
+    return {"result": generate_sql(query.text)}
+
+# ─────────────────────────
+# Gradio UI (for testing)
+# ─────────────────────────
+def launch_gradio():
+    demo = gr.Interface(
+        fn=generate_sql,
+        inputs=gr.Textbox(lines=3, label="SQL Question"),
+        outputs=gr.Textbox(lines=6, label="Generated SQL"),
+        title="SQL Generator Test UI"
+    )
+    demo.launch(server_name="0.0.0.0", server_port=7861)
+
+# Run Gradio in parallel thread
+threading.Thread(target=launch_gradio).start()