import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# 加载模型和分词器
model_name = "defog/sqlcoder-7b-2"  # 使用更新的模型以提高性能
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")  # 降低内存占用

def generate_sql(user_question, create_table_statements, instructions=""):
    prompt = f"Generate a SQL query to answer this question: `{user_question}`\n{instructions}\n\nDDL statements:\n{create_table_statements}\n<|eot_id|>"
    
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_length=150)
    sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    return sql_query

# Gradio 接口
with gr.Blocks() as demo:
    gr.Markdown("## SQL Query Generator")
    
    user_question = gr.Textbox(label="User Question", placeholder="请输入您的问题...", value="从纽约的客户那里获得的总收入是多少？")
    create_table_statements = gr.Textbox(label="Create Table Statements", placeholder="请输入DDL语句...", value="CREATE TABLE customers (id INT, city VARCHAR(50), revenue DECIMAL);")
    instructions = gr.Textbox(label="Instructions (可选)", placeholder="请输入额外说明...", value="")
    
    submit_btn = gr.Button("生成 SQL 查询")
    output = gr.Textbox(label="生成的 SQL 查询")
    
    submit_btn.click(generate_sql, inputs=[user_question, create_table_statements, instructions], outputs=output)


if __name__ == "__main__":
    demo.launch(share=True)