Update

app.py

@@ -7,6 +7,7 @@ import time
 import re
 import platform
 import openai
+import random
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -127,9 +128,9 @@ def extract_db_code(text):
     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
+def generate_dummy_db(db_info, question):
+    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.\n\n"
+    prompt = pre_prompt + db_info + "\n\nQuestion: " + question
 
     while True:
         try:
@@ -212,6 +213,10 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
 
     tokens = m.generate(**generate_kwargs)
 
+    db_code = None
+    if (num_return_sequences > 1):
+        db_code = generate_dummy_db(db_info, input_message)
+
     responses = []
     for response in tokens:
         response_text = tok.decode(response, skip_special_tokens=True)
@@ -223,7 +228,6 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
         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
@@ -232,8 +236,8 @@ def generate(input_message: str, db_info="", temperature=0.2, top_p=0.9, top_k=0
         else:
             responses.append(query)
             
-    # Choose the first response
-    output = responses[0] if responses else ""
+    # Choose a random response from responses
+    output = responses[random.randint(0, len(responses)-1)] if len(responses) > 0 else "###"
 
     if log:
         # Log the request to Firestore