text2sql-eval-toolkit 1.0.0__py3-none-any.whl

text2sql_eval_toolkit/config_args.py

@@ -0,0 +1,65 @@
+#
+# Copyright IBM Corp. 2025 - 2026
+# SPDX-License-Identifier: Apache-2.0
+#
+
+import argparse
+from text2sql_eval_toolkit.utils import get_available_benchmarks
+
+DEFAULT_MODEL_NAMES = [
+    "wxai:meta-llama/llama-3-3-70b-instruct",
+    "wxai:ibm/granite-4-h-small",
+    "wxai:meta-llama/llama-4-maverick-17b-128e-instruct-fp8",
+    "wxai:openai/gpt-oss-120b",
+    # "vllm:Qwen/Qwen2.5-1.5B-Instruct",
+    # "rits:microsoft/Phi-4-reasoning",
+    # "anthropic:claude-opus-4-1-20250805"
+]
+
+# Default models for agentic baselines (single model to reduce runtime)
+DEFAULT_AGENTIC_MODELS = ["wxai:openai/gpt-oss-120b"]
+
+
+def add_common_arguments(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
+    parser.add_argument(
+        "benchmark_id",
+        help="Benchmark ID. Available benchmarks: "
+        + ", ".join(get_available_benchmarks()),
+    )
+    parser.add_argument(
+        "--model_names",
+        nargs="+",
+        default=DEFAULT_MODEL_NAMES,
+        help=F"Optional list of model names. Defaults to {DEFAULT_MODEL_NAMES}.",
+    )
+    parser.add_argument(
+        "--decoding_method",
+        type=str,
+        default="greedy",
+        help="Decoding method (e.g., greedy, sampling). Default: greedy. "
+        "Note: Only supported by WatsonX legacy API; ignored by Chat APIs (WatsonX Chat, Claude, VLLM).",
+    )
+    parser.add_argument(
+        "--max_new_tokens",
+        type=int,
+        default=256,
+        help="Maximum number of new tokens to generate. Default: 256. "
+        "Automatically converted to 'max_tokens' for Chat APIs.",
+    )
+    parser.add_argument(
+        "--stop_sequences",
+        nargs="+",
+        default=[],
+        help="Stop sequences. Pass each one as a separate string. Default: []. "
+        "Note: Supported by Claude and VLLM, but not by WatsonX Chat API.",
+    )
+    parser.add_argument(
+        "--agentic_models",
+        nargs="+",
+        default=None,
+        help=f"List of models to use for agentic baselines when using --run_all_baselines or --pipeline_type agentic. "
+        f"Defaults to: {', '.join(DEFAULT_AGENTIC_MODELS)}. "
+        "If not specified, uses all models from --model_names for run_experiment.py, "
+        f"or DEFAULT_AGENTIC_MODELS for run_all_benchmarks.py.",
+    )
+    return parser

text2sql_eval_toolkit/data/__init__.py

@@ -0,0 +1,4 @@
+#
+# Package data for text2sql_eval_toolkit.
+#
+

text2sql_eval_toolkit/data/benchmarks.json

@@ -0,0 +1,69 @@
+{
+    "bird_mini_dev_sqlite": {
+        "name": "bird_mini_dev_sqlite",
+        "description": "BIRD-SQL Mini-Dev in SQLite https://github.com/bird-bench/mini_dev",
+        "data": "benchmarks/bird_mini_dev_sqlite.json",
+        "schema": "benchmarks/bird_mini_dev_sqlite-schema.json",
+        "predictions": "results/bird_mini_dev_sqlite-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/bird/dev_databases"
+        }
+    },
+    "bird_mini_dev_postgres": {
+        "name": "bird_mini_dev_postgres",
+        "description": "BIRD-SQL Mini-Dev in PostgreSQL https://github.com/bird-bench/mini_dev",
+        "data": "benchmarks/bird_mini_dev_postgres.json",
+        "schema": "benchmarks/bird_mini_dev_postgres-schema.json",
+        "predictions": "results/bird_mini_dev_postgres-predictions.json",
+        "db_engine": {
+            "db_type": "postgres",
+            "schema_name": "public",
+            "connection_string_env_var": "POSTGRES_CONNECTION_STRING"
+        }
+    },
+    "beaver": {
+        "name": "beaver",
+        "description": "Beaver benchmark https://peterbaile.github.io/beaver/",
+        "data": "benchmarks/beaver.json",
+        "schema": "benchmarks/beaver-schema.json",
+        "predictions": "results/beaver-predictions.json",
+        "db_engine": {
+            "db_type": "mysql",
+            "connection_string_env_var": "MYSQL_CONNECTION_STRING"
+        }
+    },
+    "archer_en_dev": {
+        "name": "archer_en_dev",
+        "description": "Archer English Dev Set https://sig4kg.github.io/archer-bench/",
+        "data": "benchmarks/archer_en_dev.json",
+        "schema": "benchmarks/archer-schema.json",
+        "predictions": "results/archer_en_dev-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/archer/database"
+        }
+    },
+    "spider_dev": {
+        "name": "spider_dev",
+        "description": "Spider Dev Set - Full 1,034 questions https://yale-lily.github.io/spider",
+        "data": "benchmarks/spider-dev-converted.json",
+        "schema": "benchmarks/spider-dev-schema.json",
+        "predictions": "results/spider_dev-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/spider/database"
+        }
+    },
+    "spider_realistic": {
+        "name": "spider_realistic",
+        "description": "Spider Realistic Dataset In Structure-Grounded Pretraining for Text-to-SQL https://zenodo.org/records/5205322",
+        "data": "benchmarks/spider-realistic.json",
+        "schema": "benchmarks/spider-dev-schema.json",
+        "predictions": "results/spider_realistic-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/spider/database"
+        }
+    }
+}

