app.py

import os
import gradio as gr
import sqlite3
import sqlparse
import requests
import time
import re
import platform
import openai
import random
import concurrent.futures
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):
    print(text)
    text = text.replace(".print", "")
    pattern = r'```(?:\w+)?\s?(.*?)```'
    matches = re.findall(pattern, text, re.DOTALL)
    return [match.strip() for match in matches]

def extract_from_code_block(text):
    # First try to match text inside ```
    pattern = r'```(?:\w+)?\s?(.*?)```'
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()
    
    # If no match in ```, try to match SQL-like text
    pattern = r'\((SELECT .*?)\)'
    match = re.search(pattern, text, re.DOTALL)
    if match:
        return match.group(1).strip()

    # Return None if no match found
    return ""

def generate_dummy_db(db_info, question):
    pre_prompt = """
    Generate a SQLite database with dummy data for this database from the DB Layout. Your task is to generate just a database, no queries. For each input do the following:
        1. Breakdown the Question into small pieces and explain what the question is asking for.
        2. Write code to create the specified dummy database using the same exact table and column names used from the DB Layout. Insert dummy data relevant to the Question. Output the datbase code in a single code block. Don't write any queries or SELECT statements in the code.
    """
    prompt = pre_prompt + "\n\nDB Layout:" + db_info + "\n\nQuestion: " + question

    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 10 seconds before retrying...')
            time.sleep(10)

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(f"Query: {query}\tResult: {cursor.fetchall()}")

        # Close the connection
        conn.close()

        # If everything executed without errors, return True
        return True

    except Exception as e:
        print(f"Query: {query}\tError encountered: {e}")
        return False

def choose_best_query(queries, question):
    pre_prompt = """
    Given a list of queries. Your task is to choose just a single query which satisfies the question the most with the least amount of filters, groupings, and conditions. For each input do the following:
        1. Breakdown the list of queries into small pieces and explain what each query is doing.
        2. Breakdown the question peice by piece and explain what each part of the question is asking for. If asking to order by, pay close attention to which order the question is asking for.
        3. Output the most relevant query to the question in a single markdown code block. The user will use regex to extract the SQL query, so make sure it is in a code block.
    """
    prompt = pre_prompt + "\n\nQuestion: " + question + "\n\nQueries:" + "\n\n".join(queries)

    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']
            print(response_text)
            
            query = extract_from_code_block(response_text)

            return query
            
        except Exception as e:
            print(f'Error occurred: {str(e)}')
            print('Waiting for 10 seconds before retrying...')
            time.sleep(10)


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):
    if num_return_sequences > num_beams:
        gr.Warning("Num return sequences must be less than or equal to num beams.")
    
    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,
    )

    # Generate dummy database code if num_return_sequences > 1 in a separate thread
    db_code_future = None
    if num_return_sequences > 1:
        with concurrent.futures.ThreadPoolExecutor() as executor:
            db_code_future = executor.submit(generate_dummy_db, db_info, input_message)
    
    # Generate the SQL query
    tokens = m.generate(**generate_kwargs)

    # Wait for the dummy database code to finish generating
    if db_code_future:
        db_code = db_code_future.result()

    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
            success = test_query_on_dummy_db(db_code, query)
            
            if success:
                responses.append(query)
        else:
            responses.append(query)
            
    # Choose the best query if num_return_sequences > 1
    if num_return_sequences > 1:
        query = choose_best_query(responses, input_message)
        # Format again
        query = format(query) if format_sql else query
        responses = [query]

    output = responses[0]

    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, outputs=output_box)
    #  if platform.system() == "Windows" or platform.system() == "Darwin" else True

    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)