fixed extraction and cleaned output

duckdb-nsql/eval/predict.py

@@ -294,7 +294,6 @@ def predict(
         for prediction, model_response in generated_responses:
             prediction = re.sub(r"[\s\t\n]+", " ", prediction)
             token_lengths.append(len(tokenizer(prediction).input_ids))
-            console.print(f"[blue]Prompt:[/blue] {model_response.final_prompt}")
             console.print(f"[red]Prediction:[/red] {prediction}")
             if data[i].get("query") or data[i].get("sql"):
                 console.print(

duckdb-nsql/eval/prompt_formatters.py

@@ -65,22 +65,77 @@ class RajkumarFormatter:
 
     @classmethod
     def format_model_output(cls, output_sql: str, prompt: str) -> str:
-        """Format model output."""
-        time.sleep(10)
-        clean_sql = (output_sql
-            .replace('sql\n', '')
-            .replace('```sql\n', '')
-            .replace('```duckdb\n', '')
-            .replace('```\n', '')
-            .replace('```', '')).strip()
-
-        if clean_sql.find(';') != -1:
-            clean_sql[:clean_sql.find(';')].strip()
-
-        if not clean_sql.endswith(";"):
-            clean_sql += ";"
-
-        return clean_sql
+        def clean_code_block(block: str) -> str:
+            """Clean a code block by removing markdown syntax and extra whitespace."""
+            # Remove markdown indicators and common SQL prefixes
+            cleaned = (block
+                       .replace('```sql\n', '')
+                       .replace('```duckdb\n', '')
+                       .replace('```\n', '')
+                       .replace('```', '')
+                       .strip())
+
+            return cleaned
+
+        def ensure_semicolon(sql: str) -> str:
+            """Ensure the SQL query ends with exactly one semicolon."""
+            sql = sql.strip()
+            # Remove any existing trailing semicolons
+            while sql.endswith(';'):
+                sql = sql[:-1].strip()
+            # Add back exactly one semicolon
+            return sql + ";"
+
+        # First, try to find SQL-specific code blocks
+        sql_blocks = []
+        start_pos = 0
+        while True:
+            start = output_sql.find('```sql', start_pos)
+            if start == -1:
+                start = output_sql.find('```duckdb', start_pos)
+            if start == -1:
+                break
+
+            end = output_sql.find('```', start + 4)
+            if end == -1:
+                break
+
+            sql_blocks.append(output_sql[start:end+3])
+            start_pos = end + 3
+
+        # If SQL blocks found, use the last one
+        if sql_blocks:
+            return ensure_semicolon(clean_code_block(sql_blocks[-1]))
+
+        # If no SQL blocks, look for generic code blocks
+        generic_blocks = []
+        start_pos = 0
+        while True:
+            start = output_sql.find('```', start_pos)
+            if start == -1:
+                break
+
+            end = output_sql.find('```', start + 3)
+            if end == -1:
+                break
+
+            block = output_sql[start:end+3]
+            # Skip if this is actually an SQL block (we already handled those)
+            if not block.startswith('```sql') and not block.startswith('```duckdb'):
+                generic_blocks.append(block)
+            start_pos = end + 3
+
+        # If generic blocks found, use the last one
+        if generic_blocks:
+            return ensure_semicolon(clean_code_block(generic_blocks[-1]))
+
+        # If no code blocks found at all, take everything up to first semicolon
+        semicolon_pos = output_sql.find(';')
+        if semicolon_pos != -1:
+            return ensure_semicolon(output_sql[:semicolon_pos].strip())
+
+        # If no semicolon found, use the entire text
+        return ensure_semicolon(output_sql.strip())
 
     @classmethod
     def format_gold_output(cls, output_sql: str) -> str:
@@ -503,7 +558,7 @@ DuckDB Other Keywords:
 `DESC`: The keyword `desc` is used in SQL to describe the DESCENDING order in which query results should be sorted, often associated with the `ORDER BY` clause., Examples: ['SELECT name FROM employees ORDER BY salary DESC;', 'CREATE TABLE example (id INTEGER, name VARCHAR);', 'DESCRIBE example;', 'DESCRIBE SELECT * FROM example WHERE id < 10 ORDER BY name DESC;']
 `IS`: The "IS" keyword is used in SQL to perform tests on values to check for NULL values or to use as part of statements like IS DISTINCT FROM which can handle NULLs in equality comparisons., Examples: ['SELECT 4 IS DISTINCT FROM NULL;', 'SELECT 4 IS NOT DISTINCT FROM 4;', 'SELECT NULL IS NULL;', 'SELECT NULL IS NOT NULL;']
 `IN`: The `IN` keyword is used in SQL to specify a list of discrete values for a column to match against, typically in a `WHERE` clause, allowing for multiple specific conditions to be evaluated at once., Examples: ["SELECT * FROM employees WHERE department IN ('HR', 'Engineering', 'Marketing');", 'SELECT id, name FROM students WHERE grade IN (10, 11, 12);', "DELETE FROM orders WHERE order_status IN ('Cancelled', 'Returned');", "UPDATE items SET status = 'Unavailable' WHERE item_id IN (1001, 1002, 1003);", "SELECT * FROM logs WHERE severity IN ('ERROR', 'CRITICAL') ORDER BY timestamp DESC;"]
-`ALL`: The `ALL` keyword in SQL specifies that operations should retain all duplicate rows, as seen in commands like `UNION ALL`, `INTERSECT ALL`, and `EXCEPT ALL`, which follow bag semantics instead of eliminating duplicates., Examples: ['UNION ALL\n\n```sql\nSELECT * FROM range(2) t1(x)\nUNION ALL\nSELECT * FROM range(3) t2(x);\n```\nThis example demonstrates using `UNION ALL` to combine rows from two queries without eliminating duplicates.', 'INTERSECT ALL\n\n```sql\nSELECT unnest([5, 5, 6, 6, 6, 6, 7, 8]) AS x\nINTERSECT ALL\nSELECT unnest([5, 6, 6, 7, 7, 9]);\n```\nThis example shows using `INTERSECT ALL` to select rows that are present in both result sets, keeping duplicate values.', 'EXCEPT ALL\n\n```sql\nSELECT unnest([5, 5, 6, 6, 6, 6, 7, 8]) AS x\nEXCEPT ALL\nSELECT unnest([5, 6, 6, 7, 7, 9]);\n```\nThis example illustrates `EXCEPT ALL`, which selects all rows present in the first query but not in the second, without removing duplicates.', 'ORDER BY ALL\n\n```sql\nSELECT *\nFROM addresses\nORDER BY ALL;\n```\nThis SQL command uses `ORDER BY ALL` to sort the result set by all columns sequentially from left to right.']
+`ALL`: The `ALL` keyword in SQL specifies that operations should retain all duplicate rows, as seen in commands like `UNION ALL`, `INTERSECT ALL`, and `EXCEPT ALL`, which follow bag semantics instead of eliminating duplicates., Examples: ['UNION ALL ```sql\nSELECT * FROM range(2) t1(x)\nUNION ALL\nSELECT * FROM range(3) t2(x);\n```\nThis example demonstrates using `UNION ALL` to combine rows from two queries without eliminating duplicates.', 'INTERSECT ALL ```sql\nSELECT unnest([5, 5, 6, 6, 6, 6, 7, 8]) AS x\nINTERSECT ALL\nSELECT unnest([5, 6, 6, 7, 7, 9]);\n```\nThis example shows using `INTERSECT ALL` to select rows that are present in both result sets, keeping duplicate values.', 'EXCEPT ALL ```sql\nSELECT unnest([5, 5, 6, 6, 6, 6, 7, 8]) AS x\nEXCEPT ALL\nSELECT unnest([5, 6, 6, 7, 7, 9]);\n```\nThis example illustrates `EXCEPT ALL`, which selects all rows present in the first query but not in the second, without removing duplicates.', 'ORDER BY ALL ```sql\nSELECT *\nFROM addresses\nORDER BY ALL;\n```\nThis SQL command uses `ORDER BY ALL` to sort the result set by all columns sequentially from left to right.']
 `LIKE`: The `LIKE` expression is used to determine if a string matches a specified pattern, allowing wildcard characters such as `_` to represent any single character and `%` to match any sequence of characters., Examples: ["SELECT 'abc' LIKE 'abc'; -- true", "SELECT 'abc' LIKE 'a%'; -- true", "SELECT 'abc' LIKE '_b_'; -- true", "SELECT 'abc' LIKE 'c'; -- false", "SELECT 'abc' LIKE 'c%'; -- false", "SELECT 'abc' LIKE '%c'; -- true", "SELECT 'abc' NOT LIKE '%c'; -- false", "SELECT 'abc' ILIKE '%C'; -- true"]
 `IF`: The `IF` keyword is used in conditional constructs, most commonly found in the `IF NOT EXISTS` or `IF EXISTS` clauses in SQL commands to prevent errors or skip operations, such as creating or dropping a database or table when certain conditions are met., Examples: ['CREATE DATABASE IF NOT EXISTS ducks_db;', 'CREATE TABLE IF NOT EXISTS t1 (i INTEGER, j INTEGER);', 'CREATE TABLE t1 (id INTEGER, PRIMARY KEY (id), UNIQUE (id)) IF NOT EXISTS;']
 `EXISTS`: The `EXISTS` operator is used to determine if a subquery returns any rows, returning `true` if at least one row exists and `false` otherwise., Examples: ["SELECT EXISTS (FROM grades WHERE course = 'Math') AS math_grades_present;", "SELECT EXISTS (FROM grades WHERE course = 'History') AS history_grades_present;"]
@@ -528,10 +583,9 @@ Question:
 Here is the question or an instruction the user provided:
 {question}
 
-Write a DuckDB SQL query for the given question!
+Write a DuckDB SQL query for the given question! Make sure to only response with the SQL query and wrap it in ```sql\n<ANSWER>``` markdown code tags.
 
 Answer:
-```
 """
 
     @classmethod

duckdb-nsql/eval/text_to_sql.py

@@ -219,11 +219,6 @@ def _run_manifest(
 ) -> TextToSQLModelResponse:
     """Run manifest for prompt format."""
     logger.info(f"PARAMS: {manifest_params}")
-    if isinstance(prompt, list):
-        for p in prompt:
-            logger.info(f"PROMPT: {p['role']}: {p['content']}")
-    else:
-        logger.info(f"PROMPT: {prompt}")
     start_time = time.time()
     # Run result
     response = cast(
@@ -248,7 +243,6 @@ def _run_manifest(
         cast(str, response.get_response()), prompt
     )
 
-    logger.info(f"RAW OUTPUT: {response.get_response()}")
     for token in stop_sequences:
         sql_query = sql_query.split(token)[0]
     logger.info(f"OUTPUT: {sql_query}")

evaluation_logic.py

@@ -86,8 +86,8 @@ def save_evaluation(inference_api, model_name, prompt_format, custom_prompt, met
 def run_prediction(inference_api, model_name, prompt_format, custom_prompt, output_file):
     dataset_path = str(eval_dir / "data/dev.json")
     table_meta_path = str(eval_dir / "data/tables.json")
-    stop_tokens = [';']
-    max_tokens = 32000
+    stop_tokens = ['`<|dummy|>`']
+    max_tokens = 1000
     temperature = 0
     num_beams = -1
     manifest_client = inference_api

requirements.txt

@@ -28,4 +28,4 @@ diffusers
 deepspeed
 sentence_transformers
 tqdm
-pydantic
+pydantic