update

app.py

@@ -3,14 +3,13 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from example_strings import example1, example2, example3
 
 
-# tokenizer6B = AutoTokenizer.from_pretrained(f"NumbersStation/nsql-6B")
-# model6B = AutoModelForCausalLM.from_pretrained(f"NumbersStation/nsql-6B")
-
-# tokenizer2B = AutoTokenizer.from_pretrained(f"NumbersStation/nsql-2B")
-# model2B = AutoModelForCausalLM.from_pretrained(f"NumbersStation/nsql-2B")
-
-# tokenizer350M = AutoTokenizer.from_pretrained(f"NumbersStation/nsql-2B")
-# model350M = AutoModelForCausalLM.from_pretrained(f"NumbersStation/nsql-2B")
+template_str = """{table_schemas}
+\n \n
+{task_spec}
+\n \n
+{prompt}
+\n \n
+SELECT"""
 
 
 def load_model(model_name: str):
@@ -19,8 +18,15 @@ def load_model(model_name: str):
     return tokenizer, model
 
 
-def infer(input_text, model_choice):
+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))
@@ -31,7 +37,7 @@ description = """The NSQL model family was published by [Numbers Station](https:
 - [nsql-2B](https://huggingface.co/NumbersStation/nsql-2B)
 - [nsql-350M]((https://huggingface.co/NumbersStation/nsql-350M))  
 
-This demo let's you choose which one you want to use and provides the three examples you can also find in their model cards.  
+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.
@@ -41,10 +47,11 @@ iface = gr.Interface(
     title="Text to SQL with NSQL",
     description=description,
     fn=infer, 
-    inputs=["text", 
-            gr.Dropdown(["nsql-6B", "nsql-2B", "nsql-350M"], value="nsql-6B")], 
+    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, "nsql-350M"],
-              [example2, "nsql-2B"],
-              [example3, "nsql-350M"]])
-iface.launch()
+    examples=[example1, example2, example3])
+iface.launch()

example_strings.py

@@ -1,4 +1,4 @@
-example1 = """CREATE TABLE stadium (
+example1 = ["""CREATE TABLE stadium (
     stadium_id number,
     location text,
     name text,
@@ -29,28 +29,20 @@ CREATE TABLE concert (
 CREATE TABLE singer_in_concert (
     concert_id number,
     singer_id text
-)
-
--- Using valid SQLite, answer the following questions for the tables provided above.
-
--- What is the maximum, the average, and the minimum capacity of stadiums ?
+)""",
+"Using valid SQLite, answer the following questions for the tables provided above.",
+"What is the maximum, the average, and the minimum capacity of stadiums ?"]
 
-SELECT"""
-
-example2 = """CREATE TABLE stadium (
+example2 = ["""CREATE TABLE stadium (
     stadium_id number,
     location text,
     name text,
     capacity number,
-)
-
--- Using valid SQLite, answer the following questions for the tables provided above.
+)""",
+"Using valid SQLite, answer the following questions for the tables provided above.",
+"how many stadiums in total?"]
 
--- how many stadiums in total?
-
-SELECT"""
-
-example3 = """CREATE TABLE work_orders (
+example3 = ["""CREATE TABLE work_orders (
     ID NUMBER,
     CREATED_AT TEXT,
     COST FLOAT,
@@ -59,10 +51,6 @@ example3 = """CREATE TABLE work_orders (
     IS_OPEN BOOLEAN,
     IS_OVERDUE BOOLEAN,
     COUNTRY_NAME TEXT,
-)
-
--- Using valid SQLite, answer the following questions for the tables provided above.
-
--- how many work orders are open?
-
-SELECT"""
+)""",
+"Using valid SQLite, answer the following questions for the tables provided above.",
+"how many work orders are open?"]