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


demo = gr.Interface(
    fn=generate_sql,
    inputs=[
        gr.Markdown("## SQL Query Generator"),
        gr.Textbox(label="User Question", placeholder="请输入您的问题...", value="从纽约的客户那里获得的总收入是多少？"),
        gr.Textbox(label="Create Table Statements", placeholder="请输入DDL语句...", value="CREATE TABLE customers (id INT, city VARCHAR(50), revenue DECIMAL);"),
        gr.Textbox(label="Instructions (可选)", placeholder="请输入额外说明...", value="")
    ],
    outputs="text",
)


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