johko/NSQL-Text-To-SQL

1import gradio as gr2from transformers import AutoTokenizer, AutoModelForCausalLM3from example_strings import example1, example2, example34 5 6template_str = """{table_schemas}7\n \n8{task_spec}9\n \n10{prompt}11\n \n12SELECT"""13 14 15def load_model(model_name: str):16    tokenizer = AutoTokenizer.from_pretrained(f"NumbersStation/{model_name}")17    model = AutoModelForCausalLM.from_pretrained(f"NumbersStation/{model_name}")18    return tokenizer, model19 20 21def build_complete_prompt(table_schemas: str, task_spec: str, prompt: str) -> str:22    return template_str.format(table_schemas=table_schemas, task_spec=task_spec, prompt=prompt)23 24 25def infer(table_schemas: str, task_spec: str, prompt: str, model_choice: str = "nsql-350M"):26    tokenizer, model = load_model(model_choice)27    28    input_text = build_complete_prompt(table_schemas, task_spec, prompt)29    30    input_ids = tokenizer(input_text, return_tensors="pt").input_ids31    generated_ids = model.generate(input_ids, max_length=500)32    return (tokenizer.decode(generated_ids[0], skip_special_tokens=True))33 34 35description = """The NSQL model family was published by [Numbers Station](https://www.numbersstation.ai/) and is available in three flavors:  36- [nsql-6B](https://huggingface.co/NumbersStation/nsql-6B)37- [nsql-2B](https://huggingface.co/NumbersStation/nsql-2B)38- [nsql-350M]((https://huggingface.co/NumbersStation/nsql-350M))  39 40This demo let's you choose from all of them and provides the three examples you can also find in their model cards.  41 42In general you should first provide the table schemas of the tables you have questions about and then prompt it with a natural language question.  43The model will then generate a SQL query that you can run against your database.44"""45 46iface = gr.Interface(47    title="Text to SQL with NSQL",48    description=description,49    fn=infer, 50    inputs=[gr.Text(label="Table schemas", placeholder="Insert your table schemas here"),51            gr.Text(label="Specify Task", value="Using valid SQLite, answer the following questions for the tables provided above."), 52            gr.Text(label="Prompt", placeholder="Put your natural language prompt here"), 53            gr.Dropdown(["nsql-6B", "nsql-2B", "nsql-350M"], value="nsql-6B")54            ], 55    outputs="text",56    examples=[example1, example2, example3])57iface.launch()58