Spaces:
Runtime error
Runtime error
| """Sampling utilities for QualiVec.""" | |
| import pandas as pd | |
| import numpy as np | |
| from typing import Optional, Union, Literal | |
| class Sampler: | |
| """Handles sampling mechanisms for QualiVec.""" | |
| def __init__(self, verbose: bool = True): | |
| """Initialize the Sampler. | |
| Args: | |
| verbose: Whether to print status messages. | |
| """ | |
| self.verbose = verbose | |
| def sample(self, | |
| df: pd.DataFrame, | |
| sampling_type: Literal["random", "stratified"] = "random", | |
| sample_size: Union[int, float] = 0.1, | |
| stratify_column: Optional[str] = None, | |
| seed: Optional[int] = None, | |
| label_column: str = "Label") -> pd.DataFrame: | |
| """Sample data from a DataFrame. | |
| Args: | |
| df: DataFrame to sample from. | |
| sampling_type: Type of sampling ("random" or "stratified"). | |
| sample_size: Size of the sample. If float, interpreted as a fraction. | |
| stratify_column: Column to stratify by (required for stratified sampling). | |
| seed: Random seed for reproducibility. | |
| label_column: Name of the label column to add to the output. | |
| Returns: | |
| DataFrame containing the sampled data. | |
| Raises: | |
| ValueError: If parameters are invalid. | |
| """ | |
| # Set random seed if provided | |
| if seed is not None: | |
| np.random.seed(seed) | |
| # Calculate sample size if given as a fraction | |
| if isinstance(sample_size, float): | |
| if not 0 < sample_size <= 1: | |
| raise ValueError("Sample size as fraction must be between 0 and 1.") | |
| n_samples = int(len(df) * sample_size) | |
| else: | |
| if not 0 < sample_size <= len(df): | |
| raise ValueError(f"Sample size must be between 1 and {len(df)}.") | |
| n_samples = sample_size | |
| if self.verbose: | |
| print(f"Sampling {n_samples} rows ({n_samples/len(df):.1%} of data)...") | |
| # Perform sampling | |
| if sampling_type == "random": | |
| sample = df.sample(n=n_samples, random_state=seed) | |
| elif sampling_type == "stratified": | |
| if stratify_column is None: | |
| raise ValueError("stratify_column must be provided for stratified sampling.") | |
| if stratify_column not in df.columns: | |
| raise ValueError(f"Stratification column '{stratify_column}' not found in DataFrame.") | |
| # Check for NaN values in stratification column | |
| if df[stratify_column].isna().any(): | |
| raise ValueError(f"NaN values found in stratification column '{stratify_column}'.") | |
| # Calculate the proportion for each stratum | |
| strata = df[stratify_column].value_counts(normalize=True) | |
| # Create empty sample DataFrame | |
| sample = pd.DataFrame(columns=df.columns) | |
| # Sample from each stratum | |
| for stratum, proportion in strata.items(): | |
| stratum_df = df[df[stratify_column] == stratum] | |
| stratum_samples = max(1, int(n_samples * proportion)) | |
| stratum_sample = stratum_df.sample(n=min(stratum_samples, len(stratum_df)), | |
| random_state=seed) | |
| sample = pd.concat([sample, stratum_sample]) | |
| if self.verbose: | |
| print(f"Stratified sampling based on '{stratify_column}':") | |
| for stratum, count in sample[stratify_column].value_counts().items(): | |
| print(f" - {stratum}: {count} samples ({count/n_samples:.1%})") | |
| else: | |
| raise ValueError(f"Unknown sampling type: {sampling_type}") | |
| # Add empty label column for manual annotation | |
| if label_column not in sample.columns: | |
| sample[label_column] = None | |
| if self.verbose: | |
| print(f"Created sample with {len(sample)} rows.") | |
| return sample | |