File size: 2,103 Bytes
3741f67 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import argparse
def generate_prompt(question, prompt_file="prompt.md", metadata_file="metadata.sql"):
with open(prompt_file, "r") as f:
prompt = f.read()
with open(metadata_file, "r") as f:
table_metadata_string = f.read()
prompt = prompt.format(
user_question=question, table_metadata_string=table_metadata_string
)
return prompt
def get_tokenizer_model(model_name):
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
trust_remote_code=True,
torch_dtype=torch.float16,
device_map="auto",
use_cache=True,
)
return tokenizer, model
def run_inference(question, prompt_file="prompt.md", metadata_file="metadata.sql"):
tokenizer, model = get_tokenizer_model("defog/sqlcoder")
prompt = generate_prompt(question, prompt_file, metadata_file)
# make sure the model stops generating at triple ticks
eos_token_id = tokenizer.convert_tokens_to_ids(["```"])[0]
pipe = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
max_new_tokens=300,
do_sample=False,
num_beams=5, # do beam search with 5 beams for high quality results
)
generated_query = (
pipe(
prompt,
num_return_sequences=1,
eos_token_id=eos_token_id,
pad_token_id=eos_token_id,
)[0]["generated_text"]
.split("```sql")[-1]
.split("```")[0]
.split(";")[0]
.strip()
+ ";"
)
return generated_query
if __name__ == "__main__":
# Parse arguments
parser = argparse.ArgumentParser(description="Run inference on a question")
parser.add_argument("-q","--question", type=str, help="Question to run inference on")
args = parser.parse_args()
question = args.question
print("Loading a model and generating a SQL query for answering your question...")
print(run_inference(question)) |