Update app.py

app.py

@@ -5,32 +5,29 @@ 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")  # 使用半精度以降低内存占用
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")  # 降低内存占用
 
-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}\nThe following SQL query best answers the question `{user_question}`:"
+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)
-    
-    # 生成输出
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=150)
-    
-    # 解码输出
+    outputs = model.generate(**inputs, max_length=150)
     sql_query = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
     return sql_query
 
-# 创建 Gradio 接口
+# 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语句...")
-    sql_output = 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=sql_output)
+    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__":