text_plain_output_1.png
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

from typing import List

tokenizer = AutoTokenizer.from_pretrained('juierror/flan-t5-text2sql-with-schema')
model = AutoModelForSeq2SeqLM.from_pretrained('juierror/flan-t5-text2sql-with-schema')

def prepare_input(question: str, table: List[str]):
    table_prefix = 'table:'
    question_prefix = 'question:'
    join_table = ','.join(table)
    inputs = f'{question_prefix} {question} {table_prefix} {join_table}'
    input_ids = tokenizer(inputs, max_length=700, return_tensors='pt').input_ids
    return input_ids

def inference(question: str, table: List[str]) -> str:
    input_data = prepare_input(question=question, table=table)
    input_data = input_data.to(model.device)
    outputs = model.generate(inputs=input_data, num_beams=10, top_k=10, max_length=700)
    result = tokenizer.decode(token_ids=outputs[0], skip_special_tokens=True)
    return result

def infer(prompt):
    input_data = tokenizer(prompt, max_length=700, return_tensors='pt').input_ids
    outputs = model.generate(inputs=input_data, num_beams=10, top_k=10, max_length=700)
    result = tokenizer.decode(token_ids=outputs[0], skip_special_tokens=True)
    return result

prompt1 = 'Task Description: In this task, your goal is to generate a SQL query using the input sentence and table schema given. The generated SQL query should be able to verify the given sentence without requiring any additional knowledge.\n\nExample 1:\nTable Schema: \'Author\', \'Publication date\', \'Genre\', \'Publisher\', \'Pages\', \'ISBN\'\nInput: The book "To Kill a Mockingbird" was written by Harper Lee, published by J. B. Lippincott Company in 1960, and has 281 pages.\nOutput: SELECT * FROM table WHERE title = "To Kill a Mockingbird" AND Author = "Harper Lee" AND Publisher = "J. B. Lippincott Company" AND Publication date = "1960" AND Pages = 281\n\nExample 2:\nTable Schema: \'Player\', \'Team\', \'Position\', \'Number\', \'Age\', \'Height\', \'Weight\', \'Nationality\'\nInput: Lionel Messi, who is from Argentina, plays for Paris Saint-Germain and wears the number 30 jersey.\nOutput: SELECT * FROM table WHERE Player = "Lionel Messi" AND Nationality = "Argentina" AND Team = "Paris Saint-Germain" AND Number = 30\n\nExample 3:\nTable Schema: \'Movie Title\', \'Director\', \'Lead Actor\', \'Lead Actress\', \'Release Year\', \'Budget\', \'Box Office Collection\'\nInput: The movie "Forrest Gump" was directed by Robert Zemeckis, starred Tom Hanks and Robin Wright, was released in 1994, and had a budget of $55 million.\nOutput: SELECT * FROM table WHERE "Movie Title" = "Forrest Gump" AND Director = "Robert Zemeckis" AND "Lead Actor" = "Tom Hanks" AND "Lead Actress" = "Robin Wright" AND "Release Year" = 1994 AND Budget = "$55 million"\n\nNow complete the following:\nTable Schema: \'title\', \'Birth Name\', \'Born\', \'Residence\', \'Occupation\', \'Years active\', \'Spouse(s)\', \'Partner(s)\', \'Children\'\nInput:  Chris Pratt has been in his current profession since the turn of the century.\nOutput: \n'

print(infer(prompt1))