chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,467 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Data utilities for Text-to-SQL evaluation with Ragas.
|
||||
|
||||
This module provides CLI tools to download and prepare datasets for
|
||||
text-to-SQL evaluation workflows.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List
|
||||
|
||||
# Load environment variables from ragas root
|
||||
try:
|
||||
from dotenv import load_dotenv
|
||||
# Load .env from ragas root directory (3 levels up from this file)
|
||||
ragas_root = Path(__file__).parent.parent.parent.parent
|
||||
env_path = ragas_root / ".env"
|
||||
load_dotenv(env_path)
|
||||
except ImportError:
|
||||
# dotenv is optional, continue without it
|
||||
pass
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(levelname)s: %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError
|
||||
except ImportError:
|
||||
logger.error("huggingface_hub is required. Install with: pip install huggingface_hub")
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
import pandas as pd
|
||||
from pandas import DataFrame
|
||||
except ImportError:
|
||||
logger.error("pandas is required. Install with: pip install pandas")
|
||||
sys.exit(1)
|
||||
|
||||
# Import validation functions from validate_sql_dataset.py
|
||||
try:
|
||||
from .validate_sql_dataset import execute_and_validate_query
|
||||
except ImportError:
|
||||
logger.error("validate_sql_dataset.py not found in the same directory")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def download_booksql_dataset() -> bool:
|
||||
"""
|
||||
Download the BookSQL dataset from Hugging Face Hub to ./BookSQL-files directory.
|
||||
|
||||
Returns:
|
||||
bool: True if download successful, False otherwise
|
||||
|
||||
Note:
|
||||
This dataset is gated and requires accepting terms on the Hugging Face Hub.
|
||||
You need to:
|
||||
1. Visit https://huggingface.co/datasets/Exploration-Lab/BookSQL
|
||||
2. Accept the terms and conditions
|
||||
3. Authenticate with: huggingface-cli login
|
||||
"""
|
||||
repo_id = "Exploration-Lab/BookSQL"
|
||||
local_dir = "BookSQL-files"
|
||||
|
||||
# Create local directory if it doesn't exist
|
||||
Path(local_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
logger.info(f"Downloading BookSQL dataset to {local_dir}")
|
||||
logger.info(f"Repository: {repo_id}")
|
||||
|
||||
try:
|
||||
# Download the entire repository
|
||||
downloaded_path = snapshot_download(
|
||||
repo_id=repo_id,
|
||||
repo_type="dataset",
|
||||
local_dir=local_dir,
|
||||
local_dir_use_symlinks=False # Create actual files, not symlinks
|
||||
)
|
||||
|
||||
logger.info(f"Successfully downloaded dataset to: {downloaded_path}")
|
||||
|
||||
# List downloaded files
|
||||
dataset_path = Path(local_dir)
|
||||
files = list(dataset_path.rglob("*"))
|
||||
logger.info(f"Downloaded {len(files)} files")
|
||||
for file in sorted(files)[:5]: # Show first 5 files
|
||||
if file.is_file():
|
||||
logger.info(f" {file.relative_to(dataset_path)}")
|
||||
if len(files) > 5:
|
||||
logger.info(f" ... and {len(files) - 5} more files")
|
||||
|
||||
return True
|
||||
|
||||
except GatedRepoError:
|
||||
logger.error("This dataset is gated and requires authentication")
|
||||
logger.error("Please follow these steps:")
|
||||
logger.error("1. Visit: https://huggingface.co/datasets/Exploration-Lab/BookSQL")
|
||||
logger.error("2. Accept the terms and conditions")
|
||||
logger.error("3. Run: huggingface-cli login")
|
||||
logger.error("4. Try downloading again")
|
||||
return False
|
||||
|
||||
except RepositoryNotFoundError:
|
||||
logger.error(f"Repository '{repo_id}' not found")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error downloading dataset: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def validate_query_data(query_data: Dict[str, Any], require_data: bool = False) -> bool:
|
||||
"""
|
||||
Validate a single query by executing it against the database.
|
||||
|
||||
Args:
|
||||
query_data: Dictionary containing query information (query, sql, level, split)
|
||||
require_data: If True, only accept queries that return actual data
|
||||
|
||||
Returns:
|
||||
bool: True if query is valid (and optionally returns data), False otherwise
|
||||
"""
|
||||
try:
|
||||
result = execute_and_validate_query(query_data)
|
||||
|
||||
if not result['execution_success']:
|
||||
return False
|
||||
|
||||
if require_data:
|
||||
# Only accept queries that return actual data (not empty or null values)
|
||||
return result.get('result_type') == 'has_data'
|
||||
else:
|
||||
# Accept any successful query execution
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Error validating query: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def load_and_clean_data(input_file: str) -> DataFrame:
|
||||
"""
|
||||
Load JSON data and remove duplicates.
|
||||
|
||||
Args:
|
||||
input_file: Path to the BookSQL train.json file
|
||||
|
||||
Returns:
|
||||
DataFrame: Cleaned train data with duplicates removed
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If input file doesn't exist
|
||||
json.JSONDecodeError: If JSON is invalid
|
||||
"""
|
||||
input_path = Path(input_file)
|
||||
|
||||
if not input_path.exists():
|
||||
raise FileNotFoundError(f"Input file '{input_file}' not found")
|
||||
|
||||
logger.info(f"Loading data from {input_file}")
|
||||
|
||||
# Load JSON data
|
||||
with open(input_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
logger.info(f"Loaded {len(data)} total records")
|
||||
|
||||
# Convert to DataFrame and filter for train split
|
||||
df = pd.DataFrame(data)
|
||||
train_df = df[df['split'] == 'train'].copy()
|
||||
logger.info(f"Found {len(train_df)} train records")
|
||||
|
||||
# Remove duplicates based on Query + SQL combination
|
||||
original_count = len(train_df)
|
||||
train_df = train_df.drop_duplicates(subset=['Query', 'SQL'], keep='first')
|
||||
duplicate_count = original_count - len(train_df)
|
||||
|
||||
if duplicate_count > 0:
|
||||
logger.info(f"Removed {duplicate_count} duplicate records")
|
||||
logger.info(f"{len(train_df)} unique records remaining")
|
||||
|
||||
# Show difficulty distribution
|
||||
level_counts = train_df['Levels'].value_counts()
|
||||
logger.info("Difficulty distribution after deduplication:")
|
||||
for level, count in level_counts.items():
|
||||
logger.info(f" {level}: {count} records")
|
||||
|
||||
return train_df
|
||||
|
||||
|
||||
def sample_by_difficulty(data: DataFrame, level: str, samples_per_level: int, random_seed: int) -> DataFrame:
|
||||
"""
|
||||
Sample data for a specific difficulty level.
|
||||
|
||||
Args:
|
||||
data: DataFrame containing the data
|
||||
level: Difficulty level ('easy', 'medium', 'hard')
|
||||
samples_per_level: Number of samples to take
|
||||
random_seed: Random seed for reproducible sampling
|
||||
|
||||
Returns:
|
||||
DataFrame: Sampled data for the specified level
|
||||
"""
|
||||
level_data = data[data['Levels'] == level]
|
||||
|
||||
if len(level_data) == 0:
|
||||
logger.warning(f"No '{level}' records found, skipping")
|
||||
return pd.DataFrame()
|
||||
|
||||
if len(level_data) < samples_per_level:
|
||||
logger.warning(f"Only {len(level_data)} '{level}' records available, using all")
|
||||
return level_data
|
||||
else:
|
||||
sampled = level_data.sample(n=samples_per_level, random_state=random_seed)
|
||||
logger.info(f"Sampled {len(sampled)} '{level}' records")
|
||||
return sampled
|
||||
|
||||
|
||||
def validate_samples(data: DataFrame, level: str, samples_per_level: int,
|
||||
random_seed: int, require_data: bool = False) -> DataFrame:
|
||||
"""
|
||||
Sample and validate data for a specific difficulty level.
|
||||
|
||||
Args:
|
||||
data: DataFrame containing the data
|
||||
level: Difficulty level ('easy', 'medium', 'hard')
|
||||
samples_per_level: Number of samples to find
|
||||
random_seed: Random seed for reproducible sampling
|
||||
require_data: If True, only include queries that return data
|
||||
|
||||
Returns:
|
||||
DataFrame: Validated samples for the specified level
|
||||
"""
|
||||
level_data = data[data['Levels'] == level]
|
||||
|
||||
if len(level_data) == 0:
|
||||
logger.warning(f"No '{level}' records found, skipping")
|
||||
return pd.DataFrame()
|
||||
|
||||
logger.info(f"Validating '{level}' queries to find {samples_per_level} valid samples")
|
||||
|
||||
# Shuffle data for random sampling during validation
|
||||
shuffled_data = level_data.sample(frac=1, random_state=random_seed).reset_index(drop=True)
|
||||
|
||||
valid_samples = []
|
||||
checked_count = 0
|
||||
|
||||
for idx, row in shuffled_data.iterrows():
|
||||
checked_count += 1
|
||||
|
||||
# Prepare query data for validation
|
||||
query_data = {
|
||||
'index': idx,
|
||||
'query': row['Query'],
|
||||
'sql': row['SQL'],
|
||||
'level': row['Levels'],
|
||||
'split': row['split']
|
||||
}
|
||||
|
||||
if validate_query_data(query_data, require_data):
|
||||
valid_samples.append(row)
|
||||
|
||||
# Stop if we have enough samples
|
||||
if len(valid_samples) >= samples_per_level:
|
||||
break
|
||||
|
||||
if len(valid_samples) == 0:
|
||||
logger.warning(f"No valid '{level}' queries found, skipping this level")
|
||||
return pd.DataFrame()
|
||||
elif len(valid_samples) < samples_per_level:
|
||||
logger.warning(f"Only found {len(valid_samples)} valid '{level}' queries out of {samples_per_level} requested")
|
||||
else:
|
||||
logger.info(f"Found {len(valid_samples)} valid '{level}' queries")
|
||||
|
||||
return pd.DataFrame(valid_samples) if valid_samples else pd.DataFrame()
|
||||
|
||||
|
||||
def save_results(data: DataFrame, output_dir: str, output_filename: str, random_seed: int) -> bool:
|
||||
"""
|
||||
Save final dataset to CSV.
|
||||
|
||||
Args:
|
||||
data: Final dataset to save
|
||||
output_dir: Directory to save the output CSV
|
||||
output_filename: Name of the output CSV file
|
||||
random_seed: Random seed for final shuffle
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
if data.empty:
|
||||
logger.error("No data to save")
|
||||
return False
|
||||
|
||||
# Create output directory
|
||||
output_path = Path(output_dir)
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Final duplicate check
|
||||
pre_final_count = len(data)
|
||||
data = data.drop_duplicates(subset=['Query', 'SQL'], keep='first')
|
||||
final_duplicate_count = pre_final_count - len(data)
|
||||
|
||||
if final_duplicate_count > 0:
|
||||
logger.warning(f"Removed {final_duplicate_count} duplicates from final sample")
|
||||
|
||||
# Shuffle the final dataset
|
||||
data = data.sample(frac=1, random_state=random_seed).reset_index(drop=True)
|
||||
|
||||
# Save to CSV
|
||||
output_file_path = output_path / output_filename
|
||||
data.to_csv(output_file_path, index=False)
|
||||
|
||||
logger.info(f"Saved {len(data)} records to {output_file_path}")
|
||||
logger.info("Final distribution:")
|
||||
for level, count in data['Levels'].value_counts().items():
|
||||
logger.info(f" {level}: {count} records")
|
||||
|
||||
return True
|
||||
|
||||
|
||||
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:
|
||||
"""
|
||||
Create a balanced sample dataset from BookSQL train.json.
|
||||
|
||||
This function orchestrates the data loading, sampling, validation, and saving process.
|
||||
|
||||
Args:
|
||||
input_file: Path to the BookSQL train.json file
|
||||
output_dir: Directory to save the output CSV
|
||||
output_filename: Name of the output CSV file
|
||||
samples_per_level: Number of samples per difficulty level (easy, medium, hard)
|
||||
random_seed: Random seed for reproducible sampling
|
||||
validate_queries: If True, validate SQL queries before including them
|
||||
require_data: If True (and validate_queries=True), only include queries that return data
|
||||
|
||||
Returns:
|
||||
bool: True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Step 1: Load and clean data
|
||||
train_df = load_and_clean_data(input_file)
|
||||
|
||||
# Step 2: Sample data for each difficulty level
|
||||
sampled_dfs = []
|
||||
|
||||
if validate_queries:
|
||||
logger.info("Validation enabled - testing SQL queries before including them in sample")
|
||||
if require_data:
|
||||
logger.info("Only including queries that return actual data")
|
||||
|
||||
for level in ['easy', 'medium', 'hard']:
|
||||
if validate_queries:
|
||||
sampled = validate_samples(train_df, level, samples_per_level, random_seed, require_data)
|
||||
else:
|
||||
sampled = sample_by_difficulty(train_df, level, samples_per_level, random_seed)
|
||||
|
||||
if not sampled.empty:
|
||||
sampled_dfs.append(sampled)
|
||||
|
||||
if not sampled_dfs:
|
||||
logger.error("No data could be sampled")
|
||||
return False
|
||||
|
||||
# Step 3: Combine all sampled data
|
||||
final_df = pd.concat(sampled_dfs, ignore_index=True)
|
||||
|
||||
# Step 4: Save results
|
||||
return save_results(final_df, output_dir, output_filename, random_seed)
|
||||
|
||||
except FileNotFoundError:
|
||||
logger.error(f"Input file '{input_file}' not found")
|
||||
logger.error("Tip: Run with --download-data first to download the BookSQL dataset")
|
||||
return False
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Invalid JSON in {input_file}: {e}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing data: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def main():
|
||||
"""Main CLI entry point."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Data utilities for Text-to-SQL evaluation",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Examples:
|
||||
%(prog)s --download-data # Download BookSQL dataset
|
||||
%(prog)s --create-sample # Create sample CSV (15 per level)
|
||||
%(prog)s --create-sample --samples 5 # Create sample with 5 per level
|
||||
%(prog)s --create-sample --validate # Create sample with SQL validation
|
||||
%(prog)s --create-sample --validate --require-data # Only queries that return data
|
||||
"""
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--download-data",
|
||||
action="store_true",
|
||||
help="Download the BookSQL dataset to ./BookSQL-files directory"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--create-sample",
|
||||
action="store_true",
|
||||
help="Create a balanced sample CSV from BookSQL train.json"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--samples",
|
||||
type=int,
|
||||
default=15,
|
||||
help="Number of samples per difficulty level (default: 15)"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--validate",
|
||||
action="store_true",
|
||||
help="Validate SQL queries before including them in the sample"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--require-data",
|
||||
action="store_true",
|
||||
help="Only include queries that return actual data (requires --validate)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.download_data:
|
||||
success = download_booksql_dataset()
|
||||
sys.exit(0 if success else 1)
|
||||
elif args.create_sample:
|
||||
# Validate argument combinations
|
||||
if args.require_data and not args.validate:
|
||||
logger.error("--require-data requires --validate to be enabled")
|
||||
sys.exit(1)
|
||||
|
||||
success = create_sample_dataset(
|
||||
samples_per_level=args.samples,
|
||||
validate_queries=args.validate,
|
||||
require_data=args.require_data
|
||||
)
|
||||
sys.exit(0 if success else 1)
|
||||
else:
|
||||
parser.print_help()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user