Update inference.py

2024-01-31 11:23:09 +08:00 · 2024-01-31 11:23:09 +08:00 · 243102a264
parent c167d340ad
commit 243102a264
1 changed files with 1 additions and 2 deletions
--- a/inference.py
+++ b/inference.py
@ -48,8 +48,7 @@ def run_inference(question, prompt_file="prompt.md", metadata_file="metadata.sql
            eos_token_id=eos_token_id,
            pad_token_id=eos_token_id,
        )[0]["generated_text"]
-        .split("# SQL\n")[-1]
-        .split("```")[0]
+        .split("```")[1]
        .split(";")[0]
        .strip()
        + ";"