shankardev07/Context-Aware-Text-to-SQL

1import gradio as gr2import torch3from transformers import AutoTokenizer, AutoModelForCausalLM4from peft import PeftModel5 6print("Loading Model for Web UI (CPU Mode)...")7base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"8 9# ✅ FIXED: Using "." because files are in the same folder as app.py10tokenizer = AutoTokenizer.from_pretrained(".")11 12# Load Base Model without 4-bit quantization (CPU friendly)13base_model = AutoModelForCausalLM.from_pretrained(14    base_model_name,15    device_map="cpu",16    torch_dtype=torch.float32 17)18 19# ✅ FIXED: Using "." here as well20model = PeftModel.from_pretrained(base_model, ".")21device = "cpu"22 23def generate_sql(schema, question):24    prompt = f"""### Instruction:25Generate ONLY SQL query.26 27### Schema:28{schema}29 30### Question:31{question}32 33### SQL:34"""35    inputs = tokenizer(prompt, return_tensors="pt")36    inputs = {k: v.to(device) for k, v in inputs.items()}37 38    outputs = model.generate(39        **inputs,40        max_new_tokens=120,41        temperature=0.1,42        do_sample=False,43        pad_token_id=tokenizer.eos_token_id44    )45 46    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)47    generated_sql = generated.split("### SQL:")[-1].strip()48    49    return generated_sql50 51# Web UI52with gr.Blocks(theme=gr.themes.Soft()) as demo:53    gr.Markdown(54        """55        # 🚀 Context-Aware Text-to-SQL Generator56        **Model:** TinyLlama-1.1B + QLoRA (Fine-tuned) | **Accuracy:** 37.0%57        58        *Enter your database schema (CREATE TABLE) and ask a question in plain English to generate the SQL query.*59        """60    )61    62    with gr.Row():63        with gr.Column():64            schema_input = gr.Textbox(65                lines=8, 66                label="Database Schema (CREATE TABLE context)", 67                placeholder="CREATE TABLE employees (id INT, name VARCHAR, department VARCHAR, salary INT);"68            )69            question_input = gr.Textbox(70                lines=2, 71                label="Your Question", 72                placeholder="Find the names of all employees in the IT department who earn more than 50000."73            )74            submit_btn = gr.Button("Generate SQL ⚡", variant="primary")75            76        with gr.Column():77            sql_output = gr.Code(label="Generated SQL Output", language="sql")78            79    submit_btn.click(fn=generate_sql, inputs=[schema_input, question_input], outputs=sql_output)80    81    gr.Examples(82        examples=[83            [84                "CREATE TABLE stadium (stadium_id INT, capacity INT); CREATE TABLE singer (stadium_id INT, singer_name VARCHAR);", 85                "How many singers performed in a stadium with a capacity greater than 5000?"86            ]87        ],88        inputs=[schema_input, question_input]89    )90 91demo.launch()