Overview
This module provides CLI tools for downloading, cleaning, sampling, validating, and saving balanced text-to-SQL evaluation datasets from the BookSQL Hugging Face repository.
Description
The data_utils.py module is a comprehensive data preparation pipeline for text-to-SQL evaluation with Ragas. It orchestrates several stages of dataset processing:
Dataset Download -- The download_booksql_dataset function uses huggingface_hub.snapshot_download to fetch the entire BookSQL dataset from the Exploration-Lab/BookSQL repository on Hugging Face Hub. It handles gated repository authentication errors, repository-not-found errors, and creates actual files rather than symlinks.
Data Loading and Cleaning -- The load_and_clean_data function reads the BookSQL train.json file, filters for train-split records, removes duplicates based on the (Query, SQL) pair, and reports the difficulty distribution across easy, medium, and hard levels.
Sampling -- Two sampling strategies are available:
- sample_by_difficulty performs simple random sampling of a specified number of records per difficulty level with a fixed random seed for reproducibility.
- validate_samples performs sampling with SQL execution validation. It shuffles data by difficulty level and iterates through records, executing each SQL query against the database via execute_and_validate_query (imported from validate_sql_dataset). Only queries that execute successfully (and optionally return non-empty, non-null data) are included, up to the requested sample count.
Dataset Creation -- The create_sample_dataset orchestrator function chains loading, sampling (with optional validation), combination via pd.concat, and saving. The save_results function performs a final deduplication check, shuffles the combined dataset, and writes it to CSV.
The CLI supports two top-level commands:
- --download-data triggers the Hugging Face dataset download.
- --create-sample creates a balanced sample CSV with configurable --samples count per level, optional --validate for SQL execution testing, and --require-data to filter for queries that return actual data rows.
Usage
Import this module when you need to prepare a balanced, optionally validated text-to-SQL evaluation dataset from BookSQL. Use download_booksql_dataset to obtain the raw data, create_sample_dataset to produce a ready-to-use evaluation CSV, or individual functions like load_and_clean_data and validate_samples for custom pipelines.
Code Reference
Source Location
Signature
def download_booksql_dataset() -> bool
def validate_query_data(query_data: Dict[str, Any], require_data: bool = False) -> bool
def load_and_clean_data(input_file: str) -> DataFrame
def sample_by_difficulty(data: DataFrame, level: str, samples_per_level: int, random_seed: int) -> DataFrame
def validate_samples(data: DataFrame, level: str, samples_per_level: int,
random_seed: int, require_data: bool = False) -> DataFrame
def save_results(data: DataFrame, output_dir: str, output_filename: str, random_seed: int) -> bool
def create_sample_dataset(
input_file: str = "BookSQL-files/BookSQL/train.json",
output_dir: str = "datasets",
output_filename: str = "booksql_sample.csv",
samples_per_level: int = 10,
random_seed: int = 42,
validate_queries: bool = False,
require_data: bool = False,
) -> bool
def main() -> None
Import
from ragas_examples.text2sql.data_utils import (
download_booksql_dataset,
create_sample_dataset,
load_and_clean_data,
sample_by_difficulty,
validate_samples,
validate_query_data,
save_results,
)
I/O Contract
Inputs
download_booksql_dataset
| Name |
Type |
Required |
Description
|
| (none) |
-- |
-- |
No parameters; downloads from the hardcoded Exploration-Lab/BookSQL repo
|
load_and_clean_data
| Name |
Type |
Required |
Description
|
| input_file |
str |
Yes |
Path to the BookSQL train.json file
|
sample_by_difficulty
| Name |
Type |
Required |
Description
|
| data |
DataFrame |
Yes |
DataFrame containing the cleaned data
|
| level |
str |
Yes |
Difficulty level: "easy", "medium", or "hard"
|
| samples_per_level |
int |
Yes |
Number of samples to take per level
|
| random_seed |
int |
Yes |
Random seed for reproducible sampling
|
validate_samples
| Name |
Type |
Required |
Description
|
| data |
DataFrame |
Yes |
DataFrame containing the cleaned data
|
| level |
str |
Yes |
Difficulty level: "easy", "medium", or "hard"
|
| samples_per_level |
int |
Yes |
Number of valid samples to find
|
| random_seed |
int |
Yes |
Random seed for reproducible shuffling
|
| require_data |
bool |
No |
If True, only include queries that return actual data rows (default: False)
|
create_sample_dataset
| Name |
Type |
Required |
Description
|
| input_file |
str |
No |
Path to train.json (default: "BookSQL-files/BookSQL/train.json")
|
| output_dir |
str |
No |
Directory for output CSV (default: "datasets")
|
| output_filename |
str |
No |
Name of output CSV file (default: "booksql_sample.csv")
|
| samples_per_level |
int |
No |
Samples per difficulty level (default: 10)
|
| random_seed |
int |
No |
Random seed (default: 42)
|
| validate_queries |
bool |
No |
Enable SQL query validation (default: False)
|
| require_data |
bool |
No |
Only include data-returning queries when validating (default: False)
|
Outputs
download_booksql_dataset
| Name |
Type |
Description
|
| return |
bool |
True if dataset download succeeded, False otherwise
|
load_and_clean_data
| Name |
Type |
Description
|
| return |
DataFrame |
Cleaned train-split DataFrame with duplicates removed
|
sample_by_difficulty / validate_samples
| Name |
Type |
Description
|
| return |
DataFrame |
Sampled (and optionally validated) DataFrame for the specified difficulty level
|
create_sample_dataset / save_results
| Name |
Type |
Description
|
| return |
bool |
True if the pipeline completed successfully and CSV was saved
|
Usage Examples
CLI Usage
# Download the BookSQL dataset from Hugging Face
python data_utils.py --download-data
# Create a sample dataset with 15 samples per difficulty level
python data_utils.py --create-sample --samples 15
# Create a validated sample (only queries that execute successfully)
python data_utils.py --create-sample --validate
# Create a validated sample, only including queries that return data
python data_utils.py --create-sample --validate --require-data
Programmatic Usage
from data_utils import create_sample_dataset, load_and_clean_data
# Create a balanced sample dataset with validation
success = create_sample_dataset(
input_file="BookSQL-files/BookSQL/train.json",
output_dir="datasets",
output_filename="booksql_validated.csv",
samples_per_level=10,
random_seed=42,
validate_queries=True,
require_data=True,
)
if success:
print("Dataset created successfully")
# Or use individual steps for a custom pipeline
df = load_and_clean_data("BookSQL-files/BookSQL/train.json")
print(f"Loaded {len(df)} cleaned records")
Related Pages
