Update inference.py

defog-ai · Jan 31, 2024 · 243102a · 243102a
1 parent c167d34
commit 243102a
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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()
         + ";"