import gradio as gr
from transformers import pipeline

# 加载 SQLCoder 模型
sqlcoder = pipeline("text2text-generation", model="defog/sqlcoder-7b")

def generate_sql(user_question, create_table_statements):
    prompt = f"Generate a SQL query to answer this question: \"{user_question}\"\nDDL statements:\n{create_table_statements}\n"
    result = sqlcoder(prompt)
    return result[0]['generated_text']

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("## SQL Query Generator")
    user_question = gr.Textbox(label="User Question", placeholder="请输入您的问题...")
    create_table_statements = gr.Textbox(label="DDL Statements", placeholder="请输入表的 DDL 语句...")
    output_sql = gr.Textbox(label="Generated SQL Query", interactive=False)

    submit_btn = gr.Button("Generate SQL")
    submit_btn.click(generate_sql, inputs=[user_question, create_table_statements], outputs=output_sql)


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