vibrantlabsai--ragas/examples/ragas_examples/text2sql/data_utils.py

#!/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()