Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from example_strings import example1, example2, example3 | |
| template_str = """{table_schemas} | |
| \n \n | |
| {task_spec} | |
| \n \n | |
| {prompt} | |
| \n \n | |
| SELECT""" | |
| def load_model(model_name: str): | |
| tokenizer = AutoTokenizer.from_pretrained(f"NumbersStation/{model_name}") | |
| model = AutoModelForCausalLM.from_pretrained(f"NumbersStation/{model_name}") | |
| return tokenizer, model | |
| def build_complete_prompt(table_schemas: str, task_spec: str, prompt: str) -> str: | |
| return template_str.format(table_schemas=table_schemas, task_spec=task_spec, prompt=prompt) | |
| def infer(table_schemas: str, task_spec: str, prompt: str, model_choice: str = "nsql-350M"): | |
| tokenizer, model = load_model(model_choice) | |
| input_text = build_complete_prompt(table_schemas, task_spec, prompt) | |
| input_ids = tokenizer(input_text, return_tensors="pt").input_ids | |
| generated_ids = model.generate(input_ids, max_length=500) | |
| return (tokenizer.decode(generated_ids[0], skip_special_tokens=True)) | |
| description = """The NSQL model family was published by [Numbers Station](https://www.numbersstation.ai/) and is available in three flavors: | |
| - [nsql-6B](https://huggingface.co/NumbersStation/nsql-6B) | |
| - [nsql-2B](https://huggingface.co/NumbersStation/nsql-2B) | |
| - [nsql-350M]((https://huggingface.co/NumbersStation/nsql-350M)) | |
| For now you can only use the 350M version of the model here, as the file size of the other models exceeds the max memory available in spaces. | |
| In general you should first provide the table schemas of the tables you have questions about and then prompt it with a natural language question. | |
| The model will then generate a SQL query that you can run against your database. | |
| """ | |
| iface = gr.Interface( | |
| title="Text to SQL with NSQL", | |
| description=description, | |
| fn=infer, | |
| inputs=[gr.Text(label="Table schemas", placeholder="Insert your table schemas here"), | |
| gr.Text(label="Specify Task", value="Using valid SQLite, answer the following questions for the tables provided above."), | |
| gr.Text(label="Prompt", placeholder="Put your natural language prompt here"), | |
| # gr.Dropdown(["nsql-6B", "nsql-2B", "nsql-350M"], value="nsql-6B") | |
| ], | |
| outputs="text", | |
| examples=[example1, example2, example3]) | |
| iface.launch() | |