Tests exec on db before output

app.py

@@ -1,18 +1,18 @@
 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,
-    TextIteratorStreamer
 )
-from threading import Event, Thread
 # Additional Firebase imports
 import firebase_admin
 from firebase_admin import credentials, firestore
@@ -20,7 +20,6 @@ import json
 import base64
 import torch
 
-
 print(f"Running on {platform.system()}")
 
 if platform.system() == "Windows" or platform.system() == "Darwin":
@@ -33,7 +32,25 @@ initial_model = "WizardLM/WizardCoder-15B-V1.0"
 lora_model = "richardr1126/spider-skeleton-wizard-coder-qlora"
 dataset = "richardr1126/spider-skeleton-context-instruct"
 
-# Firebase code
+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')
@@ -84,7 +101,7 @@ def log_rating_to_firestore(input_message, db_info, temperature, response_text,
     }
     doc_ref.set(log_data)
     gr.Info("Thanks for your feedback!")
-# End Firebase code
+############### End Firebase code ###############
 
 def format(text):
     # Split the text by "|", and get the last element in the list which should be the final query
@@ -105,23 +122,63 @@ def format(text):
 
     return final_query_markdown
 
-model_name = os.getenv("HF_MODEL_NAME", None)
-tok = AutoTokenizer.from_pretrained(model_name)
+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()
 
-max_new_tokens = 1024
+        # 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)
 
-print(f"Starting to load the model {model_name}")
+        # Run the provided test query against the database
+        cursor.execute(query)
+        print(cursor.fetchall())
 
-m = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map=0,
-    #load_in_8bit=True,
-)
+        # Close the connection
+        conn.close()
 
-m.config.pad_token_id = m.config.eos_token_id
-m.generation_config.pad_token_id = m.config.eos_token_id
+        # If everything executed without errors, return True
+        return True
 
-print(f"Successfully loaded the model {model_name} into memory")
+    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):
@@ -139,7 +196,6 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
 
     input_ids = tok(messages, return_tensors="pt").input_ids
     input_ids = input_ids.to(m.device)
-    #streamer = TextIteratorStreamer(tok, timeout=1000.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
         input_ids=input_ids,
         max_new_tokens=max_new_tokens,
@@ -154,15 +210,6 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
         do_sample=do_sample,
     )
 
-    #stream_complete = Event()
-
-    # def generate_and_signal_complete():
-    #     m.generate(**generate_kwargs)
-    #     stream_complete.set()
-
-    # t1 = Thread(target=generate_and_signal_complete)
-    # t1.start()
-
     tokens = m.generate(**generate_kwargs)
 
     responses = []
@@ -172,14 +219,21 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
         # Only take what comes after ### Response:
         response_text = response_text.split("### Response:")[1].strip()
 
-        formatted_text = format(response_text) if format_sql else response_text
+        query = format(response_text) if format_sql else response_text
         if (num_return_sequences > 1):
-            formatted_text = formatted_text.replace("\n", " ").replace("\t", " ").strip()
-
-        responses.append(formatted_text)
+            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)
+            if success:
+                responses.append(query)
+        else:
+            responses.append(query)
+            
 
     # Concat responses to be a single string seperated by a newline
-    output = "\n".join(responses)
+    #output = "\n".join(responses)
+    output = responses[0] if responses else ""
 
     if log:
         # Log the request to Firestore
@@ -219,7 +273,8 @@ with gr.Blocks(theme='gradio/soft') as demo:
         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):
-            num_return_sequences = gr.Slider(label="Num Return Sequences", minimum=1, maximum=5, value=1, step=1)
+            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)
         

requirements.txt

@@ -8,4 +8,5 @@ scipy
 transformers
 accelerate
 sqlparse
-firebase_admin
+firebase_admin
+openai