Update app.py

app.py

@@ -5,7 +5,13 @@ import torch
 # 加载模型和分词器
 model_name = "defog/llama-3-sqlcoder-8b"  # 使用更新的模型以提高性能
 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,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
+    device_map="auto",
+    use_cache=True,
+)
 
 def generate_sql(user_question, instructions, create_table_statements):
     prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>
@@ -20,11 +26,24 @@ The following SQL query best answers the question `{user_question}`:
 ```sql
 """
     
-    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
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    generated_ids = model.generate(
+        **inputs,
+        num_return_sequences=1,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id,
+        max_new_tokens=400,
+        do_sample=False,
+        num_beams=1,
+    )
+    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+
+    torch.cuda.empty_cache()
+    torch.cuda.synchronize()
+    # empty cache so that you do generate more results w/o memory crashing
+    # particularly important on Colab – memory management is much more straightforward
+    # when running on an inference service
+    return sqlparse.format(outputs[0].split("[SQL]")[-1], reindent=True)
 
 
 question = f"What are our top 3 products by revenue in the New York region?"