From 243102a2645d5278a6397de7b3e206835371c626 Mon Sep 17 00:00:00 2001
From: Rishabh Srivastava <rishabhsriv@gmail.com>
Date: Wed, 31 Jan 2024 11:23:09 +0800
Subject: [PATCH] Update inference.py

---
 inference.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/inference.py b/inference.py
index 2d6c4c1..d18ab7a 100644
--- 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()
         + ";"