Module neuralops.data.glue_loader

Local GLUE Loader — loads pre-cached GLUE datasets from disk first, then HF Hub.

Functions

def get_calibration_dataloader(task_name: str = 'sst2',
tokenizer: Any = None,
batch_size: int = 8,
num_samples: int = 128,
seq_len: int = 128,
split: str = 'train') ‑> Any

Expand source code

def get_calibration_dataloader(
    task_name: str = "sst2",
    tokenizer: Any = None,
    batch_size: int = 8,
    num_samples: int = 128,
    seq_len: int = 128,
    split: str = "train",
) -> Any:
    """
    Builds a calibration DataLoader from local disk-cached GLUE data.
    Used by data-aware SVD compression algorithms (ASVD, A3, FRanDI, DynamicRank,
    SwiftSVD, SliceGPT, etc.) to collect activation covariance and Fisher information.
    """
    from torch.utils.data import DataLoader

    t_name = task_name.lower().replace("-", "_")
    ds_dict = load_glue_local(t_name)
    if split not in ds_dict:
        split = list(ds_dict.keys())[0]
    ds = ds_dict[split]

    if num_samples and len(ds) > num_samples:
        ds = ds.select(range(min(num_samples, len(ds))))

    if tokenizer is not None:
        key1, key2 = GLUE_TASK_KEYS.get(t_name, ("sentence", None))
        if key1 not in ds.column_names:
            cols = [c for c in ds.column_names if c not in ("label", "labels", "idx")]
            key1 = cols[0] if cols else "text"
            key2 = None

        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token or "[PAD]"

        def tokenize_fn(examples):
            args = (
                (examples[key1],)
                if key2 is None or key2 not in examples
                else (examples[key1], examples[key2])
            )
            return tokenizer(*args, padding="max_length", truncation=True, max_length=seq_len)

        remove_cols = [c for c in ds.column_names if c not in ("label", "labels")]
        ds = ds.map(tokenize_fn, batched=True, remove_columns=remove_cols)
        keep = [
            c for c in ["input_ids", "attention_mask", "token_type_ids"] if c in ds.column_names
        ]
        ds.set_format(type="torch", columns=keep)

    return DataLoader(ds, batch_size=batch_size, shuffle=False)

Builds a calibration DataLoader from local disk-cached GLUE data. Used by data-aware SVD compression algorithms (ASVD, A3, FRanDI, DynamicRank, SwiftSVD, SliceGPT, etc.) to collect activation covariance and Fisher information.

def get_num_labels_and_is_regression(task_name: str) ‑> tuple[int, bool]

Expand source code

def get_num_labels_and_is_regression(task_name: str) -> tuple[int, bool]:
    """Return (num_labels, is_regression) for GLUE tasks."""
    t = task_name.lower()
    if t == "stsb":
        return 1, True
    mapping = {
        "cola": 2,
        "sst2": 2,
        "mrpc": 2,
        "qqp": 2,
        "mnli": 3,
        "qnli": 2,
        "rte": 2,
        "wnli": 2,
    }
    return mapping.get(t, 2), False

Return (num_labels, is_regression) for GLUE tasks.

def load_glue_local(task_name: str, base_path: str = None) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_glue_local(task_name: str, base_path: str = None) -> DatasetDict:
    """
    Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub.
    Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.
    """
    from datasets import load_from_disk

    candidates = [
        GLUE_LOCAL_DIR / task_name,
        Path.cwd() / "data" / "glue_data" / task_name,
    ]
    for cand in candidates:
        if cand.exists():
            return load_from_disk(str(cand))

    from datasets import load_dataset as hf_load_dataset

    try:
        return hf_load_dataset("glue", task_name)
    except Exception:
        return hf_load_dataset("nyu-mll/glue", task_name)

Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub. Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.

def load_glue_split(task: str,
split: str = 'validation',
tokenizer_name: str = 'bert-base-uncased',
max_seq_len: int = 128) ‑> datasets.arrow_dataset.Dataset | None

Expand source code

def load_glue_split(
    task: str,
    split: str = "validation",
    tokenizer_name: str = "bert-base-uncased",
    max_seq_len: int = 128,
) -> Dataset | None:
    """Load GLUE split from disk cache or HF Hub. Returns torch-format Dataset."""
    task = task.lower()
    if task not in GLUE_TASK_HF_NAME:
        logger.error("Unknown GLUE task: %s", task)
        return None

    key1, key2 = GLUE_TASK_TEXT_KEYS[task]
    hf_path, hf_config = GLUE_TASK_HF_NAME[task]

    dataset = None
    # 1. Try local disk
    local_path = GLUE_LOCAL_DIR / task
    if local_path.exists():
        try:
            from datasets import load_from_disk

            raw = load_from_disk(str(local_path))
            if split in raw:
                dataset = raw[split]
                logger.info("Loaded '%s/%s' from local disk.", task, split)
        except Exception as e:
            logger.warning("Local disk load failed for '%s/%s': %s", task, split, e)

    # 2. Try HF Hub
    if dataset is None:
        try:
            from datasets import load_dataset

            raw = load_dataset(hf_path, hf_config)
            if split in raw:
                dataset = raw[split]
                logger.info("Loaded '%s/%s' from HF Hub.", task, split)
        except Exception as e:
            logger.warning("HF Hub load failed for '%s/%s': %s", task, split, e)
            return None

    if dataset is None:
        return None

    # 3. Tokenize
    try:
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
            tokenizer.pad_token_id = tokenizer.eos_token_id

        def tokenize_fn(examples):
            args = (examples[key1],) if key2 is None else (examples[key1], examples[key2])
            return tokenizer(*args, padding="max_length", truncation=True, max_length=max_seq_len)

        remove_cols = [c for c in dataset.column_names if c not in {"label", "labels"}]
        dataset = dataset.map(tokenize_fn, batched=True, remove_columns=remove_cols)
        if "label" in dataset.column_names and "labels" not in dataset.column_names:
            dataset = dataset.rename_column("label", "labels")
        if "t5" in tokenizer_name.lower() and "token_type_ids" in dataset.column_names:
            dataset = dataset.remove_columns(["token_type_ids"])
        keep = [
            c
            for c in ["input_ids", "attention_mask", "token_type_ids", "labels"]
            if c in dataset.column_names
        ]
        dataset.set_format(type="torch", columns=keep)
        return dataset
    except Exception as e:
        logger.error("Tokenization failed for '%s/%s': %s", task, split, e)
        return None

Load GLUE split from disk cache or HF Hub. Returns torch-format Dataset.

def load_glue_task(task_name: str, base_path: str = None) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_glue_local(task_name: str, base_path: str = None) -> DatasetDict:
    """
    Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub.
    Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.
    """
    from datasets import load_from_disk

    candidates = [
        GLUE_LOCAL_DIR / task_name,
        Path.cwd() / "data" / "glue_data" / task_name,
    ]
    for cand in candidates:
        if cand.exists():
            return load_from_disk(str(cand))

    from datasets import load_dataset as hf_load_dataset

    try:
        return hf_load_dataset("glue", task_name)
    except Exception:
        return hf_load_dataset("nyu-mll/glue", task_name)

Load a GLUE task from local Parquet files, falling back to HuggingFace Datasets hub. Compatible with local, Kaggle, Colab, and remote GPU/HPC environments.