app.py

import warnings
warnings.filterwarnings("ignore")

import torch
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM

# Optimize CPU
torch.set_num_threads(1)

app = FastAPI(title="SQL Generator API")

BASE_MODEL = "distilgpt2"

print("Loading model...")

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)

model.eval()

print("Model ready")

# ─────────────────────────
# Request Schema
# ─────────────────────────
class Query(BaseModel):
    text: str

# ─────────────────────────
# SQL FILTER
# ─────────────────────────
SQL_KEYWORDS = [
    "sql", "database", "table", "select", "insert",
    "update", "delete", "join", "group by",
    "postgres", "mysql", "sqlite", "query"
]

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.
"""

def generate_sql(user_input: str):
    if not user_input.strip():
        return "Empty input."

    if not is_sql_related(user_input):
        return "Only SQL-related queries allowed."

    prompt = f"{SYSTEM_PROMPT}\nUser: {user_input}\nSQL:"

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=80,
            temperature=0.2,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    text = tokenizer.decode(output[0], skip_special_tokens=True)
    result = text.split("SQL:")[-1].strip().split("\n")[0]

    return result

# ─────────────────────────
# Routes
# ─────────────────────────
@app.get("/")
def root():
    return {"status": "API is running"}

@app.post("/generate")
def generate(query: Query):
    result = generate_sql(query.text)
    return {"result": result}