text2sql_eval_toolkit/data/test-benchmarks.json

@@ -0,0 +1,69 @@
+{
+    "bird_mini_dev_sqlite_test_50": {
+        "name": "bird_mini_dev_sqlite_test_50",
+        "description": "Test sample of 50 BIRD-SQL Mini-Dev in SQLite https://github.com/bird-bench/mini_dev",
+        "data": "benchmarks/test_benchmarks/bird_mini_dev_sqlite_test_50.json",
+        "schema": "benchmarks/bird_mini_dev_sqlite-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/bird_mini_dev_sqlite_test_50-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/bird/dev_databases"
+        }
+    },
+    "bird_mini_dev_postgres_test_50": {
+        "name": "bird_mini_dev_postgres_test_50",
+        "description": "Test sample of 50 BIRD-SQL Mini-Dev in PostgreSQL https://github.com/bird-bench/mini_dev",
+        "data": "benchmarks/test_benchmarks/bird_mini_dev_postgres_test_50.json",
+        "schema": "benchmarks/bird_mini_dev_postgres-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/bird_mini_dev_postgres_test_50-predictions.json",
+        "db_engine": {
+            "db_type": "postgres",
+            "schema_name": "public",
+            "connection_string_env_var": "POSTGRES_CONNECTION_STRING"
+        }
+    },
+    "spider_dev_test_50": {
+        "name": "spider_dev_test_50",
+        "description": "Test sample of 50 Spider Dev questions https://yale-lily.github.io/spider",
+        "data": "benchmarks/test_benchmarks/spider_dev_test_50.json",
+        "schema": "benchmarks/spider-dev-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/spider_dev_test_50-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/spider/database"
+        }
+    },
+    "beaver_test_10": {
+        "name": "beaver_test_10",
+        "description": "Test sample of 10 Beaver benchmark questions https://peterbaile.github.io/beaver/",
+        "data": "benchmarks/test_benchmarks/beaver_test_10.json",
+        "schema": "benchmarks/beaver-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/beaver_test_10-predictions.json",
+        "db_engine": {
+            "db_type": "mysql",
+            "connection_string_env_var": "MYSQL_CONNECTION_STRING"
+        }
+    },
+    "archer_en_dev_test_10": {
+        "name": "archer_en_dev_test_10",
+        "description": "Test sample of 10 Archer English Dev questions https://sig4kg.github.io/archer-bench/",
+        "data": "benchmarks/test_benchmarks/archer_en_dev_test_10.json",
+        "schema": "benchmarks/archer-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/archer_en_dev_test_10-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/dbs/archer/database"
+        }
+    },
+    "bird_sqlite_test_benchmark": {
+        "name": "bird_sqlite_test_benchmark",
+        "description": "3 records from BIRD mini-dev used for testing this code base",
+        "data": "benchmarks/test_benchmarks/bird_sqlite_test_benchmark.json",
+        "schema": "benchmarks/test_benchmarks/bird_sqlite_test_benchmark-schema.json",
+        "predictions": "benchmarks/test_benchmarks/results/bird_sqlite_test_benchmark-predictions.json",
+        "db_engine": {
+            "db_type": "sqlite",
+            "db_folder": "benchmarks/test_benchmarks/db"
+        }
+    }
+}

text2sql_eval_toolkit/env_loader.py

@@ -0,0 +1,55 @@
+#
+# Copyright IBM Corp. 2025 - 2026
+# SPDX-License-Identifier: Apache-2.0
+#
+
+"""
+Environment variable loader for text2sql-eval-toolkit.
+
+Automatically loads environment variables from .env file in project root.
+"""
+
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+
+
+def load_env():
+    """
+    Load environment variables from .env file in the project root.
+    
+    This function looks for .env starting from the current working directory
+    and searches upwards through parent directories. This allows it to work
+    correctly when the toolkit is used as a dependency in other projects.
+    
+    Returns:
+        bool: True if .env file was found and loaded, False otherwise
+    """
+    # Start from current working directory and search upwards
+    current_dir = Path.cwd()
+    
+    # Search upwards for .env file (max 10 levels)
+    for _ in range(10):
+        env_path = current_dir / ".env"
+        if env_path.exists():
+            load_dotenv(env_path, override=False)  # Don't override existing env vars
+            return True
+        
+        # Move to parent directory
+        parent = current_dir.parent
+        if parent == current_dir:  # Reached root
+            break
+        current_dir = parent
+    
+    # Try to load from ~/.env as fallback
+    home_env = Path.home() / ".env"
+    if home_env.exists():
+        load_dotenv(home_env, override=False)
+        return True
+    
+    return False
+
+
+# Auto-load on import
+load_env()
+

text2sql_eval_toolkit/evaluation/__init__.py

@@ -0,0 +1,26 @@
+#
+# Evaluation subpackage public API for text2sql_eval_toolkit.
+#
+
+from unitxt.text2sql_utils import (
+    compare_result_dfs,
+    compare_dfs_bird_eval_logic,
+    is_sqlglot_parsable,
+    is_sqlparse_parsable,
+    sqlglot_parsed_queries_equivalent,
+    sqlglot_optimized_equivalence,
+    sqlparse_queries_equivalent,
+    sql_exact_match,
+)
+
+__all__ = [
+    "compare_result_dfs",
+    "compare_dfs_bird_eval_logic",
+    "is_sqlglot_parsable",
+    "is_sqlparse_parsable",
+    "sqlglot_parsed_queries_equivalent",
+    "sqlglot_optimized_equivalence",
+    "sqlparse_queries_equivalent",
+    "sql_exact_match",
+]
+