Fix

.gitattributes

@@ -1,34 +1,34 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,4 +1,4 @@
-venv/
-.venv/
-env/
+venv/
+.venv/
+env/
 .env

README.md

@@ -1,52 +1,52 @@
----
-title: SQL Skeleton WizardCoder Demo
-emoji: 🕷️☠️🧙‍♂️
-colorFrom: gray
-colorTo: purple
-sdk: gradio
-sdk_version: 3.37.0
-app_file: app.py
-pinned: true
-license: bigcode-openrail-m
-tags:
-  - sql
-  - spider
-  - text-to-sql
-  - sql demo
----
-
-### Spider Skeleton WizardCoder Demo
-
-A demo of [Spider Skeleton Wizard Coder](https://huggingface.co/richardr1126/spider-skeleton-wizard-coder-merged/).
-
-## Citations
-
-```
-@misc{luo2023wizardcoder,
-      title={WizardCoder: Empowering Code Large Language Models with Evol-Instruct}, 
-      author={Ziyang Luo and Can Xu and Pu Zhao and Qingfeng Sun and Xiubo Geng and Wenxiang Hu and Chongyang Tao and Jing Ma and Qingwei Lin and Daxin Jiang},
-      year={2023},
-}
-```
-```
-@article{yu2018spider,
-  title={Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task},
-  author={Yu, Tao and Zhang, Rui and Yang, Kai and Yasunaga, Michihiro and Wang, Dongxu and Li, Zifan and Ma, James and Li, Irene and Yao, Qingning and Roman, Shanelle and others},
-  journal={arXiv preprint arXiv:1809.08887},
-  year={2018}
-}
-```
-```
-@article{dettmers2023qlora,
-  title={QLoRA: Efficient Finetuning of Quantized LLMs},
-  author={Dettmers, Tim and Pagnoni, Artidoro and Holtzman, Ari and Zettlemoyer, Luke},
-  journal={arXiv preprint arXiv:2305.14314},
-  year={2023}
-}
-```
-
-## Disclaimer
-
-The resources, including code, data, and model weights, associated with this project are restricted for academic research purposes only and cannot be used for commercial purposes. The content produced by any version of WizardCoder is influenced by uncontrollable variables such as randomness, and therefore, the accuracy of the output cannot be guaranteed by this project. This project does not accept any legal liability for the content of the model output, nor does it assume responsibility for any losses incurred due to the use of associated resources and output results.
-
+---
+title: SQL Skeleton WizardCoder Demo
+emoji: 🕷️☠️🧙‍♂️
+colorFrom: gray
+colorTo: purple
+sdk: gradio
+sdk_version: 3.37.0
+app_file: app.py
+pinned: true
+license: bigcode-openrail-m
+tags:
+  - sql
+  - spider
+  - text-to-sql
+  - sql demo
+---
+
+### Spider Skeleton WizardCoder Demo
+
+A demo of [Spider Skeleton Wizard Coder](https://huggingface.co/richardr1126/spider-skeleton-wizard-coder-merged/).
+
+## Citations
+
+```
+@misc{luo2023wizardcoder,
+      title={WizardCoder: Empowering Code Large Language Models with Evol-Instruct}, 
+      author={Ziyang Luo and Can Xu and Pu Zhao and Qingfeng Sun and Xiubo Geng and Wenxiang Hu and Chongyang Tao and Jing Ma and Qingwei Lin and Daxin Jiang},
+      year={2023},
+}
+```
+```
+@article{yu2018spider,
+  title={Spider: A large-scale human-labeled dataset for complex and cross-domain semantic parsing and text-to-sql task},
+  author={Yu, Tao and Zhang, Rui and Yang, Kai and Yasunaga, Michihiro and Wang, Dongxu and Li, Zifan and Ma, James and Li, Irene and Yao, Qingning and Roman, Shanelle and others},
+  journal={arXiv preprint arXiv:1809.08887},
+  year={2018}
+}
+```
+```
+@article{dettmers2023qlora,
+  title={QLoRA: Efficient Finetuning of Quantized LLMs},
+  author={Dettmers, Tim and Pagnoni, Artidoro and Holtzman, Ari and Zettlemoyer, Luke},
+  journal={arXiv preprint arXiv:2305.14314},
+  year={2023}
+}
+```
+
+## Disclaimer
+
+The resources, including code, data, and model weights, associated with this project are restricted for academic research purposes only and cannot be used for commercial purposes. The content produced by any version of WizardCoder is influenced by uncontrollable variables such as randomness, and therefore, the accuracy of the output cannot be guaranteed by this project. This project does not accept any legal liability for the content of the model output, nor does it assume responsibility for any losses incurred due to the use of associated resources and output results.
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app-ngrok.py

@@ -1,222 +1,222 @@
-import os
-import gradio as gr
-import sqlparse
-import requests
-from time import sleep
-import re
-import platform
-# Additional Firebase imports
-import firebase_admin
-from firebase_admin import credentials, firestore
-import json
-import base64
-
-print(f"Running on {platform.system()}")
-
-if platform.system() == "Windows" or platform.system() == "Darwin":
-    from dotenv import load_dotenv
-    load_dotenv()
-
-quantized_model = "richardr1126/spider-skeleton-wizard-coder-ggml"
-merged_model = "richardr1126/spider-skeleton-wizard-coder-merged"
-initial_model = "WizardLM/WizardCoder-15B-V1.0"
-lora_model = "richardr1126/spider-skeleton-wizard-coder-qlora"
-dataset = "richardr1126/spider-skeleton-context-instruct"
-
-# Firebase code
-# Initialize Firebase
-base64_string = os.getenv('FIREBASE')
-base64_bytes = base64_string.encode('utf-8')
-json_bytes = base64.b64decode(base64_bytes)
-json_data = json_bytes.decode('utf-8')
-
-firebase_auth = json.loads(json_data)
-
-# Load credentials and initialize Firestore
-cred = credentials.Certificate(firebase_auth)
-firebase_admin.initialize_app(cred)
-db = firestore.client()
-
-def log_message_to_firestore(input_message, db_info, temperature, response_text):
-    doc_ref = db.collection('logs').document()
-    log_data = {
-        'timestamp': firestore.SERVER_TIMESTAMP,
-        'temperature': temperature,
-        'db_info': db_info,
-        'input': input_message,
-        'output': response_text,
-    }
-    doc_ref.set(log_data)
-
-rated_outputs = set()  # set to store already rated outputs
-
-def log_rating_to_firestore(input_message, db_info, temperature, response_text, rating):
-    global rated_outputs
-    output_id = f"{input_message} {db_info} {response_text} {temperature}"
-
-    if output_id in rated_outputs:
-        gr.Warning("You've already rated this output!")
-        return
-    if not input_message or not response_text or not rating:
-        gr.Info("You haven't asked a question yet!")
-        return
-    
-    rated_outputs.add(output_id)
-
-    doc_ref = db.collection('ratings').document()
-    log_data = {
-        'timestamp': firestore.SERVER_TIMESTAMP,
-        'temperature': temperature,
-        'db_info': db_info,
-        'input': input_message,
-        'output': response_text,
-        'rating': rating,
-    }
-    doc_ref.set(log_data)
-    gr.Info("Thanks for your feedback!")
-# End Firebase code
-
-def format(text):
-    # Split the text by "|", and get the last element in the list which should be the final query
-    try:
-        final_query = text.split("|")[1].strip()
-    except Exception:
-        final_query = text
-
-    try:
-        # Attempt to format SQL query using sqlparse
-        formatted_query = sqlparse.format(final_query, reindent=True, keyword_case='upper')
-    except Exception:
-        # If formatting fails, use the original, unformatted query
-        formatted_query = final_query
-
-    # Convert SQL to markdown (not required, but just to show how to use the markdown module)
-    final_query_markdown = f"{formatted_query}"
-
-    return final_query_markdown
-
-def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, stop_sequence="###", log=False):
-    # Format the user's input message
-    messages = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n\nConvert text to sql: {input_message} {db_info}\n\n### Response:\n\n"
-
-    url = os.getenv("KOBOLDCPP_API_URL")
-    stop_sequence = stop_sequence.split(",")
-    stop = ["###"] + stop_sequence
-    payload = {
-        "prompt": messages,
-        "temperature": temperature,
-        "top_p": top_p,
-        "top_k": top_k,
-        "top_a": 0,
-        "n": 1,
-        "max_context_length": 2048,
-        "max_length": 512,
-        "rep_pen": repetition_penalty,
-        "sampler_order": [6,0,1,3,4,2,5],
-        "stop_sequence": stop,
-    }
-    headers = {
-        "Content-Type": "application/json",
-        "ngrok-skip-browser-warning": "1"  # added this line
-    }
-
-    for _ in range(3): # Try 3 times
-        try:
-            response = requests.post(url, json=payload, headers=headers)
-            response_text = response.json()["results"][0]["text"]
-            response_text = response_text.replace("\n", "").replace("\t", " ")
-            if response_text and response_text[-1] == ".":
-                response_text = response_text[:-1]
-
-            output = format(response_text) if format_sql else response_text
-
-            if log:
-                # Log the request to Firestore
-                log_message_to_firestore(input_message, db_info, temperature, output if format_sql else response_text)
-
-            return output
-
-            
-        except Exception as e:
-            print(f'Error occurred: {str(e)}')
-            print('Waiting for 10 seconds before retrying...')
-            gr.Warning("Error occurred, retrying, the sever may be down...")
-            sleep(10)
-
-# Gradio UI Code
-with gr.Blocks(theme='gradio/soft') as demo:
-    # Elements stack vertically by default just define elements in order you want them to stack
-    header = gr.HTML("""
-        <h1 style="text-align: center">SQL Skeleton WizardCoder Demo</h1>
-        <h3 style="text-align: center">🕷️☠️🧙‍♂️ Generate SQL queries from Natural Language 🕷️☠️🧙‍♂️</h3>
-        <div style="max-width: 450px; margin: auto; text-align: center">
-            <p style="font-size: 12px; text-align: center">⚠️ Should take 30-60s to generate. Please rate the response, it helps a lot. If you get a blank output, the model server is currently down, please try again another time.</p>
-        </div>
-    """)
-
-    output_box = gr.Code(label="Generated SQL", lines=2, interactive=False)
-
-    with gr.Row():
-        rate_up = gr.Button("👍", variant="secondary")
-        rate_down = gr.Button("👎", variant="secondary")
-
-    input_text = gr.Textbox(lines=3, placeholder='Write your question here...', label='NL Input')
-    db_info = gr.Textbox(lines=4, placeholder='Make sure to place your tables information inside || for better results. Example: | table_01 : column_01 , column_02 | table_02 : column_01 , column_02 | ...', label='Database Info')
-    format_sql = gr.Checkbox(label="Format SQL + Remove Skeleton", value=True, interactive=True)
-    
-    with gr.Row():
-        run_button = gr.Button("Generate SQL", variant="primary")
-        clear_button = gr.ClearButton(variant="secondary")
-
-    with gr.Accordion("Options", open=False):
-        temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.2, step=0.1)
-        top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.0, maximum=1.0, value=0.9, step=0.01)
-        top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, value=0, step=1)
-        repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.08, step=0.01)
-        stop_sequence = gr.Textbox(lines=1, value="Explanation,Note", label='Extra Stop Sequence')
-    
-    info = gr.HTML(f"""
-        <p>🌐 Leveraging the <a href='https://huggingface.co/{quantized_model}'><strong>4-bit GGML version</strong></a> of <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a> model.</p>
-        <p>🔗 How it's made: <a href='https://huggingface.co/{initial_model}'><strong>{initial_model}</strong></a> was finetuned to create <a href='https://huggingface.co/{lora_model}'><strong>{lora_model}</strong></a>, then merged together to create <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a>.</p>
-        <p>📉 Fine-tuning was performed using QLoRA techniques on the <a href='https://huggingface.co/datasets/{dataset}'><strong>{dataset}</strong></a> dataset. You can view training metrics on the <a href='https://huggingface.co/{lora_model}'><strong>QLoRa adapter HF Repo</strong></a>.</p>
-        <p>📊 All inputs/outputs are logged to Firebase to see how the model is doing. You can also leave a rating for each generated SQL the model produces, which gets sent to the database as well.</a></p>
-    """)
-
-    examples = gr.Examples([
-        ["What is the average, minimum, and maximum age of all singers from France?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-        ["How many students have dogs?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid | pets.pettype = 'Dog' |"],
-    ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence], fn=generate, cache_examples=False if platform.system() == "Windows" or platform.system() == "Darwin" else True, outputs=output_box)
-
-    with gr.Accordion("More Examples", open=False):
-        examples = gr.Examples([
-            ["What is the average weight of pets of all students?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["How many male singers performed in concerts in the year 2023?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["For students who have pets, how many pets does each student have? List their ids instead of names.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["Show location and name for all stadiums with a capacity between 5000 and 10000.", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["What are the number of concerts that occurred in the stadium with the largest capacity ?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["Which student has the oldest pet?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["List the names of all singers who performed in a concert with the theme 'Rock'", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["List all students who don't have pets.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-        ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence], fn=generate, cache_examples=False, outputs=output_box)
-
-
-    readme_content = requests.get(f"https://huggingface.co/{merged_model}/raw/main/README.md").text
-    readme_content = re.sub('---.*?---', '', readme_content, flags=re.DOTALL) #Remove YAML front matter
-
-    with gr.Accordion("📖 Model Readme", open=True):
-        readme = gr.Markdown(
-            readme_content,
-        )
-    
-    with gr.Accordion("Disabled Options:", open=False):
-        log = gr.Checkbox(label="Log to Firebase", value=True, interactive=False)
-    
-    # When the button is clicked, call the generate function, inputs are taken from the UI elements, outputs are sent to outputs elements
-    run_button.click(fn=generate, inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence, log], outputs=output_box, api_name="txt2sql")
-    clear_button.add([input_text, db_info, output_box])
-
-    # Firebase code - for rating the generated SQL (remove if you don't want to use Firebase)
-    rate_up.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_up])
-    rate_down.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_down])
-
+import os
+import gradio as gr
+import sqlparse
+import requests
+from time import sleep
+import re
+import platform
+# Additional Firebase imports
+import firebase_admin
+from firebase_admin import credentials, firestore
+import json
+import base64
+
+print(f"Running on {platform.system()}")
+
+if platform.system() == "Windows" or platform.system() == "Darwin":
+    from dotenv import load_dotenv
+    load_dotenv()
+
+quantized_model = "richardr1126/spider-skeleton-wizard-coder-ggml"
+merged_model = "richardr1126/spider-skeleton-wizard-coder-merged"
+initial_model = "WizardLM/WizardCoder-15B-V1.0"
+lora_model = "richardr1126/spider-skeleton-wizard-coder-qlora"
+dataset = "richardr1126/spider-skeleton-context-instruct"
+
+# Firebase code
+# Initialize Firebase
+base64_string = os.getenv('FIREBASE')
+base64_bytes = base64_string.encode('utf-8')
+json_bytes = base64.b64decode(base64_bytes)
+json_data = json_bytes.decode('utf-8')
+
+firebase_auth = json.loads(json_data)
+
+# Load credentials and initialize Firestore
+cred = credentials.Certificate(firebase_auth)
+firebase_admin.initialize_app(cred)
+db = firestore.client()
+
+def log_message_to_firestore(input_message, db_info, temperature, response_text):
+    doc_ref = db.collection('logs').document()
+    log_data = {
+        'timestamp': firestore.SERVER_TIMESTAMP,
+        'temperature': temperature,
+        'db_info': db_info,
+        'input': input_message,
+        'output': response_text,
+    }
+    doc_ref.set(log_data)
+
+rated_outputs = set()  # set to store already rated outputs
+
+def log_rating_to_firestore(input_message, db_info, temperature, response_text, rating):
+    global rated_outputs
+    output_id = f"{input_message} {db_info} {response_text} {temperature}"
+
+    if output_id in rated_outputs:
+        gr.Warning("You've already rated this output!")
+        return
+    if not input_message or not response_text or not rating:
+        gr.Info("You haven't asked a question yet!")
+        return
+    
+    rated_outputs.add(output_id)
+
+    doc_ref = db.collection('ratings').document()
+    log_data = {
+        'timestamp': firestore.SERVER_TIMESTAMP,
+        'temperature': temperature,
+        'db_info': db_info,
+        'input': input_message,
+        'output': response_text,
+        'rating': rating,
+    }
+    doc_ref.set(log_data)
+    gr.Info("Thanks for your feedback!")
+# End Firebase code
+
+def format(text):
+    # Split the text by "|", and get the last element in the list which should be the final query
+    try:
+        final_query = text.split("|")[1].strip()
+    except Exception:
+        final_query = text
+
+    try:
+        # Attempt to format SQL query using sqlparse
+        formatted_query = sqlparse.format(final_query, reindent=True, keyword_case='upper')
+    except Exception:
+        # If formatting fails, use the original, unformatted query
+        formatted_query = final_query
+
+    # Convert SQL to markdown (not required, but just to show how to use the markdown module)
+    final_query_markdown = f"{formatted_query}"
+
+    return final_query_markdown
+
+def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, stop_sequence="###", log=False):
+    # Format the user's input message
+    messages = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n\nConvert text to sql: {input_message} {db_info}\n\n### Response:\n\n"
+
+    url = os.getenv("KOBOLDCPP_API_URL")
+    stop_sequence = stop_sequence.split(",")
+    stop = ["###"] + stop_sequence
+    payload = {
+        "prompt": messages,
+        "temperature": temperature,
+        "top_p": top_p,
+        "top_k": top_k,
+        "top_a": 0,
+        "n": 1,
+        "max_context_length": 2048,
+        "max_length": 512,
+        "rep_pen": repetition_penalty,
+        "sampler_order": [6,0,1,3,4,2,5],
+        "stop_sequence": stop,
+    }
+    headers = {
+        "Content-Type": "application/json",
+        "ngrok-skip-browser-warning": "1"  # added this line
+    }
+
+    for _ in range(3): # Try 3 times
+        try:
+            response = requests.post(url, json=payload, headers=headers)
+            response_text = response.json()["results"][0]["text"]
+            response_text = response_text.replace("\n", "").replace("\t", " ")
+            if response_text and response_text[-1] == ".":
+                response_text = response_text[:-1]
+
+            output = format(response_text) if format_sql else response_text
+
+            if log:
+                # Log the request to Firestore
+                log_message_to_firestore(input_message, db_info, temperature, output if format_sql else response_text)
+
+            return output
+
+            
+        except Exception as e:
+            print(f'Error occurred: {str(e)}')
+            print('Waiting for 10 seconds before retrying...')
+            gr.Warning("Error occurred, retrying, the sever may be down...")
+            sleep(10)
+
+# Gradio UI Code
+with gr.Blocks(theme='gradio/soft') as demo:
+    # Elements stack vertically by default just define elements in order you want them to stack
+    header = gr.HTML("""
+        <h1 style="text-align: center">SQL Skeleton WizardCoder Demo</h1>
+        <h3 style="text-align: center">🕷️☠️🧙‍♂️ Generate SQL queries from Natural Language 🕷️☠️🧙‍♂️</h3>
+        <div style="max-width: 450px; margin: auto; text-align: center">
+            <p style="font-size: 12px; text-align: center">⚠️ Should take 30-60s to generate. Please rate the response, it helps a lot. If you get a blank output, the model server is currently down, please try again another time.</p>
+        </div>
+    """)
+
+    output_box = gr.Code(label="Generated SQL", lines=2, interactive=False)
+
+    with gr.Row():
+        rate_up = gr.Button("👍", variant="secondary")
+        rate_down = gr.Button("👎", variant="secondary")
+
+    input_text = gr.Textbox(lines=3, placeholder='Write your question here...', label='NL Input')
+    db_info = gr.Textbox(lines=4, placeholder='Make sure to place your tables information inside || for better results. Example: | table_01 : column_01 , column_02 | table_02 : column_01 , column_02 | ...', label='Database Info')
+    format_sql = gr.Checkbox(label="Format SQL + Remove Skeleton", value=True, interactive=True)
+    
+    with gr.Row():
+        run_button = gr.Button("Generate SQL", variant="primary")
+        clear_button = gr.ClearButton(variant="secondary")
+
+    with gr.Accordion("Options", open=False):
+        temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.2, step=0.1)
+        top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.0, maximum=1.0, value=0.9, step=0.01)
+        top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, value=0, step=1)
+        repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.08, step=0.01)
+        stop_sequence = gr.Textbox(lines=1, value="Explanation,Note", label='Extra Stop Sequence')
+    
+    info = gr.HTML(f"""
+        <p>🌐 Leveraging the <a href='https://huggingface.co/{quantized_model}'><strong>4-bit GGML version</strong></a> of <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a> model.</p>
+        <p>🔗 How it's made: <a href='https://huggingface.co/{initial_model}'><strong>{initial_model}</strong></a> was finetuned to create <a href='https://huggingface.co/{lora_model}'><strong>{lora_model}</strong></a>, then merged together to create <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a>.</p>
+        <p>📉 Fine-tuning was performed using QLoRA techniques on the <a href='https://huggingface.co/datasets/{dataset}'><strong>{dataset}</strong></a> dataset. You can view training metrics on the <a href='https://huggingface.co/{lora_model}'><strong>QLoRa adapter HF Repo</strong></a>.</p>
+        <p>📊 All inputs/outputs are logged to Firebase to see how the model is doing. You can also leave a rating for each generated SQL the model produces, which gets sent to the database as well.</a></p>
+    """)
+
+    examples = gr.Examples([
+        ["What is the average, minimum, and maximum age of all singers from France?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+        ["How many students have dogs?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid | pets.pettype = 'Dog' |"],
+    ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence], fn=generate, cache_examples=False if platform.system() == "Windows" or platform.system() == "Darwin" else True, outputs=output_box)
+
+    with gr.Accordion("More Examples", open=False):
+        examples = gr.Examples([
+            ["What is the average weight of pets of all students?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["How many male singers performed in concerts in the year 2023?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["For students who have pets, how many pets does each student have? List their ids instead of names.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["Show location and name for all stadiums with a capacity between 5000 and 10000.", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["What are the number of concerts that occurred in the stadium with the largest capacity ?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["Which student has the oldest pet?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["List the names of all singers who performed in a concert with the theme 'Rock'", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["List all students who don't have pets.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+        ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence], fn=generate, cache_examples=False, outputs=output_box)
+
+
+    readme_content = requests.get(f"https://huggingface.co/{merged_model}/raw/main/README.md").text
+    readme_content = re.sub('---.*?---', '', readme_content, flags=re.DOTALL) #Remove YAML front matter
+
+    with gr.Accordion("📖 Model Readme", open=True):
+        readme = gr.Markdown(
+            readme_content,
+        )
+    
+    with gr.Accordion("Disabled Options:", open=False):
+        log = gr.Checkbox(label="Log to Firebase", value=True, interactive=False)
+    
+    # When the button is clicked, call the generate function, inputs are taken from the UI elements, outputs are sent to outputs elements
+    run_button.click(fn=generate, inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, stop_sequence, log], outputs=output_box, api_name="txt2sql")
+    clear_button.add([input_text, db_info, output_box])
+
+    # Firebase code - for rating the generated SQL (remove if you don't want to use Firebase)
+    rate_up.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_up])
+    rate_down.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_down])
+
 demo.queue(concurrency_count=1, max_size=20).launch(debug=True)

app.py

@@ -1,325 +1,325 @@
-import os
-import gradio as gr
-import sqlite3
-import sqlparse
-import requests
-from time import sleep
-import re
-import platform
-import openai
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    StoppingCriteria,
-    StoppingCriteriaList,
-)
-# Additional Firebase imports
-import firebase_admin
-from firebase_admin import credentials, firestore
-import json
-import base64
-import torch
-
-print(f"Running on {platform.system()}")
-
-if platform.system() == "Windows" or platform.system() == "Darwin":
-    from dotenv import load_dotenv
-    load_dotenv()
-
-quantized_model = "richardr1126/spider-skeleton-wizard-coder-8bit"
-merged_model = "richardr1126/spider-skeleton-wizard-coder-merged"
-initial_model = "WizardLM/WizardCoder-15B-V1.0"
-lora_model = "richardr1126/spider-skeleton-wizard-coder-qlora"
-dataset = "richardr1126/spider-skeleton-context-instruct"
-
-model_name = os.getenv("HF_MODEL_NAME", None)
-tok = AutoTokenizer.from_pretrained(model_name)
-
-max_new_tokens = 1024
-
-print(f"Starting to load the model {model_name}")
-
-m = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map=0,
-    #load_in_8bit=True,
-)
-
-m.config.pad_token_id = m.config.eos_token_id
-m.generation_config.pad_token_id = m.config.eos_token_id
-
-print(f"Successfully loaded the model {model_name} into memory")
-
-################# Firebase code #################
-# Initialize Firebase
-base64_string = os.getenv('FIREBASE')
-base64_bytes = base64_string.encode('utf-8')
-json_bytes = base64.b64decode(base64_bytes)
-json_data = json_bytes.decode('utf-8')
-
-firebase_auth = json.loads(json_data)
-
-# Load credentials and initialize Firestore
-cred = credentials.Certificate(firebase_auth)
-firebase_admin.initialize_app(cred)
-db = firestore.client()
-
-def log_message_to_firestore(input_message, db_info, temperature, response_text):
-    doc_ref = db.collection('logs').document()
-    log_data = {
-        'timestamp': firestore.SERVER_TIMESTAMP,
-        'temperature': temperature,
-        'db_info': db_info,
-        'input': input_message,
-        'output': response_text,
-    }
-    doc_ref.set(log_data)
-
-rated_outputs = set()  # set to store already rated outputs
-
-def log_rating_to_firestore(input_message, db_info, temperature, response_text, rating):
-    global rated_outputs
-    output_id = f"{input_message} {db_info} {response_text} {temperature}"
-
-    if output_id in rated_outputs:
-        gr.Warning("You've already rated this output!")
-        return
-    if not input_message or not response_text or not rating:
-        gr.Info("You haven't asked a question yet!")
-        return
-    
-    rated_outputs.add(output_id)
-
-    doc_ref = db.collection('ratings').document()
-    log_data = {
-        'timestamp': firestore.SERVER_TIMESTAMP,
-        'temperature': temperature,
-        'db_info': db_info,
-        'input': input_message,
-        'output': response_text,
-        'rating': rating,
-    }
-    doc_ref.set(log_data)
-    gr.Info("Thanks for your feedback!")
-############### End Firebase code ###############
-
-def format(text):
-    # Split the text by "|", and get the last element in the list which should be the final query
-    try:
-        final_query = text.split("|")[1].strip()
-    except Exception:
-        final_query = text
-
-    try:
-        # Attempt to format SQL query using sqlparse
-        formatted_query = sqlparse.format(final_query, reindent=True, keyword_case='upper')
-    except Exception:
-        # If formatting fails, use the original, unformatted query
-        formatted_query = final_query
-
-    # Convert SQL to markdown (not required, but just to show how to use the markdown module)
-    final_query_markdown = f"{formatted_query}"
-
-    return final_query_markdown
-
-def extract_db_code(text):
-    pattern = r'```(?:\w+)?\s?(.*?)```'
-    matches = re.findall(pattern, text, re.DOTALL)
-    return [match.strip() for match in matches]
-
-def generate_dummy_db(db_info, question, query):
-    pre_prompt = "Generate a SQLite database with dummy data for this database, output the SQL code in a SQL code block. Make sure you add dummy data relevant to the question and query.\n\n"
-    prompt = pre_prompt + db_info + "\n\nQuestion: " + question + "\nQuery: " + query
-
-    while True:
-        try:
-            response = openai.ChatCompletion.create(
-                model="gpt-3.5-turbo",
-                messages=[
-                    {"role": "user", "content": prompt}
-                ],
-                #temperature=0.7,
-            )
-            response_text = response['choices'][0]['message']['content']
-            
-            db_code = extract_db_code(response_text)
-
-            return db_code
-            
-        except Exception as e:
-            print(f'Error occurred: {str(e)}')
-            print('Waiting for 20 seconds before retrying...')
-            time.sleep(20)
-
-def test_query_on_dummy_db(db_code, query):
-    try:
-        # Connect to an SQLite database in memory
-        conn = sqlite3.connect(':memory:')
-        cursor = conn.cursor()
-
-        # Iterate over each extracted SQL block and split them into individual commands
-        for sql_block in db_code:
-            statements = sqlparse.split(sql_block)
-            
-            # Execute each SQL command
-            for statement in statements:
-                if statement:
-                    cursor.execute(statement)
-
-        # Run the provided test query against the database
-        cursor.execute(query)
-        print(cursor.fetchall())
-
-        # Close the connection
-        conn.close()
-
-        # If everything executed without errors, return True
-        return True
-
-    except Exception as e:
-        print(f"Error encountered: {e}")
-        return False
-
-
-def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, log=False, num_return_sequences=1, num_beams=1, do_sample=False):
-    stop_token_ids = tok.convert_tokens_to_ids(["###"])
-    class StopOnTokens(StoppingCriteria):
-        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-            for stop_id in stop_token_ids:
-                if input_ids[0][-1] == stop_id:
-                    return True
-            return False
-    stop = StopOnTokens()
-
-    # Format the user's input message
-    messages = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n\nConvert text to sql: {input_message} {db_info}\n\n### Response:\n\n"
-
-    input_ids = tok(messages, return_tensors="pt").input_ids
-    input_ids = input_ids.to(m.device)
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        repetition_penalty=repetition_penalty,
-        #streamer=streamer,
-        stopping_criteria=StoppingCriteriaList([stop]),
-        num_return_sequences=num_return_sequences,
-        num_beams=num_beams,
-        do_sample=do_sample,
-    )
-
-    tokens = m.generate(**generate_kwargs)
-
-    responses = []
-    for response in tokens:
-        response_text = tok.decode(response, skip_special_tokens=True)
-
-        # Only take what comes after ### Response:
-        response_text = response_text.split("### Response:")[1].strip()
-
-        query = format(response_text) if format_sql else response_text
-        if (num_return_sequences > 1):
-            query = query.replace("\n", " ").replace("\t", " ").strip()
-            # Test against dummy database
-            db_code = generate_dummy_db(db_info, input_message, query)
-            success = test_query_on_dummy_db(db_code, query)
-            # Format again
-            query = format(query) if format_sql else query
-            if success:
-                responses.append(query)
-        else:
-            responses.append(query)
-            
-    # Choose the first response
-    output = responses[0] if responses else ""
-
-    if log:
-        # Log the request to Firestore
-        log_message_to_firestore(input_message, db_info, temperature, output)
-
-    return output
-
-# Gradio UI Code
-with gr.Blocks(theme='gradio/soft') as demo:
-    # Elements stack vertically by default just define elements in order you want them to stack
-    header = gr.HTML("""
-        <h1 style="text-align: center">SQL Skeleton WizardCoder Demo</h1>
-        <h3 style="text-align: center">🕷️☠️🧙‍♂️ Generate SQL queries from Natural Language 🕷️☠️🧙‍♂️</h3>
-        <div style="max-width: 450px; margin: auto; text-align: center">
-            <p style="font-size: 12px; text-align: center">⚠️ Should take 30-60s to generate. Please rate the response, it helps a lot. If you get a blank output, the model server is currently down, please try again another time.</p>
-        </div>
-    """)
-
-    output_box = gr.Code(label="Generated SQL", lines=2, interactive=False)
-
-    with gr.Row():
-        rate_up = gr.Button("👍", variant="secondary")
-        rate_down = gr.Button("👎", variant="secondary")
-
-    input_text = gr.Textbox(lines=3, placeholder='Write your question here...', label='NL Input')
-    db_info = gr.Textbox(lines=4, placeholder='Make sure to place your tables information inside || for better results. Example: | table_01 : column_01 , column_02 | table_02 : column_01 , column_02 | ...', label='Database Info')
-    format_sql = gr.Checkbox(label="Format SQL + Remove Skeleton", value=True, interactive=True)
-    
-    with gr.Row():
-        run_button = gr.Button("Generate SQL", variant="primary")
-        clear_button = gr.ClearButton(variant="secondary")
-
-    with gr.Accordion("Options", open=False):
-        temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.2, step=0.1)
-        top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.0, maximum=1.0, value=0.9, step=0.01)
-        top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, value=0, step=1)
-        repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.08, step=0.01)
-
-        with gr.Accordion("Generation strategies", open=False):
-            md_description = gr.Markdown("""Increasing num return sequences will increase the number of SQLs generated, but will still yield only the best output of the number of return sequences. SQLs are tested against the db info you provide.""")
-            num_return_sequences = gr.Slider(label="Number of return sequences (to generate and test)", minimum=1, maximum=5, value=1, step=1)
-            num_beams = gr.Slider(label="Num Beams", minimum=1, maximum=5, value=1, step=1)
-            do_sample = gr.Checkbox(label="Do Sample", value=False, interactive=True)
-        
-    info = gr.HTML(f"""
-        <p>🌐 Leveraging the <a href='https://huggingface.co/{quantized_model}'><strong>bitsandbytes 8-bit version</strong></a> of <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a> model.</p>
-        <p>🔗 How it's made: <a href='https://huggingface.co/{initial_model}'><strong>{initial_model}</strong></a> was finetuned to create <a href='https://huggingface.co/{lora_model}'><strong>{lora_model}</strong></a>, then merged together to create <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a>.</p>
-        <p>📉 Fine-tuning was performed using QLoRA techniques on the <a href='https://huggingface.co/datasets/{dataset}'><strong>{dataset}</strong></a> dataset. You can view training metrics on the <a href='https://huggingface.co/{lora_model}'><strong>QLoRa adapter HF Repo</strong></a>.</p>
-        <p>📊 All inputs/outputs are logged to Firebase to see how the model is doing. You can also leave a rating for each generated SQL the model produces, which gets sent to the database as well.</a></p>
-    """)
-
-    examples = gr.Examples([
-        ["What is the average, minimum, and maximum age of all singers from France?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-        ["How many students have dogs?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid | pets.pettype = 'Dog' |"],
-    ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql], fn=generate, cache_examples=False if platform.system() == "Windows" or platform.system() == "Darwin" else True, outputs=output_box)
-
-    with gr.Accordion("More Examples", open=False):
-        examples = gr.Examples([
-            ["What is the average weight of pets of all students?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["How many male singers performed in concerts in the year 2023?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["For students who have pets, how many pets does each student have? List their ids instead of names.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["Show location and name for all stadiums with a capacity between 5000 and 10000.", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["What are the number of concerts that occurred in the stadium with the largest capacity ?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["Which student has the oldest pet?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-            ["List the names of all singers who performed in a concert with the theme 'Rock'", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
-            ["List all students who don't have pets.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
-        ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql], fn=generate, cache_examples=False, outputs=output_box)
-
-
-    readme_content = requests.get(f"https://huggingface.co/{merged_model}/raw/main/README.md").text
-    readme_content = re.sub('---.*?---', '', readme_content, flags=re.DOTALL) #Remove YAML front matter
-
-    with gr.Accordion("📖 Model Readme", open=True):
-        readme = gr.Markdown(
-            readme_content,
-        )
-    
-    with gr.Accordion("Disabled Options:", open=False):
-        log = gr.Checkbox(label="Log to Firebase", value=True, interactive=False)
-    
-    # When the button is clicked, call the generate function, inputs are taken from the UI elements, outputs are sent to outputs elements
-    run_button.click(fn=generate, inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, log, num_return_sequences, num_beams, do_sample], outputs=output_box, api_name="txt2sql")
-    clear_button.add([input_text, db_info, output_box])
-
-    # Firebase code - for rating the generated SQL (remove if you don't want to use Firebase)
-    rate_up.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_up])
-    rate_down.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_down])
-
+import os
+import gradio as gr
+import sqlite3
+import sqlparse
+import requests
+import time
+import re
+import platform
+import openai
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+)
+# Additional Firebase imports
+import firebase_admin
+from firebase_admin import credentials, firestore
+import json
+import base64
+import torch
+
+print(f"Running on {platform.system()}")
+
+if platform.system() == "Windows" or platform.system() == "Darwin":
+    from dotenv import load_dotenv
+    load_dotenv()
+
+quantized_model = "richardr1126/spider-skeleton-wizard-coder-8bit"
+merged_model = "richardr1126/spider-skeleton-wizard-coder-merged"
+initial_model = "WizardLM/WizardCoder-15B-V1.0"
+lora_model = "richardr1126/spider-skeleton-wizard-coder-qlora"
+dataset = "richardr1126/spider-skeleton-context-instruct"
+
+model_name = os.getenv("HF_MODEL_NAME", None)
+tok = AutoTokenizer.from_pretrained(model_name)
+
+max_new_tokens = 1024
+
+print(f"Starting to load the model {model_name}")
+
+m = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map=0,
+    #load_in_8bit=True,
+)
+
+m.config.pad_token_id = m.config.eos_token_id
+m.generation_config.pad_token_id = m.config.eos_token_id
+
+print(f"Successfully loaded the model {model_name} into memory")
+
+################# Firebase code #################
+# Initialize Firebase
+base64_string = os.getenv('FIREBASE')
+base64_bytes = base64_string.encode('utf-8')
+json_bytes = base64.b64decode(base64_bytes)
+json_data = json_bytes.decode('utf-8')
+
+firebase_auth = json.loads(json_data)
+
+# Load credentials and initialize Firestore
+cred = credentials.Certificate(firebase_auth)
+firebase_admin.initialize_app(cred)
+db = firestore.client()
+
+def log_message_to_firestore(input_message, db_info, temperature, response_text):
+    doc_ref = db.collection('logs').document()
+    log_data = {
+        'timestamp': firestore.SERVER_TIMESTAMP,
+        'temperature': temperature,
+        'db_info': db_info,
+        'input': input_message,
+        'output': response_text,
+    }
+    doc_ref.set(log_data)
+
+rated_outputs = set()  # set to store already rated outputs
+
+def log_rating_to_firestore(input_message, db_info, temperature, response_text, rating):
+    global rated_outputs
+    output_id = f"{input_message} {db_info} {response_text} {temperature}"
+
+    if output_id in rated_outputs:
+        gr.Warning("You've already rated this output!")
+        return
+    if not input_message or not response_text or not rating:
+        gr.Info("You haven't asked a question yet!")
+        return
+    
+    rated_outputs.add(output_id)
+
+    doc_ref = db.collection('ratings').document()
+    log_data = {
+        'timestamp': firestore.SERVER_TIMESTAMP,
+        'temperature': temperature,
+        'db_info': db_info,
+        'input': input_message,
+        'output': response_text,
+        'rating': rating,
+    }
+    doc_ref.set(log_data)
+    gr.Info("Thanks for your feedback!")
+############### End Firebase code ###############
+
+def format(text):
+    # Split the text by "|", and get the last element in the list which should be the final query
+    try:
+        final_query = text.split("|")[1].strip()
+    except Exception:
+        final_query = text
+
+    try:
+        # Attempt to format SQL query using sqlparse
+        formatted_query = sqlparse.format(final_query, reindent=True, keyword_case='upper')
+    except Exception:
+        # If formatting fails, use the original, unformatted query
+        formatted_query = final_query
+
+    # Convert SQL to markdown (not required, but just to show how to use the markdown module)
+    final_query_markdown = f"{formatted_query}"
+
+    return final_query_markdown
+
+def extract_db_code(text):
+    pattern = r'```(?:\w+)?\s?(.*?)```'
+    matches = re.findall(pattern, text, re.DOTALL)
+    return [match.strip() for match in matches]
+
+def generate_dummy_db(db_info, question, query):
+    pre_prompt = "Generate a SQLite database with dummy data for this database, output the SQL code in a SQL code block. Make sure you add dummy data relevant to the question and query.\n\n"
+    prompt = pre_prompt + db_info + "\n\nQuestion: " + question + "\nQuery: " + query
+
+    while True:
+        try:
+            response = openai.ChatCompletion.create(
+                model="gpt-3.5-turbo",
+                messages=[
+                    {"role": "user", "content": prompt}
+                ],
+                #temperature=0.7,
+            )
+            response_text = response['choices'][0]['message']['content']
+            
+            db_code = extract_db_code(response_text)
+
+            return db_code
+            
+        except Exception as e:
+            print(f'Error occurred: {str(e)}')
+            print('Waiting for 20 seconds before retrying...')
+            time.sleep(20)
+
+def test_query_on_dummy_db(db_code, query):
+    try:
+        # Connect to an SQLite database in memory
+        conn = sqlite3.connect(':memory:')
+        cursor = conn.cursor()
+
+        # Iterate over each extracted SQL block and split them into individual commands
+        for sql_block in db_code:
+            statements = sqlparse.split(sql_block)
+            
+            # Execute each SQL command
+            for statement in statements:
+                if statement:
+                    cursor.execute(statement)
+
+        # Run the provided test query against the database
+        cursor.execute(query)
+        print(cursor.fetchall())
+
+        # Close the connection
+        conn.close()
+
+        # If everything executed without errors, return True
+        return True
+
+    except Exception as e:
+        print(f"Error encountered: {e}")
+        return False
+
+
+def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0, repetition_penalty=1.08, format_sql=True, log=False, num_return_sequences=1, num_beams=1, do_sample=False):
+    stop_token_ids = tok.convert_tokens_to_ids(["###"])
+    class StopOnTokens(StoppingCriteria):
+        def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+            for stop_id in stop_token_ids:
+                if input_ids[0][-1] == stop_id:
+                    return True
+            return False
+    stop = StopOnTokens()
+
+    # Format the user's input message
+    messages = f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n\nConvert text to sql: {input_message} {db_info}\n\n### Response:\n\n"
+
+    input_ids = tok(messages, return_tensors="pt").input_ids
+    input_ids = input_ids.to(m.device)
+    generate_kwargs = dict(
+        input_ids=input_ids,
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+        #streamer=streamer,
+        stopping_criteria=StoppingCriteriaList([stop]),
+        num_return_sequences=num_return_sequences,
+        num_beams=num_beams,
+        do_sample=do_sample,
+    )
+
+    tokens = m.generate(**generate_kwargs)
+
+    responses = []
+    for response in tokens:
+        response_text = tok.decode(response, skip_special_tokens=True)
+
+        # Only take what comes after ### Response:
+        response_text = response_text.split("### Response:")[1].strip()
+
+        query = format(response_text) if format_sql else response_text
+        if (num_return_sequences > 1):
+            query = query.replace("\n", " ").replace("\t", " ").strip()
+            # Test against dummy database
+            db_code = generate_dummy_db(db_info, input_message, query)
+            success = test_query_on_dummy_db(db_code, query)
+            # Format again
+            query = format(query) if format_sql else query
+            if success:
+                responses.append(query)
+        else:
+            responses.append(query)
+            
+    # Choose the first response
+    output = responses[0] if responses else ""
+
+    if log:
+        # Log the request to Firestore
+        log_message_to_firestore(input_message, db_info, temperature, output)
+
+    return output
+
+# Gradio UI Code
+with gr.Blocks(theme='gradio/soft') as demo:
+    # Elements stack vertically by default just define elements in order you want them to stack
+    header = gr.HTML("""
+        <h1 style="text-align: center">SQL Skeleton WizardCoder Demo</h1>
+        <h3 style="text-align: center">🕷️☠️🧙‍♂️ Generate SQL queries from Natural Language 🕷️☠️🧙‍♂️</h3>
+        <div style="max-width: 450px; margin: auto; text-align: center">
+            <p style="font-size: 12px; text-align: center">⚠️ Should take 30-60s to generate. Please rate the response, it helps a lot. If you get a blank output, the model server is currently down, please try again another time.</p>
+        </div>
+    """)
+
+    output_box = gr.Code(label="Generated SQL", lines=2, interactive=False)
+
+    with gr.Row():
+        rate_up = gr.Button("👍", variant="secondary")
+        rate_down = gr.Button("👎", variant="secondary")
+
+    input_text = gr.Textbox(lines=3, placeholder='Write your question here...', label='NL Input')
+    db_info = gr.Textbox(lines=4, placeholder='Make sure to place your tables information inside || for better results. Example: | table_01 : column_01 , column_02 | table_02 : column_01 , column_02 | ...', label='Database Info')
+    format_sql = gr.Checkbox(label="Format SQL + Remove Skeleton", value=True, interactive=True)
+    
+    with gr.Row():
+        run_button = gr.Button("Generate SQL", variant="primary")
+        clear_button = gr.ClearButton(variant="secondary")
+
+    with gr.Accordion("Options", open=False):
+        temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.2, step=0.1)
+        top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.0, maximum=1.0, value=0.9, step=0.01)
+        top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, value=0, step=1)
+        repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.08, step=0.01)
+
+        with gr.Accordion("Generation strategies", open=False):
+            md_description = gr.Markdown("""Increasing num return sequences will increase the number of SQLs generated, but will still yield only the best output of the number of return sequences. SQLs are tested against the db info you provide.""")
+            num_return_sequences = gr.Slider(label="Number of return sequences (to generate and test)", minimum=1, maximum=5, value=1, step=1)
+            num_beams = gr.Slider(label="Num Beams", minimum=1, maximum=5, value=1, step=1)
+            do_sample = gr.Checkbox(label="Do Sample", value=False, interactive=True)
+        
+    info = gr.HTML(f"""
+        <p>🌐 Leveraging the <a href='https://huggingface.co/{quantized_model}'><strong>bitsandbytes 8-bit version</strong></a> of <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a> model.</p>
+        <p>🔗 How it's made: <a href='https://huggingface.co/{initial_model}'><strong>{initial_model}</strong></a> was finetuned to create <a href='https://huggingface.co/{lora_model}'><strong>{lora_model}</strong></a>, then merged together to create <a href='https://huggingface.co/{merged_model}'><strong>{merged_model}</strong></a>.</p>
+        <p>📉 Fine-tuning was performed using QLoRA techniques on the <a href='https://huggingface.co/datasets/{dataset}'><strong>{dataset}</strong></a> dataset. You can view training metrics on the <a href='https://huggingface.co/{lora_model}'><strong>QLoRa adapter HF Repo</strong></a>.</p>
+        <p>📊 All inputs/outputs are logged to Firebase to see how the model is doing. You can also leave a rating for each generated SQL the model produces, which gets sent to the database as well.</a></p>
+    """)
+
+    examples = gr.Examples([
+        ["What is the average, minimum, and maximum age of all singers from France?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+        ["How many students have dogs?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid | pets.pettype = 'Dog' |"],
+    ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql], fn=generate, cache_examples=False if platform.system() == "Windows" or platform.system() == "Darwin" else True, outputs=output_box)
+
+    with gr.Accordion("More Examples", open=False):
+        examples = gr.Examples([
+            ["What is the average weight of pets of all students?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["How many male singers performed in concerts in the year 2023?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["For students who have pets, how many pets does each student have? List their ids instead of names.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["Show location and name for all stadiums with a capacity between 5000 and 10000.", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["What are the number of concerts that occurred in the stadium with the largest capacity ?", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["Which student has the oldest pet?", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+            ["List the names of all singers who performed in a concert with the theme 'Rock'", "| stadium : stadium_id , location , name , capacity , highest , lowest , average | singer : singer_id , name , country , song_name , song_release_year , age , is_male | concert : concert_id , concert_name , theme , stadium_id , year | singer_in_concert : concert_id , singer_id | concert.stadium_id = stadium.stadium_id | singer_in_concert.singer_id = singer.singer_id | singer_in_concert.concert_id = concert.concert_id |"],
+            ["List all students who don't have pets.", "| student : stuid , lname , fname , age , sex , major , advisor , city_code | has_pet : stuid , petid | pets : petid , pettype , pet_age , weight | has_pet.stuid = student.stuid | has_pet.petid = pets.petid |"],
+        ], inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql], fn=generate, cache_examples=False, outputs=output_box)
+
+
+    readme_content = requests.get(f"https://huggingface.co/{merged_model}/raw/main/README.md").text
+    readme_content = re.sub('---.*?---', '', readme_content, flags=re.DOTALL) #Remove YAML front matter
+
+    with gr.Accordion("📖 Model Readme", open=True):
+        readme = gr.Markdown(
+            readme_content,
+        )
+    
+    with gr.Accordion("Disabled Options:", open=False):
+        log = gr.Checkbox(label="Log to Firebase", value=True, interactive=False)
+    
+    # When the button is clicked, call the generate function, inputs are taken from the UI elements, outputs are sent to outputs elements
+    run_button.click(fn=generate, inputs=[input_text, db_info, temperature, top_p, top_k, repetition_penalty, format_sql, log, num_return_sequences, num_beams, do_sample], outputs=output_box, api_name="txt2sql")
+    clear_button.add([input_text, db_info, output_box])
+
+    # Firebase code - for rating the generated SQL (remove if you don't want to use Firebase)
+    rate_up.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_up])
+    rate_down.click(fn=log_rating_to_firestore, inputs=[input_text, db_info, temperature, output_box, rate_down])
+
 demo.queue(concurrency_count=1, max_size=20).launch(debug=True)

requirements.txt

@@ -1,12 +1,12 @@
-einops
-gradio
-torch
-numpy
-sentencepiece
-bitsandbytes
-scipy
-transformers
-accelerate
-sqlparse
-firebase_admin
-openai
+einops
+gradio
+numpy
+sentencepiece
+bitsandbytes
+scipy
+transformers
+accelerate
+sqlparse
+firebase_admin
+openai
+python-dotenv