feat: add TQA prompt and db_schema with INSERT INTO

prediction.py

@@ -57,6 +57,16 @@ class ModelPrediction:
             "Database Schema\n"
             "{db_schema}\n"
         )
+        self.base_prompt_QA= (
+            "Return the answer of the following question based on the provided database."
+            " Return your answer as the result of a query executed over the database."
+            " Namely, as a list of list where the first list represent the tuples and the second list the values in that tuple.\n"
+            "Return the answer in answer tag as <answer> </answer>"
+            " Question\n"
+            "{question}\n"
+            "Database Schema\n"
+            "{db_schema}\n"
+        )
 
     def _reset_pipeline(self, model_name):
         if self._model_name != model_name:
@@ -74,7 +84,7 @@ class ModelPrediction:
             return matches[-1].strip() if matches else pred
 
 
-    def make_prediction(self, question, db_schema,  model_name, prompt=None):
+    def make_prediction(self, question, db_schema,  model_name, prompt=None, task='SP'):
         if model_name not in self.model_name2pred_func:
             raise ValueError(
                 "Model not supported",
@@ -82,10 +92,13 @@ class ModelPrediction:
                 self.model_name2pred_func.keys(),
             )
         
+        if task == 'SP':
+            prompt = prompt or self.base_prompt
+        else: 
+            prompt = prompt or self.base_prompt_QA
 
-        prompt = prompt or self.base_prompt
-        #prompt = prompt.format(question=question, db_schema=db_schema)
-
+        prompt = prompt.format(question=question, db_schema=db_schema)
+        
         start_time = time.time()
         prediction = self.model_name2pred_func[model_name](prompt)
         end_time = time.time()
@@ -93,7 +106,6 @@ class ModelPrediction:
             prediction["response"]
         )
         prediction['time'] = end_time - start_time
-
         return prediction
 
    
@@ -133,14 +145,12 @@ class ModelPrediction:
             model_name = "openai/gpt-3.5-turbo-0125"
         elif "gpt-4o-mini" in model_name:
             model_name = "openai/gpt-4o-mini-2024-07-18"
-        elif "o1-mini" in model_name:
-            model_name = "openai/o1-mini-2024-09-12"
-        elif "QwQ" in model_name:
-            model_name = "together_ai/Qwen/QwQ-32B"
         elif "DeepSeek-R1-Distill-Llama-70B" in model_name:
             model_name = "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
         elif "llama-8" in model_name:
             model_name = "together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
+        elif "llama-70" in model_name:
+            model_name = "together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
         else:
             raise ValueError("Model forbidden")
 

test_get_db_schema_with_entries.py

@@ -0,0 +1,15 @@
+from utils_get_db_tables_info import utils_extract_db_schema_as_string
+
+def main():
+    db = utils_extract_db_schema_as_string(
+        db_id='', 
+        base_path='mytable_7.sqlite', 
+        normalize=True,
+        sql=None,
+        get_insert_into=True
+    )
+
+    print(db)
+
+if __name__ == '__main__':
+    main()

utils_get_db_tables_info.py

@@ -4,7 +4,7 @@ import re
 
 
 def utils_extract_db_schema_as_string(
-    db_id, base_path, normalize=False, sql: str | None = None
+    db_id, base_path, normalize=False, sql: str | None = None, get_insert_into: bool = False
 ):
     """
     Extracts the full schema of an SQLite database into a single string.
@@ -19,7 +19,7 @@ def utils_extract_db_schema_as_string(
     cursor = connection.cursor()
 
     # Get the schema entries based on the provided SQL query
-    schema_entries = _get_schema_entries(cursor, sql)
+    schema_entries = _get_schema_entries(cursor, sql, get_insert_into)
 
     # Combine all schema definitions into a single string
     schema_string = _combine_schema_entries(schema_entries, normalize)
@@ -27,27 +27,47 @@ def utils_extract_db_schema_as_string(
     return schema_string
 
 
-def _get_schema_entries(cursor, sql):
+
+def _get_schema_entries(cursor, sql=None, get_insert_into=False):
     """
-    Retrieves schema entries from the SQLite database.
+    Retrieves schema entries and optionally data entries from the SQLite database.
 
     :param cursor: SQLite cursor object.
     :param sql: Optional SQL query to filter specific tables.
-    :return: List of schema entries.
+    :param get_insert_into: Boolean flag to include INSERT INTO statements.
+    :return: List of schema and optionally data entries.
     """
+    entries = []
+
     if sql:
+        # Extract table names from the provided SQL query
         cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
         tables = [tbl[0] for tbl in cursor.fetchall() if tbl[0].lower() in sql.lower()]
-        if tables:
-            tbl_names = ", ".join(f"'{tbl}'" for tbl in tables)
-            query = f"SELECT sql FROM sqlite_master WHERE type='table' AND name IN ({tbl_names}) AND sql IS NOT NULL;"
-        else:
-            query = "SELECT sql FROM sqlite_master WHERE sql IS NOT NULL;"
     else:
-        query = "SELECT sql FROM sqlite_master WHERE sql IS NOT NULL;"
+        # Retrieve all table names
+        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+        tables = [tbl[0] for tbl in cursor.fetchall()]
+
+    for table in tables:
+        # Retrieve the CREATE TABLE statement for each table
+        cursor.execute(f"SELECT sql FROM sqlite_master WHERE type='table' AND name='{table}' AND sql IS NOT NULL;")
+        create_table_stmt = cursor.fetchone()
+        if create_table_stmt:
+            entries.append(create_table_stmt[0])
+
+        if get_insert_into:
+            # Retrieve all data from the table
+            cursor.execute(f"SELECT * FROM {table};")
+            rows = cursor.fetchall()
+            column_names = [description[0] for description in cursor.description]
+
+            # Generate INSERT INTO statements for each row
+            for row in rows:
+                values = ', '.join(f"'{str(value)}'" if isinstance(value, str) else str(value) for value in row)
+                insert_stmt = f"INSERT INTO {table} ({', '.join(column_names)}) VALUES ({values});"
+                entries.append(insert_stmt)
 
-    cursor.execute(query)
-    return cursor.fetchall()
+    return entries
 
 
 def _combine_schema_entries(schema_entries, normalize):
@@ -59,7 +79,7 @@ def _combine_schema_entries(schema_entries, normalize):
     :return: Combined schema string.
     """
     if not normalize:
-        return "\n".join(entry[0] for entry in schema_entries)
+        return "\n".join(entry for entry in schema_entries)
 
     return "\n".join(
         re.sub(
@@ -77,7 +97,7 @@ def _combine_schema_entries(schema_entries, normalize):
                         re.sub(
                             r"\s+",
                             " ",
-                            entry[0].replace("CREATE TABLE", "").replace("\t", " "),
+                            entry.replace("CREATE TABLE", "").replace("\t", " "),
                         ).strip(),
                     ),
                 ),