Update inference.py

This commit is contained in:
Rishabh Srivastava 2024-01-31 11:23:09 +08:00 committed by GitHub
parent c167d340ad
commit 243102a264
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 2 deletions

View File

@ -48,8 +48,7 @@ def run_inference(question, prompt_file="prompt.md", metadata_file="metadata.sql
eos_token_id=eos_token_id, eos_token_id=eos_token_id,
pad_token_id=eos_token_id, pad_token_id=eos_token_id,
)[0]["generated_text"] )[0]["generated_text"]
.split("# SQL\n")[-1] .split("```")[1]
.split("```")[0]
.split(";")[0] .split(";")[0]
.strip() .strip()
+ ";" + ";"