init

app.py

@@ -0,0 +1,47 @@
+import gradio as gr
+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")
+
+
+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 infer(input_text, model_choice):
+    tokenizer, model = load_model(model_choice)
+    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))
+
+
+iface = gr.Interface(
+    title="Text to SQL with NSQL",
+    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))
+    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.
+
+    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.
+    """,
+    fn=infer, 
+    inputs=["text", 
+            gr.Dropdown(["nsql-6B", "nsql-2B", "nsql-350M"], value="nsql-6B")], 
+    outputs="text",
+    examples=[[example1, "nsql-6B"],
+              [example2, "nsql-2B"],
+              [example3, "nsql-350M"]])
+iface.launch()

example_strings.py

@@ -0,0 +1,68 @@
+example1 = """CREATE TABLE stadium (
+    stadium_id number,
+    location text,
+    name text,
+    capacity number,
+    highest number,
+    lowest number,
+    average number
+)
+
+CREATE TABLE singer (
+    singer_id number,
+    name text,
+    country text,
+    song_name text,
+    song_release_year text,
+    age number,
+    is_male others
+)
+
+CREATE TABLE concert (
+    concert_id number,
+    concert_name text,
+    theme text,
+    stadium_id text,
+    year text
+)
+
+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 ?
+
+SELECT"""
+
+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.
+
+-- how many stadiums in total?
+
+SELECT"""
+
+example3 = """CREATE TABLE work_orders (
+    ID NUMBER,
+    CREATED_AT TEXT,
+    COST FLOAT,
+    INVOICE_AMOUNT FLOAT,
+    IS_DUE BOOLEAN,
+    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"""

requirements.txt

@@ -0,0 +1 @@
+transformers