NeuralOps Core ML Framework & Efficient Transformer Toolkit.

A unified, modular, and high-performance toolkit for custom attention mechanisms, model compression (SVD, quantization, pruning, distillation), distributed training, hyperparameter optimization, federated learning, data loading, benchmarking, and interactive visualization.

Sub-modules

neuralops.analysis

Analysis tools: plotting, metrics, layerwise analysis, model comparison, reporting.

neuralops.attentions

Auto-import all attention modules and export their registered classes.

neuralops.benchmarks

NeuralOps Benchmarks & Sweepers Suite …

neuralops.compression_engine

Auto-import all compression_engine modules and export their registered classes.

neuralops.config

Unified configuration system for the Efficient Transformer Toolkit …

neuralops.core

Efficient Transformer Toolkit – Core module. Provides model‑agnostic architecture detection, attention injection, compression pipeline, and base classes.

neuralops.data

Model‑agnostic data pipeline – dataset loading, tokenization, preprocessing, DataLoader creation, and augmentation for all common tasks and models.

neuralops.evaluation

Evaluation module for model performance, fairness, robustness, label quality, and A/B testing.

neuralops.experiment_logging

Experiment logging – structured logging, metrics tracking, and Trainer callbacks.

neuralops.experiment_utils

Django integration layer – background experiment runner, WebSocket broadcaster, AI assistant, and multi‑run comparison.

neuralops.main

Efficient Attention Research Framework — Full Pipeline Runner

Single entry-point for: - Training (fine-tune any HF model on GLUE / custom …

neuralops.pipeline

Pipeline package — result writing, run management, and sweep orchestration.

neuralops.results_saver

Results persistence, export, cloud upload, and experiment tracking. Orchestrated by ResultsSaver with a SaveConfig.

neuralops.runners

NeuralOps Runners — pipeline orchestration, continuous sweeps, and monitoring …

neuralops.tests

neuralops.tracking

Experiment Tracking – pluggable backends (local, TensorBoard, WandB, MLflow) with a unified API and HuggingFace Trainer callback.

neuralops.trainer

Training module – unified trainers, results management, metrics, and callbacks.

neuralops.utils

Functions

def aggregate_results(results_list) ‑> Any

Expand source code

def aggregate_results(results_list) -> Any:
    """Aggregate results."""
    if not results_list:
        return {}
    aggregated = {}
    for key in results_list[0].keys():
        values = [r[key] for r in results_list if key in r]
        if values:
            aggregated[f"{key}_mean"] = np.mean(values)
            aggregated[f"{key}_std"] = np.std(values)
            aggregated[f"{key}_values"] = values
    return aggregated

Aggregate results.

def compress_model(model: torch.nn.modules.module.Module,
config: dict[str, typing.Any],
tokenizer=None) ‑> torch.nn.modules.module.Module

Expand source code

def compress_model(
    model: nn.Module,
    config: dict[str, Any],
    tokenizer=None,
) -> nn.Module:
    """Convenience function to compress a model with a configuration dict."""
    pipeline = CompressionPipeline(model, config, tokenizer)
    return pipeline.run()

Convenience function to compress a model with a configuration dict.

def create_dataloader(dataset,
tokenizer,
config: DataConfig,
**kwargs) ‑> torch.utils.data.dataloader.DataLoader

Expand source code

def create_dataloader(dataset, tokenizer, config: DataConfig, **kwargs) -> DataLoader:
    """
    Create a DataLoader from a HuggingFace Dataset, applying preprocessing if not already tokenized.
    Uses ProductionDataCollator to guarantee rock-solid batch dimensions and GPU Tensor Core alignment.
    """
    if "input_ids" in dataset.column_names and "attention_mask" in dataset.column_names:
        processed = dataset
    else:
        processed = dataset.map(
            lambda examples: preprocess_function(examples, tokenizer, config),
            batched=True,
            remove_columns=dataset.column_names,
        )
    columns = [
        col
        for col in [
            "input_ids",
            "attention_mask",
            "labels",
            "start_positions",
            "end_positions",
            "token_type_ids",
        ]
        if col in processed.column_names
    ]
    processed.set_format(type="torch", columns=columns)

    collate_fn = kwargs.pop("collate_fn", None)
    if collate_fn is None:
        collate_fn = ProductionDataCollator(
            pad_token_id=(getattr(tokenizer, "pad_token_id", 0) if tokenizer is not None else 0),
            pad_to_multiple_of=8,
            max_length=config.max_length,
        )

    return DataLoader(
        processed,
        batch_size=config.batch_size,
        shuffle=config.shuffle,
        num_workers=config.num_workers,
        collate_fn=collate_fn,
        **kwargs,
    )

Create a DataLoader from a HuggingFace Dataset, applying preprocessing if not already tokenized. Uses ProductionDataCollator to guarantee rock-solid batch dimensions and GPU Tensor Core alignment.

def detect_architecture(model: torch.nn.modules.module.Module) ‑> str

Expand source code

def detect_architecture(model: nn.Module) -> str:
    """Return canonical architecture name."""
    config = getattr(model, "config", None)
    model_type = getattr(config, "model_type", "") if config else ""
    name = model.__class__.__name__.lower()
    if "distilbert" in model_type or "distilbert" in name:
        return "distilbert"
    if "roberta" in model_type or "roberta" in name:
        return "roberta"
    if "bert" in model_type or "bert" in name:
        return "bert"
    if "gpt2" in model_type or "gpt2" in name:
        return "gpt2"
    if "qwen" in model_type or "qwen" in name:
        return "qwen"
    if "llama" in model_type:
        return "llama"
    if "mistral" in model_type:
        return "mistral"
    if "gemma" in model_type:
        return "gemma"
    if "phi" in model_type:
        return "phi"
    if "t5" in model_type:
        return "t5"
    return "unknown"

Return canonical architecture name.

def fms_scan_triton(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, decay: torch.Tensor) ‑> torch.Tensor

Expand source code

def fms_scan_triton(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, decay: torch.Tensor
) -> torch.Tensor:
    """Triton-accelerated FMS scan function. head_dim must be <= 64."""
    B, H, L, D = q.shape
    if D > MAX_REGISTER_DIM:
        raise ValueError(f"FMS Triton head_dim {D} > {MAX_REGISTER_DIM}")

    # Ensure decay is (H,)
    if decay.dim() > 1:
        decay = decay.squeeze()
    elif decay.dim() == 0:
        decay = decay.repeat(H)

    if decay.shape[0] != H:
        raise ValueError(f"decay shape {decay.shape} does not match H={H}")

    return _FMSScanFn.apply(q, k, v, decay)

Triton-accelerated FMS scan function. head_dim must be <= 64.

def generate_synthetic_data(output_path: str,
num_samples: int = 10000,
domain: str = 'movies',
seed: int = 42) ‑> pandas.DataFrame

Expand source code

def generate_synthetic_data(
    output_path: str, num_samples: int = 10000, domain: str = "movies", seed: int = 42
) -> pd.DataFrame:
    """Generate a single-domain synthetic classification dataset."""
    random.seed(seed)
    half = num_samples // 2
    pos = _generate_batch(domain, half, label=1, is_positive=True)
    neg = _generate_batch(domain, half, label=0, is_positive=False)
    df = pd.DataFrame(pos + neg).sample(frac=1, random_state=seed).reset_index(drop=True)
    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_path, index=False)
    logger.info("Generated %d samples to %s", len(df), output_path)
    return df

Generate a single-domain synthetic classification dataset.

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_task_metrics(task_name: str) ‑> Any

Expand source code

def get_task_metrics(task_name: str) -> Any:
    """Return a compute_metrics function for a GLUE task."""
    task_name = task_name.lower()
    if task_name == "stsb":

        def fn(p) -> Any:
            """Fn."""
            preds = np.squeeze(p.predictions)
            labels = p.label_ids
            return compute_pearson_spearman(preds, labels)

    elif task_name == "cola":

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            return {
                "matthews_correlation": matthews_corrcoef(labels, preds),
                "accuracy": accuracy_score(labels, preds),
            }

    elif task_name in ("mrpc", "qqp"):

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            acc = accuracy_score(labels, preds)
            f1 = f1_score(labels, preds)
            return {"accuracy": acc, "f1": f1, "accuracy_f1": (acc + f1) / 2}

    else:

        def fn(p) -> Any:
            """Fn."""
            preds = np.argmax(p.predictions, axis=1)
            labels = p.label_ids
            return {
                "accuracy": accuracy_score(labels, preds),
                "f1_macro": f1_score(labels, preds, average="macro", zero_division=0),
                "f1_weighted": f1_score(labels, preds, average="weighted", zero_division=0),
            }

    return fn

Return a compute_metrics function for a GLUE task.

def gla_scan_triton(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, lambdas: torch.Tensor) ‑> torch.Tensor

Expand source code

def gla_scan_triton(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, lambdas: torch.Tensor
) -> torch.Tensor:
    """Triton-accelerated GLA scan function. head_dim must be <= 64."""
    B, H, L, D = q.shape
    if D > MAX_REGISTER_DIM:
        raise ValueError(f"GLA Triton head_dim {D} > {MAX_REGISTER_DIM}")
    return _GLAScanFn.apply(q, k, v, lambdas)

Triton-accelerated GLA scan function. head_dim must be <= 64.

def inject_attention(model: nn.Module,
attention_type: str,
rank: int | None = None,
copy_weights: bool = True,
**kwargs: Any) ‑> torch.nn.modules.module.Module

Expand source code

def inject_attention(
    model: nn.Module,
    attention_type: str,
    rank: int | None = None,
    copy_weights: bool = True,
    **kwargs: Any,
) -> nn.Module:
    """Universal attention layer replacement engine for HuggingFace Transformers."""
    attn_class = ATTENTION_REGISTRY.resolve(attention_type)
    if callable(attn_class) and not isinstance(attn_class, type):
        try:
            res = attn_class(model=model, rank=rank, copy_weights=copy_weights, **kwargs)
            if res is not None:
                return res
        except Exception as inj_err:
            logger.debug(
                "Dedicated injector %s failed, falling back to layer replacement: %s",
                attention_type,
                inj_err,
            )

    config = getattr(model, "config", None)
    if config is not None:
        cfg_source = getattr(config, "text_config", config)
        hidden_size = getattr(
            cfg_source,
            "hidden_size",
            getattr(
                cfg_source,
                "d_model",
                getattr(cfg_source, "n_embd", getattr(cfg_source, "dim", None)),
            ),
        )
        num_heads = getattr(
            cfg_source,
            "num_attention_heads",
            getattr(
                cfg_source,
                "num_heads",
                getattr(cfg_source, "n_heads", getattr(cfg_source, "n_head", None)),
            ),
        )
        head_dim = getattr(
            cfg_source,
            "head_dim",
            getattr(
                cfg_source,
                "d_kv",
                (hidden_size // max(1, num_heads) if (hidden_size and num_heads) else None),
            ),
        )
        num_kv_heads = getattr(cfg_source, "num_key_value_heads", num_heads)
        if hidden_size is not None:
            config.hidden_size = hidden_size
        if num_heads is not None:
            config.num_attention_heads = num_heads
        if head_dim is not None:
            config.head_dim = head_dim
        if num_kv_heads is not None:
            config.num_key_value_heads = num_kv_heads
        if not hasattr(config, "dropout"):
            dropout_val = getattr(
                config,
                "hidden_dropout_prob",
                getattr(config, "attention_probs_dropout_prob", 0.1),
            )
            config.dropout = dropout_val
    else:
        hidden_size = getattr(model, "hidden_size", None)
        num_heads = getattr(model, "num_attention_heads", None)
        head_dim = getattr(model, "head_dim", None)
        num_kv_heads = getattr(model, "num_key_value_heads", num_heads)

    arch = detect_architecture(model)
    adapter = get_adapter(arch)
    family = detect_family(model)
    family_adapter = ModelAdapter(model, family)

    is_dec = family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral")
    if config is not None:
        config.is_decoder = bool(is_dec or getattr(config, "is_decoder", False))
        if not hasattr(config, "model_type") or not config.model_type:
            config.model_type = family

    replaced = 0
    modules_list = list(get_attention_modules(model))
    if not modules_list:
        logger.warning("No attention modules found to replace for architecture '%s'", arch)

    for name, old_attn in modules_list:
        # In hybrid architectures (like Qwen3.5), only replace linear_attn if candidate explicitly supports DeltaNet
        if "linear_attn" in name and not getattr(attn_class, "supports_delta_net", False):
            continue

        # Dynamically infer dimensions from family_adapter and old_attn
        if family_adapter is not None:
            if hidden_size is None:
                hidden_size = getattr(family_adapter, "hidden_size", None)
            if num_heads is None:
                num_heads = getattr(family_adapter, "num_heads", None)
            if num_kv_heads is None:
                num_kv_heads = getattr(family_adapter, "num_key_value_heads", None)
            if head_dim is None:
                head_dim = getattr(family_adapter, "head_dim", None)

        if hidden_size is None or num_heads is None:
            for _, child in old_attn.named_children():
                if (
                    hasattr(child, "weight")
                    and isinstance(child.weight, torch.Tensor)
                    and child.weight.ndim == 2
                ):
                    if hidden_size is None:
                        hidden_size = child.weight.shape[1]
                    break
            if num_heads is None:
                num_heads = getattr(
                    old_attn,
                    "num_heads",
                    getattr(old_attn, "num_attention_heads", None),
                )
            if num_kv_heads is None:
                num_kv_heads = getattr(old_attn, "num_key_value_heads", num_heads)
            if head_dim is None and hidden_size is not None and num_heads is not None:
                head_dim = getattr(old_attn, "head_dim", hidden_size // max(1, num_heads))

        # Ensure config exists or construct minimal namespace config
        if config is None:
            from types import SimpleNamespace

            is_decoder_family_bool = bool(
                is_dec or family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral")
            )
            config = SimpleNamespace(
                hidden_size=hidden_size,
                num_attention_heads=num_heads,
                head_dim=head_dim,
                num_key_value_heads=num_kv_heads or num_heads,
                dropout=0.1,
                is_decoder=is_decoder_family_bool,
                is_causal=is_decoder_family_bool,
                model_type=family,
            )
        else:
            if family in ("gpt", "gpt2", "qwen", "qwen3_5", "llama", "mistral") or is_dec:
                config.is_decoder = True
                config.is_causal = True

        # Ensure old_attn has canonical projection attributes for custom modules inspecting pretrained_attn
        for proj_key, alias_list in [
            ("q", ["query", "q_proj", "q_lin", "Wq", "wq", "q"]),
            ("k", ["key", "k_proj", "k_lin", "Wk", "wk", "k"]),
            ("v", ["value", "v_proj", "v_lin", "Wv", "wv", "v"]),
            ("o", ["dense", "out_proj", "o_proj", "out_lin", "c_proj", "wo", "o"]),
        ]:
            try:
                existing_mod = next(
                    (getattr(old_attn, a) for a in alias_list if hasattr(old_attn, a)),
                    None,
                )
                if existing_mod is None and hasattr(old_attn, "self"):
                    existing_mod = next(
                        (
                            getattr(old_attn.self, a)
                            for a in alias_list
                            if hasattr(old_attn.self, a)
                        ),
                        None,
                    )

                if existing_mod is not None:
                    for a in alias_list:
                        if not hasattr(old_attn, a):
                            try:
                                setattr(old_attn, a, existing_mod)
                            except Exception:
                                pass
                else:
                    w, b = adapter.get_weight_bias(old_attn, proj_key)
                    if w is not None:
                        proxy = nn.Linear(
                            w.shape[1],
                            w.shape[0],
                            bias=(b is not None),
                            device=w.device,
                            dtype=w.dtype,
                        )
                        proxy.weight.data.copy_(w)
                        if b is not None:
                            proxy.bias.data.copy_(b)
                        for a in alias_list:
                            if not hasattr(old_attn, a):
                                try:
                                    setattr(old_attn, a, proxy)
                                except Exception:
                                    pass
            except Exception as ex:
                logger.debug(
                    "Failed to set projection aliases for %s on %s: %s",
                    proj_key,
                    name,
                    ex,
                )

        effective_rank = rank if rank is not None else getattr(config, "rank", 32)

        # Standardized instantiation: try passing pretrained_attn=old_attn first so modules
        # can inherit pretrained weights during initialization, then fallback without it
        attempts = [
            lambda: attn_class(config, rank=effective_rank, pretrained_attn=old_attn, **kwargs),
            lambda: attn_class(config, pretrained_attn=old_attn, **kwargs),
            lambda: attn_class(config=config, pretrained_attn=old_attn, **kwargs),
            lambda: attn_class(config, rank=effective_rank, **kwargs),
            lambda: attn_class(config, **kwargs),
            lambda: attn_class(rank=effective_rank, **kwargs),
            lambda: attn_class(**kwargs),
        ]
        new_attn = None
        for attempt in attempts:
            try:
                new_attn = attempt()
                break
            except TypeError:
                continue
            except Exception as e:
                logger.debug("Instantiation attempt failed: %s", e)
                continue

        if new_attn is None:
            raise RuntimeError(f"Could not instantiate attention class {attn_class.__name__}")

        # Do not let a legacy full-MHA replacement silently retain randomly
        # initialized K/V projections on GQA decoders.  Such a module must
        # explicitly advertise the geometry it implements.
        is_gqa = (
            family in ("qwen", "qwen3_5")
            and family_adapter.num_key_value_heads != family_adapter.num_heads
        )
        if is_gqa and not getattr(new_attn, "supports_gqa", False):
            expected = {
                "q": family_adapter.q_proj_out,
                "k": family_adapter.kv_proj_out,
                "v": family_adapter.kv_proj_out,
            }
            for key, patterns in (
                ("q", PROJ_REGEX_PATTERNS["q"]),
                ("k", PROJ_REGEX_PATTERNS["k"]),
                ("v", PROJ_REGEX_PATTERNS["v"]),
            ):
                target = _find_submodule_by_regex(new_attn, patterns)
                out_features = getattr(target, "out_features", None)
                if target is not None and out_features not in {expected[key], None}:
                    if expected[key] is not None and out_features is not None:
                        if (
                            out_features % expected[key] == 0
                            or expected[key] % out_features == 0
                            or _find_submodule_by_regex(new_attn, PROJ_REGEX_PATTERNS["o"])
                            is not None
                        ):
                            pass
                        else:
                            raise QwenIncompatibleError(
                                f"{attention_type}: {key}_projection_out={out_features}, "
                                f"requires={expected[key]} (GQA)"
                            )

        # Match device and dtype of old attention layer
        old_param = next(old_attn.parameters(), None)
        if old_param is not None:
            new_attn.to(device=old_param.device, dtype=old_param.dtype)

        if copy_weights:
            custom_copied = False
            if (
                is_dec
                or hasattr(old_attn, "c_attn")
                or "gpt" in old_attn.__class__.__name__.lower()
            ):
                old_attn.is_decoder = True
                if hasattr(new_attn, "is_decoder"):
                    new_attn.is_decoder = True
            if hasattr(new_attn, "copy_weights") and callable(new_attn.copy_weights):
                try:
                    new_attn.copy_weights(old_attn)
                    custom_copied = True
                    if is_dec and hasattr(new_attn, "is_decoder"):
                        new_attn.is_decoder = True
                    if old_param is not None:
                        new_attn.to(device=old_param.device, dtype=old_param.dtype)
                except (AttributeError, RuntimeError, ValueError) as e:
                    logger.debug("Custom copy_weights failed on %s: %s", name, e)
                except Exception as e:
                    logger.warning("Unexpected error in custom copy_weights on %s: %s", name, e)

                if not custom_copied:
                    # Detect fused QKV projection modules via regex and slice parameters
                    fused_mod = None
                    for child_name, child_mod in old_attn.named_children():
                        if FUSED_QKV_REGEX.match(child_name) and hasattr(child_mod, "weight"):
                            fused_mod = child_mod
                            break
                    if fused_mod is None:
                        for candidate in [
                            "c_attn",
                            "qkv_proj",
                            "query_key_value",
                            "in_proj",
                            "Wqkv",
                        ]:
                            mod = getattr(old_attn, candidate, None)
                            if mod is not None and hasattr(mod, "weight"):
                                fused_mod = mod
                                break

                    if fused_mod is not None:
                        # Sliced Fused QKV handling
                        sliced = _slice_fused_qkv_weights(
                            fused_mod.weight.data,
                            (
                                fused_mod.bias.data
                                if hasattr(fused_mod, "bias") and fused_mod.bias is not None
                                else None
                            ),
                            num_heads=num_heads if num_heads is not None else 12,
                            num_kv_heads=num_kv_heads or num_heads,
                            head_dim=head_dim,
                            source_mod=fused_mod,
                        )
                        for proj_key in ["q", "k", "v"]:
                            w, b = sliced[proj_key]
                            target_proj = _find_submodule_by_regex(
                                new_attn, PROJ_REGEX_PATTERNS[proj_key]
                            )
                            if target_proj is not None:
                                _copy_projection_weights(target_proj, w, b)
                            else:
                                lin1 = getattr(new_attn, f"{proj_key}_proj1", None)
                                lin2 = getattr(new_attn, f"{proj_key}_proj2", None)
                                if lin1 is not None and lin2 is not None:
                                    _copy_projection_weights(lin1, w, None)
                                    with torch.no_grad():
                                        W_f = w.detach().float()
                                        if hasattr(lin1, "weight") and hasattr(lin2, "weight"):
                                            if (
                                                W_f.shape[1] != lin1.weight.shape[1]
                                                and W_f.shape[0] == lin1.weight.shape[1]
                                            ):
                                                W_f = W_f.t()
                                            U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
                                            r = min(
                                                lin1.weight.shape[0],
                                                lin2.weight.shape[1],
                                                U.size(1),
                                                Vt.size(0),
                                            )
                                            sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                                            U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to(
                                                device=w.device, dtype=w.dtype
                                            )
                                            V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(
                                                device=w.device, dtype=w.dtype
                                            )
                                            lin1.weight.data[:r, :].copy_(V_b)
                                            lin2.weight.data[:, :r].copy_(U_b)
                                            if (
                                                b is not None
                                                and getattr(lin2, "bias", None) is not None
                                            ):
                                                lin2.bias.data.copy_(
                                                    b.to(
                                                        lin2.bias.device,
                                                        lin2.bias.dtype,
                                                    )
                                                )

                        # Output projection
                        try:
                            src_w_o, src_b_o = adapter.get_weight_bias(old_attn, "o")
                            target_o = _find_submodule_by_regex(new_attn, PROJ_REGEX_PATTERNS["o"])
                            if target_o is not None:
                                _copy_projection_weights(target_o, src_w_o, src_b_o)
                            else:
                                lin1_o = getattr(
                                    new_attn,
                                    "out_proj1",
                                    getattr(new_attn, "o_proj1", None),
                                )
                                lin2_o = getattr(
                                    new_attn,
                                    "out_proj2",
                                    getattr(new_attn, "o_proj2", None),
                                )
                                if lin1_o is not None and lin2_o is not None:
                                    with torch.no_grad():
                                        W_f = src_w_o.detach().float()
                                        if (
                                            W_f.shape[1] != lin1_o.weight.shape[1]
                                            and W_f.shape[0] == lin1_o.weight.shape[1]
                                        ):
                                            W_f = W_f.t()
                                        U, S, Vt = torch.linalg.svd(W_f, full_matrices=False)
                                        r = min(
                                            lin1_o.weight.shape[0],
                                            lin2_o.weight.shape[1],
                                            U.size(1),
                                            Vt.size(0),
                                        )
                                        sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                                        U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to(
                                            device=src_w_o.device, dtype=src_w_o.dtype
                                        )
                                        V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(
                                            device=src_w_o.device, dtype=src_w_o.dtype
                                        )
                                        min_out = min(lin2_o.weight.shape[0], U_b.shape[0])
                                        min_in = min(lin1_o.weight.shape[1], V_b.shape[1])
                                        lin1_o.weight.data[:r, :min_in].copy_(V_b[:, :min_in])
                                        lin2_o.weight.data[:min_out, :r].copy_(U_b[:min_out, :])
                                        if (
                                            src_b_o is not None
                                            and getattr(lin2_o, "bias", None) is not None
                                        ):
                                            min_b = min(lin2_o.bias.shape[0], src_b_o.shape[0])
                                            lin2_o.bias.data[:min_b].copy_(
                                                src_b_o[:min_b].to(
                                                    lin2_o.bias.device,
                                                    lin2_o.bias.dtype,
                                                )
                                            )
                        except (ValueError, AttributeError) as e:
                            logger.debug("Could not copy output projection for %s: %s", name, e)
                    else:
                        # Standard decoupled Q, K, V, O projections
                        for proj_key in ["q", "k", "v", "o"]:
                            src_weight, src_bias = None, None
                            try:
                                src_weight, src_bias = adapter.get_weight_bias(old_attn, proj_key)
                            except (ValueError, AttributeError) as e:
                                logger.debug(
                                    "Could not get %s projection weights for %s: %s",
                                    proj_key,
                                    name,
                                    e,
                                )

                            target_proj = _find_submodule_by_regex(
                                new_attn, PROJ_REGEX_PATTERNS[proj_key]
                            )
                            if target_proj is not None:
                                if src_weight is not None:
                                    _copy_projection_weights(target_proj, src_weight, src_bias)
                                elif proj_key == "o":
                                    # Encoder self-attention (BERT/RoBERTa .self) has no internal out_proj
                                    # because the outer attention container already applies output.dense.
                                    # Initialize new_attn.out_proj to Identity so context passes through transparently.
                                    with torch.no_grad():
                                        if (
                                            hasattr(target_proj, "weight")
                                            and target_proj.weight.ndim == 2
                                        ):
                                            if (
                                                target_proj.weight.shape[0]
                                                == target_proj.weight.shape[1]
                                            ):
                                                nn.init.eye_(target_proj.weight)
                                            else:
                                                nn.init.zeros_(target_proj.weight)
                                                min_d = min(
                                                    target_proj.weight.shape[0],
                                                    target_proj.weight.shape[1],
                                                )
                                                target_proj.weight.data[
                                                    :min_d, :min_d
                                                ].fill_diagonal_(1.0)
                                        elif hasattr(target_proj, "set_svd_factors") or (
                                            hasattr(target_proj, "U") and hasattr(target_proj, "V")
                                        ):
                                            # Low-rank / factorized linear layer: set factors to truncated identity
                                            d_out = getattr(
                                                target_proj,
                                                "out_features",
                                                getattr(
                                                    target_proj,
                                                    "d_out",
                                                    hidden_size or 768,
                                                ),
                                            )
                                            d_in = getattr(
                                                target_proj,
                                                "in_features",
                                                getattr(
                                                    target_proj,
                                                    "d_in",
                                                    hidden_size or 768,
                                                ),
                                            )
                                            eye_mat = torch.eye(
                                                d_out,
                                                d_in,
                                                device=(
                                                    target_proj.U.device
                                                    if hasattr(target_proj, "U")
                                                    else None
                                                ),
                                            )
                                            _copy_projection_weights(target_proj, eye_mat, None)
                                        if (
                                            hasattr(target_proj, "bias")
                                            and target_proj.bias is not None
                                        ):
                                            nn.init.zeros_(target_proj.bias)
                            else:
                                # Support indexed dual projections (e.g. aegis_hybrid, svd_linear_attention: q_proj1, q_proj2)
                                lin1 = getattr(
                                    new_attn,
                                    f"{proj_key}_proj1",
                                    getattr(new_attn, f"{proj_key}1", None),
                                )
                                lin2 = getattr(
                                    new_attn,
                                    f"{proj_key}_proj2",
                                    getattr(new_attn, f"{proj_key}2", None),
                                )
                                if proj_key == "o":
                                    lin1 = lin1 or getattr(
                                        new_attn,
                                        "out_proj1",
                                        getattr(new_attn, "o_proj1", None),
                                    )
                                    lin2 = lin2 or getattr(
                                        new_attn,
                                        "out_proj2",
                                        getattr(new_attn, "o_proj2", None),
                                    )
                                if lin1 is not None and lin2 is not None:
                                    if src_weight is not None:
                                        with torch.no_grad():
                                            W_f = src_weight.detach().float()
                                            if hasattr(lin1, "weight") and hasattr(lin2, "weight"):
                                                if (
                                                    W_f.shape[1] != lin1.weight.shape[1]
                                                    and W_f.shape[0] == lin1.weight.shape[1]
                                                ):
                                                    W_f = W_f.t()
                                                U, S, Vt = torch.linalg.svd(
                                                    W_f, full_matrices=False
                                                )
                                                r = min(
                                                    lin1.weight.shape[0],
                                                    lin2.weight.shape[1],
                                                    U.size(1),
                                                    Vt.size(0),
                                                )
                                                sqrt_S = torch.sqrt(torch.clamp(S[:r], min=1e-12))
                                                U_b = (U[:, :r] * sqrt_S.unsqueeze(0)).to(
                                                    device=src_weight.device,
                                                    dtype=src_weight.dtype,
                                                )
                                                V_b = (sqrt_S.unsqueeze(1) * Vt[:r, :]).to(
                                                    device=src_weight.device,
                                                    dtype=src_weight.dtype,
                                                )
                                                min_out = min(lin2.weight.shape[0], U_b.shape[0])
                                                min_in = min(lin1.weight.shape[1], V_b.shape[1])
                                                lin1.weight.data[:r, :min_in].copy_(V_b[:, :min_in])
                                                lin2.weight.data[:min_out, :r].copy_(
                                                    U_b[:min_out, :]
                                                )
                                                if (
                                                    src_bias is not None
                                                    and getattr(lin2, "bias", None) is not None
                                                ):
                                                    min_b = min(
                                                        lin2.bias.shape[0],
                                                        src_bias.shape[0],
                                                    )
                                                    lin2.bias.data[:min_b].copy_(
                                                        src_bias[:min_b].to(
                                                            lin2.bias.device,
                                                            lin2.bias.dtype,
                                                        )
                                                    )
                                    elif proj_key == "o":
                                        # Identity for dual out_proj
                                        with torch.no_grad():
                                            if hasattr(lin1, "weight") and hasattr(lin2, "weight"):
                                                nn.init.zeros_(lin1.weight)
                                                nn.init.zeros_(lin2.weight)
                                                r = min(
                                                    lin1.weight.shape[0],
                                                    lin2.weight.shape[1],
                                                    lin1.weight.shape[1],
                                                    lin2.weight.shape[0],
                                                )
                                                lin1.weight.data[:r, :r].fill_diagonal_(1.0)
                                                lin2.weight.data[:r, :r].fill_diagonal_(1.0)
                                                if getattr(lin2, "bias", None) is not None:
                                                    nn.init.zeros_(lin2.bias)

        from neuralops.utils.replace_attention import UniversalAttentionWrapper

        return_tuple = (
            "self_attn" in name.lower()
            or "attention" in name.lower()
            or getattr(old_attn, "is_decoder", False)
            or family == "gpt"
            or name.endswith(".attn")
        )
        wrapped_attn = UniversalAttentionWrapper(
            new_attn, layer_name=name, return_tuple=return_tuple
        )
        replace_module(model, name, wrapped_attn)
        replaced += 1

    logger.info(
        "Successfully injected '%s' into %d transformer blocks.",
        attention_type,
        replaced,
    )
    if not hasattr(model, "compression_stats"):
        model.compression_stats = {}
    model.compression_stats["num_compressed"] = replaced
    model.compression_stats["total_params_after"] = sum(p.numel() for p in model.parameters())
    return model

Universal attention layer replacement engine for HuggingFace Transformers.

def load_benchmark_dataset(dataset_name: str,
tokenizer: Any = None,
seq_len: int = 2048,
batch_size: int = 8,
split: str | None = None,
max_samples: int = 1000,
streaming: bool = True,
cache_dir: str | None = None,
shuffle: bool = False,
num_workers: int = 0,
use_synthetic: bool = False,
seed: int = 42,
return_dataloader: bool = True,
**kwargs: Any) ‑> torch.utils.data.dataloader.DataLoader | list[dict[str, typing.Any]]

Expand source code

def load_benchmark_dataset(
    dataset_name: str,
    tokenizer: Any = None,
    seq_len: int = 2048,
    batch_size: int = 8,
    split: str | None = None,
    max_samples: int = 1000,
    streaming: bool = True,
    cache_dir: str | None = None,
    shuffle: bool = False,
    num_workers: int = 0,
    use_synthetic: bool = False,
    seed: int = 42,
    return_dataloader: bool = True,
    **kwargs: Any,
) -> DataLoader | list[dict[str, Any]]:
    """
    Unified loader for 12+ NLP / LLM benchmark datasets.

    Supports:
      - wikitext2 / wikitext103
      - c4 / c4_en
      - ptb
      - pg19
      - gsm8k
      - humaneval
      - mmlu
      - pile / pile_val
      - arc (ARC-Easy & ARC-Challenge)
      - hellaswag
      - lambada
      - glue (sst2, mrpc, qqp, mnli, qnli, rte, wnli, cola, stsb)

    Args:
        dataset_name: Benchmark name or alias (e.g. 'wikitext2', 'c4', 'gsm8k', 'arc', 'glue/sst2').
        tokenizer: HuggingFace tokenizer instance (optional).
        seq_len: Maximum sequence length for tokenization.
        batch_size: DataLoader batch size.
        split: Dataset split ('train', 'validation', 'test'). Defaults to benchmark standard.
        max_samples: Maximum number of samples to load/stream.
        streaming: Whether to stream from HuggingFace Hub to avoid large downloads.
        cache_dir: Local directory to cache datasets.
        shuffle: Whether to shuffle samples.
        num_workers: Number of workers for DataLoader.
        use_synthetic: Force synthetic data generation for offline testing.
        seed: Random seed for reproducibility.
        return_dataloader: If True, returns PyTorch DataLoader; otherwise returns list of dicts.
        **kwargs: Additional arguments passed to `datasets.load_dataset`.

    Returns:
        torch.utils.data.DataLoader or List[Dict[str, Any]].
    """
    clean_name = dataset_name.lower().replace("-", "_")
    info = BenchmarkDatasetRegistry.get(clean_name)

    # Fallback to general lookup if not matched directly
    if info is None and "/" in dataset_name:
        prefix, subtask = dataset_name.split("/", 1)
        if prefix.lower() == "glue":
            info = BenchmarkDatasetRegistry.get(f"glue_{subtask}")

    if info is None:
        logger.warning(
            "Dataset '%s' not in official benchmark registry. Creating generic fallback.",
            dataset_name,
        )
        info = BenchmarkDatasetInfo(
            name=clean_name,
            aliases=[],
            category="custom",
            default_hf_path=dataset_name,
            default_split=split or "validation",
            text_formatter=lambda r: {"text": r.get("text", str(r))},
            synthetic_generator=lambda n, s: [
                {
                    "text": f"Synthetic custom text sample #{i} for {dataset_name}.",
                    "id": i,
                }
                for i in range(n)
            ],
        )

    target_split = split or info.default_split
    raw_samples: list[dict[str, Any]] = []

    # Attempt online loading if not forced synthetic
    if not use_synthetic:
        try:
            from datasets import load_dataset as hf_load_dataset

            hf_kwargs = dict(kwargs)
            if info.default_subset:
                args = (info.default_hf_path, info.default_subset)
            else:
                args = (info.default_hf_path,)

            if cache_dir:
                hf_kwargs["cache_dir"] = cache_dir

            logger.info("Attempting to load benchmark '%s' from HF (%s)...", info.name, args)

            try:
                # First try streaming if requested
                if streaming:
                    hf_ds = hf_load_dataset(*args, split=target_split, streaming=True, **hf_kwargs)
                    count = 0
                    for row in hf_ds:
                        formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                        # Filter out empty text rows (e.g. WikiText empty headers)
                        if (
                            formatted.get("text", "").strip()
                            or "prompt" in formatted
                            or "question" in formatted
                        ):
                            raw_samples.append(formatted)
                            count += 1
                            if count >= max_samples:
                                break
                else:
                    hf_ds = hf_load_dataset(*args, split=target_split, **hf_kwargs)
                    sub_ds = hf_ds.select(range(min(len(hf_ds), max_samples)))
                    for row in sub_ds:
                        formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                        raw_samples.append(formatted)

            except Exception as inner_e:
                logger.warning(
                    "Streaming/split load failed for %s (%s). Retrying non-streaming or default split...",
                    info.name,
                    inner_e,
                )
                hf_ds = hf_load_dataset(*args, **hf_kwargs)
                actual_split = target_split if target_split in hf_ds else list(hf_ds.keys())[0]
                ds_split = hf_ds[actual_split]
                sub_ds = ds_split.select(range(min(len(ds_split), max_samples)))
                for row in sub_ds:
                    formatted = info.text_formatter(row) if info.text_formatter else dict(row)
                    raw_samples.append(formatted)

        except Exception as e:
            logger.warning(
                "Failed to load '%s' from HuggingFace Hub: %s. Falling back to synthetic generator.",
                dataset_name,
                str(e),
            )
            raw_samples = []

    # If empty or failed, run synthetic generator
    if not raw_samples:
        logger.info(
            "Generating %d synthetic samples for benchmark '%s' (seed=%d)...",
            max_samples,
            info.name,
            seed,
        )
        if info.synthetic_generator:
            raw_samples = info.synthetic_generator(max_samples, seed)
        else:
            raw_samples = [
                {"text": f"Synthetic benchmark text #{i} for {info.name}", "id": i}
                for i in range(max_samples)
            ]

    if shuffle:
        rng = random.Random(seed)
        rng.shuffle(raw_samples)

    if not return_dataloader:
        return raw_samples

    # Build PyTorch Dataset and DataLoader
    is_lm = info.category in ("language_modeling", "code")
    torch_ds = BenchmarkTorchDataset(
        samples=raw_samples,
        tokenizer=tokenizer,
        seq_len=seq_len,
        is_causal_lm=is_lm,
    )

    loader = DataLoader(
        torch_ds,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
    )
    return loader

Unified loader for 12+ NLP / LLM benchmark datasets.

Supports -----= - wikitext2 / wikitext103 - c4 / c4_en - ptb - pg19 - gsm8k - humaneval - mmlu - pile / pile_val - arc (ARC-Easy & ARC-Challenge) - hellaswag - lambada - glue (sst2, mrpc, qqp, mnli, qnli, rte, wnli, cola, stsb)

Args

-----=

dataset_name

Benchmark name or alias (e.g. 'wikitext2', 'c4', 'gsm8k', 'arc', 'glue/sst2').

tokenizer

HuggingFace tokenizer instance (optional).

seq_len

Maximum sequence length for tokenization.

batch_size

DataLoader batch size.

split

Dataset split ('train', 'validation', 'test'). Defaults to benchmark standard.

max_samples

Maximum number of samples to load/stream.

streaming

Whether to stream from HuggingFace Hub to avoid large downloads.

cache_dir

Local directory to cache datasets.

shuffle

Whether to shuffle samples.

num_workers

Number of workers for DataLoader.

use_synthetic

Force synthetic data generation for offline testing.

seed

Random seed for reproducibility.

return_dataloader

If True, returns PyTorch DataLoader; otherwise returns list of dicts.

**kwargs

Additional arguments passed to datasets.load_dataset.

Returns -----= torch.utils.data.DataLoader or List[Dict[str, Any]].

def load_dataset(config: ExperimentConfig) ‑> datasets.dataset_dict.DatasetDict

Expand source code

def load_dataset(config: ExperimentConfig) -> DatasetDict:
    """
    Load a dataset based on ExperimentConfig.
    Returns DatasetDict with at least 'train' and optionally 'validation'/'test' splits.
    """
    name = config.dataset_name.lower()

    # GLUE tasks (local Parquet)
    if name in {"cola", "sst2", "mrpc", "qqp", "stsb", "mnli", "qnli", "rte", "wnli"}:
        return _load_glue(config)

    # Synthetic data generation
    elif name == "synthetic":
        return _load_synthetic(config)

    # Custom file (CSV, JSON, Parquet)
    elif name.endswith((".csv", ".json", ".jsonl", ".parquet")):
        return _load_custom(config)

    # Multi‑choice dataset (ARC, HellaSwag, etc.)
    elif name == "multi_choice":
        return _load_multi_choice(config)

    # Fallback: HuggingFace dataset hub
    else:
        return _load_hub(config)

Load a dataset based on ExperimentConfig. Returns DatasetDict with at least 'train' and optionally 'validation'/'test' splits.

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 preprocess_function(examples,
tokenizer,
config: DataConfig) ‑> Any

Expand source code

def preprocess_function(examples, tokenizer, config: DataConfig) -> Any:
    """
    Generic preprocessing function handling all TaskTypes.
    Returns a dictionary with 'input_ids', 'attention_mask', and appropriate labels.
    """
    if config.custom_preprocess is not None:
        return config.custom_preprocess(examples, tokenizer, config)

    text1 = examples.get(config.text_column_1) if config.text_column_1 else None
    text2 = examples.get(config.text_column_2) if config.text_column_2 else None

    # Fallback auto-detection for text columns if specified column is missing
    if text1 is None:
        for col in [
            "sentence",
            "text",
            "sentence1",
            "premise",
            "question",
            "inputs",
            "input",
        ]:
            if col in examples:
                text1 = examples[col]
                break
        if text1 is None and hasattr(examples, "keys"):
            for k in examples.keys():
                if k not in ("label", "labels", "idx", "id"):
                    text1 = examples[k]
                    break

    if text2 is None and config.text_column_2:
        for col in ["sentence2", "hypothesis", "context"]:
            if col in examples:
                text2 = examples[col]
                break

    if text1 is not None:
        if isinstance(text1, list):
            text1 = [str(x) if x is not None else "" for x in text1]
        elif not isinstance(text1, str):
            text1 = str(text1)

    if text2 is not None:
        if isinstance(text2, list):
            text2 = [str(x) if x is not None else "" for x in text2]
        elif not isinstance(text2, str):
            text2 = str(text2)

    tokenizer_kwargs = {
        "truncation": config.truncation,
        "padding": config.padding,
        "max_length": config.max_length,
        "return_tensors": config.return_tensors,
    }

    if config.task_type == TaskType.LANGUAGE_MODELING:
        result = (
            tokenizer(text1, **tokenizer_kwargs) if text1 else tokenizer(text2, **tokenizer_kwargs)
        )
        return result

    if config.task_type == TaskType.TOKEN_CLASSIFICATION:
        tokenized = tokenizer(
            examples[config.tokens_column],
            truncation=config.truncation,
            padding=config.padding,
            max_length=config.max_length,
            is_split_into_words=True,
            return_tensors=config.return_tensors,
        )
        labels = []
        for i, label_list in enumerate(examples[config.tags_column]):
            word_ids = tokenized.word_ids(batch_index=i)
            label_ids = [-100] * len(word_ids)
            previous_word_idx = None
            for j, word_idx in enumerate(word_ids):
                if word_idx is None:
                    continue
                if word_idx != previous_word_idx:
                    label_ids[j] = label_list[word_idx]
                    previous_word_idx = word_idx
                else:
                    label_ids[j] = -100
            labels.append(label_ids)
        tokenized["labels"] = labels
        return tokenized

    if config.task_type in (
        TaskType.SINGLE_CLASSIFICATION,
        TaskType.PAIR_CLASSIFICATION,
        TaskType.REGRESSION,
    ):
        if text1 is not None and text2 is not None:
            tokenized = tokenizer(text1, text2, **tokenizer_kwargs)
        else:
            text_input = text1 if text1 is not None else text2
            tokenized = tokenizer(text_input, **tokenizer_kwargs)
        label_col = config.label_column
        if label_col not in examples:
            for candidate in ["label", "labels", "target", "gold_label"]:
                if candidate in examples:
                    label_col = candidate
                    break
        if label_col in examples:
            if config.task_type == TaskType.REGRESSION:
                tokenized["labels"] = [float(l) for l in examples[label_col]]
            else:
                tokenized["labels"] = [int(l) for l in examples[label_col]]
        return tokenized

    if config.task_type == TaskType.QA:
        tokenized = tokenizer(
            examples[config.question_column],
            examples[config.context_column],
            truncation="only_second" if config.truncation else False,
            padding=config.padding,
            max_length=config.max_length,
            return_offsets_mapping=True,
            return_tensors=config.return_tensors,
        )
        start_positions, end_positions = [], []
        for i, (ans, context) in enumerate(
            zip(examples[config.answer_column], examples[config.context_column])
        ):
            if isinstance(ans, dict):
                start_char = ans["answer_start"]
                end_char = start_char + len(ans["text"])
            else:
                start_char = ans
                end_char = start_char + len(examples[config.answer_column][i])
            offsets = tokenized["offset_mapping"][i]
            sequence_ids = tokenized.sequence_ids(i)
            start_token = end_token = None
            for j, (start, end) in enumerate(offsets):
                if sequence_ids[j] != 1:
                    continue
                if start <= start_char < end:
                    start_token = j
                if start < end_char <= end:
                    end_token = j
                    break
            if start_token is None:
                start_token = 0
            if end_token is None:
                end_token = len(offsets) - 1
            start_positions.append(start_token)
            end_positions.append(end_token)
        tokenized["start_positions"] = start_positions
        tokenized["end_positions"] = end_positions
        tokenized.pop("offset_mapping")
        return tokenized

    raise ValueError(f"Unsupported task type: {config.task_type}")

Generic preprocessing function handling all TaskTypes. Returns a dictionary with 'input_ids', 'attention_mask', and appropriate labels.

def run_universal_generation(output_dir: str | Path,
total_samples_per_domain: int = 5000,
domains: dict[str, Any] | None = None,
seed: int = 42) ‑> pandas.DataFrame

Expand source code

def run_universal_generation(
    output_dir: str | Path,
    total_samples_per_domain: int = 5000,
    domains: dict[str, Any] | None = None,
    seed: int = 42,
) -> pd.DataFrame:
    """Generate synthetic dataset across multiple domains and combine them."""
    random.seed(seed)
    domains = domains or DEFAULT_DOMAINS
    out_dir = Path(output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    all_dfs = []
    for domain_key in domains:
        half = total_samples_per_domain // 2
        pos = _generate_batch(domain_key, half, label=1, is_positive=True)
        neg = _generate_batch(domain_key, half, label=0, is_positive=False)
        df = pd.DataFrame(pos + neg).sample(frac=1, random_state=seed).reset_index(drop=True)
        df.to_csv(out_dir / f"data_{domain_key}.csv", index=False)
        all_dfs.append(df)

    final_df = (
        pd.concat(all_dfs, ignore_index=True)
        .sample(frac=1, random_state=seed)
        .reset_index(drop=True)
    )
    final_df.to_csv(out_dir / "combined_master_dataset.csv", index=False)
    logger.info(
        "Universal generation complete: %d total samples across %d domains.",
        len(final_df),
        len(domains),
    )
    return final_df

Generate synthetic dataset across multiple domains and combine them.

def selective_scan_fn(u: torch.Tensor,
delta: torch.Tensor,
A: torch.Tensor,
B: torch.Tensor,
C: torch.Tensor,
D: torch.Tensor | None = None,
delta_softplus: bool = True) ‑> torch.Tensor

Expand source code

def selective_scan_fn(
    u: torch.Tensor,
    delta: torch.Tensor,
    A: torch.Tensor,
    B: torch.Tensor,
    C: torch.Tensor,
    D: torch.Tensor | None = None,
    delta_softplus: bool = True,
) -> torch.Tensor:
    """
    Triton selective scan functional interface.
    Shapes:
        u:     (B, L, D)
        delta: (B, L, D)
        A:     (D, N)
        B:     (B, L, N)
        C:     (B, L, N)
        D:     (D,) optional
    """
    return TritonSelectiveScanFunction.apply(u, delta, A, B, C, D, delta_softplus)

Triton selective scan functional interface.

Shapes -----= u: (B, L, D) delta: (B, L, D) A: (D, N) B: (B, L, N) C: (B, L, N) D: (D,) optional

def triton_dispatch(kernel_name: str, *args: Any, **kwargs: Any) ‑> Any

Expand source code

def triton_dispatch(kernel_name: str, *args: Any, **kwargs: Any) -> Any:
    """
    Unified router that executes high-performance Triton kernels when supported,
    or transparently falls back to optimized PyTorch implementations.
    """
    if kernel_name == "linear_attn":
        from neuralops.core.triton_kernels.linear_attn import fwd, is_supported

        q, k, v = args[0], args[1], args[2]
        head_dim = q.shape[-1]
        seq_len = q.shape[-2]
        suitable, reason = check_triton_suitability(
            q, kernel_name, head_dim=head_dim, seq_len=seq_len
        )
        if suitable and is_supported(q, k, v, **kwargs):
            try:
                return fwd(q, k, v, **kwargs)
            except Exception as e:
                log_fallback_reason(kernel_name, f"execution error: {e}")
                return linear_attn_pytorch_fallback(q, k, v, **kwargs)
        else:
            if not suitable:
                log_fallback_reason(kernel_name, reason or "unsupported configuration")
            else:
                log_fallback_reason(kernel_name, "kernel.is_supported returned False")
            return linear_attn_pytorch_fallback(q, k, v, **kwargs)

    elif kernel_name == "sparse_attn":
        from neuralops.core.triton_kernels.sparse_attn import fwd, is_supported

        q, k, v = args[0], args[1], args[2]
        head_dim = q.shape[-1]
        seq_len = q.shape[-2]
        suitable, reason = check_triton_suitability(
            q, kernel_name, head_dim=head_dim, seq_len=seq_len
        )
        if suitable and is_supported(q, k, v, **kwargs):
            try:
                return fwd(q, k, v, **kwargs)
            except Exception as e:
                log_fallback_reason(kernel_name, f"execution error: {e}")
                return sparse_attn_pytorch_fallback(q, k, v, **kwargs)
        else:
            if not suitable:
                log_fallback_reason(kernel_name, reason or "unsupported configuration")
            else:
                log_fallback_reason(kernel_name, "kernel.is_supported returned False")
            return sparse_attn_pytorch_fallback(q, k, v, **kwargs)

    elif kernel_name == "low_rank":
        from neuralops.core.triton_kernels.low_rank import fwd, is_supported

        kw = dict(kwargs)
        x = args[0]
        U = args[1] if len(args) > 1 else kw.pop("U", None)
        V = args[2] if len(args) > 2 else kw.pop("V", None)
        bias = args[3] if len(args) > 3 else kw.pop("bias", None)
        suitable, reason = check_triton_suitability(x, kernel_name)
        if suitable and is_supported(x, U, V, bias=bias, **kw):
            try:
                return fwd(x, U, V, bias=bias, **kw)
            except Exception as e:
                log_fallback_reason(kernel_name, f"execution error: {e}")
                return low_rank_pytorch_fallback(x, U, V, bias=bias, **kw)
        else:
            if not suitable:
                log_fallback_reason(kernel_name, reason or "unsupported configuration")
            else:
                log_fallback_reason(kernel_name, "kernel.is_supported returned False")
            return low_rank_pytorch_fallback(x, U, V, bias=bias, **kw)

    elif kernel_name == "state_update":
        from neuralops.core.triton_kernels.state_update import fwd, is_supported

        kw = dict(kwargs)
        u = args[0]
        delta = args[1] if len(args) > 1 else kw.pop("delta", None)
        A = args[2] if len(args) > 2 else kw.pop("A", None)
        B = args[3] if len(args) > 3 else kw.pop("B", None)
        C = args[4] if len(args) > 4 else kw.pop("C", None)
        D = args[5] if len(args) > 5 else kw.pop("D", None)
        suitable, reason = check_triton_suitability(u, kernel_name)
        if suitable and is_supported(u, delta, A, B, C, D=D, **kw):
            try:
                return fwd(u, delta, A, B, C, D=D, **kw)
            except Exception as e:
                log_fallback_reason(kernel_name, f"execution error: {e}")
                return state_update_pytorch_fallback(u, delta, A, B, C, D=D, **kw)
        else:
            if not suitable:
                log_fallback_reason(kernel_name, reason or "unsupported configuration")
            else:
                log_fallback_reason(kernel_name, "kernel.is_supported returned False")
            return state_update_pytorch_fallback(u, delta, A, B, C, D=D, **kw)

    elif kernel_name == "gqa":
        from neuralops.core.triton_kernels.gqa import fwd, is_supported

        kw = dict(kwargs)
        k = args[0]
        v = args[1]
        num_kv_heads = args[2] if len(args) > 2 else kw.pop("num_kv_heads", None)
        num_heads = args[3] if len(args) > 3 else kw.pop("num_heads", None)
        suitable, reason = check_triton_suitability(k, kernel_name)
        if suitable and is_supported(k, v, num_kv_heads, num_heads, **kw):
            try:
                return fwd(k, v, num_kv_heads, num_heads, **kw)
            except Exception as e:
                log_fallback_reason(kernel_name, f"execution error: {e}")
                return gqa_pytorch_fallback(k, v, num_kv_heads, num_heads, **kw)
        else:
            if not suitable:
                log_fallback_reason(kernel_name, reason or "unsupported configuration")
            else:
                log_fallback_reason(kernel_name, "kernel.is_supported returned False")
            return gqa_pytorch_fallback(k, v, num_kv_heads, num_heads, **kw)

    elif kernel_name is None:
        raise ValueError("triton_kernel cannot be None when calling triton_dispatch")
    else:
        raise ValueError(
            f"Unknown triton kernel: '{kernel_name}'. Supported kernels: "
            "['linear_attn', 'sparse_attn', 'low_rank', 'state_update', 'gqa']"
        )

Unified router that executes high-performance Triton kernels when supported, or transparently falls back to optimized PyTorch implementations.

Classes

class ABTestResult (metric: str,
mean_a: float,
mean_b: float,
difference: float,
p_value: float,
significant: bool,
better_model: str,
confidence_interval: tuple[float, float])

Expand source code

@dataclass
class ABTestResult:
    metric: str
    mean_a: float
    mean_b: float
    difference: float
    p_value: float
    significant: bool
    better_model: str
    confidence_interval: tuple[float, float]

ABTestResult(metric: str, mean_a: float, mean_b: float, difference: float, p_value: float, significant: bool, better_model: str, confidence_interval: tuple[float, float])

Instance variables

var better_model : str

The type of the None singleton.

var confidence_interval : tuple[float, float]

The type of the None singleton.

var difference : float

The type of the None singleton.

var mean_a : float

The type of the None singleton.

var mean_b : float

The type of the None singleton.

var metric : str

The type of the None singleton.

var p_value : float

The type of the None singleton.

var significant : bool

The type of the None singleton.

class ABTester (alpha: float = 0.05)

Expand source code

class ABTester:
    def __init__(self, alpha: float = 0.05) -> None:
        self.alpha = alpha

    def compare(
        self, results_a: dict[str, list[float]], results_b: dict[str, list[float]]
    ) -> list[ABTestResult]:
        """Compare."""
        outcomes = []
        for metric in results_a:
            if metric not in results_b:
                continue
            a = np.array(results_a[metric])
            b = np.array(results_b[metric])
            t_stat, p_val = stats.ttest_ind(a, b, equal_var=False)
            mean_a, mean_b = np.mean(a), np.mean(b)
            diff = mean_b - mean_a
            se = np.sqrt(np.var(a) / len(a) + np.var(b) / len(b))
            ci = (diff - 1.96 * se, diff + 1.96 * se)
            outcomes.append(
                ABTestResult(
                    metric=metric,
                    mean_a=mean_a,
                    mean_b=mean_b,
                    difference=diff,
                    p_value=p_val,
                    significant=p_val < self.alpha,
                    better_model="B" if diff > 0 else "A",
                    confidence_interval=ci,
                )
            )
        return outcomes

Methods

def compare(self, results_a: dict[str, list[float]], results_b: dict[str, list[float]]) ‑> list[ABTestResult]

Expand source code

def compare(
    self, results_a: dict[str, list[float]], results_b: dict[str, list[float]]
) -> list[ABTestResult]:
    """Compare."""
    outcomes = []
    for metric in results_a:
        if metric not in results_b:
            continue
        a = np.array(results_a[metric])
        b = np.array(results_b[metric])
        t_stat, p_val = stats.ttest_ind(a, b, equal_var=False)
        mean_a, mean_b = np.mean(a), np.mean(b)
        diff = mean_b - mean_a
        se = np.sqrt(np.var(a) / len(a) + np.var(b) / len(b))
        ci = (diff - 1.96 * se, diff + 1.96 * se)
        outcomes.append(
            ABTestResult(
                metric=metric,
                mean_a=mean_a,
                mean_b=mean_b,
                difference=diff,
                p_value=p_val,
                significant=p_val < self.alpha,
                better_model="B" if diff > 0 else "A",
                confidence_interval=ci,
            )
        )
    return outcomes

Compare.

class BaseAttention (config: Any | None = None, rank: int | None = None, **kwargs: Any)

Expand source code

class BaseAttention(nn.Module):
    """Every custom attention module must inherit from this."""

    def __init__(self, config: Any | None = None, rank: int | None = None, **kwargs: Any) -> None:
        super().__init__()
        self.config = config
        self.rank = rank
        self.num_heads = (
            kwargs.get("num_heads")
            or kwargs.get("num_attention_heads")
            or (
                getattr(
                    config,
                    "num_attention_heads",
                    getattr(config, "num_heads", getattr(config, "n_heads", 12)),
                )
                if config is not None
                else 12
            )
        )
        self.hidden_size = (
            kwargs.get("hidden_size")
            or kwargs.get("dim")
            or kwargs.get("d_model")
            or (
                getattr(
                    config,
                    "hidden_size",
                    getattr(config, "d_model", getattr(config, "dim", 768)),
                )
                if config is not None
                else 768
            )
        )
        self.head_dim = getattr(
            config,
            "d_kv",
            getattr(
                config,
                "head_dim",
                self.hidden_size // self.num_heads if self.num_heads > 0 else 64,
            ),
        )

    @abstractmethod
    def forward(
        self,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
        **kwargs,
    ) -> tuple[Any, Any | None]:
        """Must return (output, attention_weights_or_None)."""
        raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

    def copy_weights(self, old_attn: Any) -> None:
        """Universal fallback: copy pretrained weights from old attention layer."""
        if old_attn is None:
            return
        import torch

        from neuralops.core.adapters import UniversalAdapter

        adapter = UniversalAdapter()

        try:
            w_q, b_q = adapter.get_weight_bias(old_attn, "q")
            w_k, b_k = adapter.get_weight_bias(old_attn, "k")
            w_v, b_v = adapter.get_weight_bias(old_attn, "v")
            w_o, b_o = adapter.get_weight_bias(old_attn, "o")

            with torch.no_grad():
                dim = self.hidden_size
                for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                    if hasattr(self, q_name):
                        mod = getattr(self, q_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                            mod.weight.copy_(w_q)
                            if b_q is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_q)

                for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                    if hasattr(self, k_name):
                        mod = getattr(self, k_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                            mod.weight.copy_(w_k)
                            if b_k is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_k)

                for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                    if hasattr(self, v_name):
                        mod = getattr(self, v_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                            mod.weight.copy_(w_v)
                            if b_v is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_v)

                for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                    if hasattr(self, o_name):
                        mod = getattr(self, o_name)
                        if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                            mod.weight.copy_(w_o)
                            if b_o is not None and getattr(mod, "bias", None) is not None:
                                mod.bias.copy_(b_o)

                # Handle fused in_proj
                if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                    if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                        self.in_proj.weight[:dim].copy_(w_v[:dim])
                        self.in_proj.weight[dim:].copy_(w_q[:dim])
                    elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                        self.in_proj.weight[:dim].copy_(w_q[:dim])
                        self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                        self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
        except (AttributeError, RuntimeError, ValueError) as e:
            logger.debug("BaseAttention.copy_weights failed: %s", e)
        except Exception as e:
            logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

    def extra_repr(self) -> str:
        """Extra repr."""
        return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Every custom attention module must inherit from this.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Subclasses

• AdaptiveSpanAttention
• AdaptiveSparseAttention
• AdaptiveSVDAttention
• AdaSplashAttention
• AegisHybridAttention
• AegisSpectronFP8Attention
• AegisAttentionV2
• AegisAttention
• AegisAttentionV3
• ASVDAttention
• BAMUMoEBlock
• BlockTransformerHIBA
• ChelaSlayAttention
• ClusterAttention
• CompressedAttentionLayer
• CompressedDistilBertAttention
• CosformerAttention
• DecayLinearAttention
• EATBlock
• ExpertChoiceMoEAttention
• FactorisedMatrixState
• FlashAttention3
• FlashAttention3FP8
• FullAttention
• FusedSVDKernelAttention
• FUSIONATTN
• GalerkinAttention
• GatedDeltaNet
• GatedDeltaNetBlock
• GatedLinearAttention
• GoldFinchAHNBlock
• GroupedQueryAttention
• GraniteHybridBlock
• GriffinSSDBlock
• HALOReGLAAttention
• HammingAttention
• HGRN2GatedDeltaNet
• HGRN2Layer
• HGRN2OuterBlock
• HLAAttention
• HybridAttention
• HymbaTransXBlock
• InfiniAttention
• Jamba2Block
• JambaInfoBlock
• KDAAttention
• LEKAV1Attention
• LEKAv2Attention
• LEKAV3Attention
• LinformerAttention
• LoLAAttention
• LoRAAttention
• LSHAttention
• MambaSVDFlash
• MinimalNOVAAttention
• MinimalSVDAttention
• MixAttentionMoSKA
• MLAAttention
• MoBAAttention
• NASCAttention
• NEXUSAttention
• NativeHybridAttention
• NativeSparseAttention
• NystromAttention
• OptimizedHybridAttention
• ORIONAttention
• PADReAttention
• PerformerAttention
• PHANTOMAttention
• PrimalAttention
• ProgressiveTokenPruning
• QCQAAttention
• RALALinearAttention
• REAPMoEAttention
• ReGLAAttention
• ReplaceAttention
• ResSVDAttention
• RoutingAttention
• RPCAttention
• SALTAttention
• SambaMoEBlock
• SelectiveSSM
• SelectiveSSM
• SOAAttention
• SpanExpandedAttention
• SpectraV2Attention
• SpectraformerAttention
• SpectralMambaAttention
• SpectronAttention
• SqueezingHeadsDistiller
• SVDLinearAttention
• SVDCompressedAttention
• SVDLinearAttention
• SVDLLMV2Attention
• SVDMamba
• SVDProjectionAttention
• YourCustomAttention
• TLinFormerAttention
• TokenCompressedAttention
• TPAttention
• TritonCosformerAttention
• TritonSVDLinearAttention
• UnifiedRoPE
• SelectiveMamba
• ViGAttention
• VSpecAttention
• WeightedGroupedQueryAttention
• XAttention
• ZeroSAttention
• NLISoftmaxAttention
• ReplaceAttention

Methods

def copy_weights(self, old_attn: Any) ‑> None

Expand source code

def copy_weights(self, old_attn: Any) -> None:
    """Universal fallback: copy pretrained weights from old attention layer."""
    if old_attn is None:
        return
    import torch

    from neuralops.core.adapters import UniversalAdapter

    adapter = UniversalAdapter()

    try:
        w_q, b_q = adapter.get_weight_bias(old_attn, "q")
        w_k, b_k = adapter.get_weight_bias(old_attn, "k")
        w_v, b_v = adapter.get_weight_bias(old_attn, "v")
        w_o, b_o = adapter.get_weight_bias(old_attn, "o")

        with torch.no_grad():
            dim = self.hidden_size
            for q_name in ["W_q", "q_proj", "q_lin", "Wq", "query", "q"]:
                if hasattr(self, q_name):
                    mod = getattr(self, q_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_q.shape:
                        mod.weight.copy_(w_q)
                        if b_q is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_q)

            for k_name in ["W_k", "k_proj", "k_lin", "Wk", "key", "k"]:
                if hasattr(self, k_name):
                    mod = getattr(self, k_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_k.shape:
                        mod.weight.copy_(w_k)
                        if b_k is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_k)

            for v_name in ["W_v", "v_proj", "v_lin", "Wv", "value", "v"]:
                if hasattr(self, v_name):
                    mod = getattr(self, v_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_v.shape:
                        mod.weight.copy_(w_v)
                        if b_v is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_v)

            for o_name in ["out", "out_proj", "o_proj", "out_lin", "proj", "dense"]:
                if hasattr(self, o_name):
                    mod = getattr(self, o_name)
                    if hasattr(mod, "weight") and mod.weight.shape == w_o.shape:
                        mod.weight.copy_(w_o)
                        if b_o is not None and getattr(mod, "bias", None) is not None:
                            mod.bias.copy_(b_o)

            # Handle fused in_proj
            if hasattr(self, "in_proj") and hasattr(self.in_proj, "weight"):
                if self.in_proj.weight.shape == torch.Size([dim * 2, dim]):
                    self.in_proj.weight[:dim].copy_(w_v[:dim])
                    self.in_proj.weight[dim:].copy_(w_q[:dim])
                elif self.in_proj.weight.shape == torch.Size([dim * 3, dim]):
                    self.in_proj.weight[:dim].copy_(w_q[:dim])
                    self.in_proj.weight[dim : 2 * dim].copy_(w_k[:dim])
                    self.in_proj.weight[2 * dim :].copy_(w_v[:dim])
    except (AttributeError, RuntimeError, ValueError) as e:
        logger.debug("BaseAttention.copy_weights failed: %s", e)
    except Exception as e:
        logger.warning("Unexpected error in BaseAttention.copy_weights: %s", e)

Universal fallback: copy pretrained weights from old attention layer.

def extra_repr(self) ‑> str

Expand source code

def extra_repr(self) -> str:
    """Extra repr."""
    return f"num_heads={self.num_heads}, head_dim={self.head_dim}, rank={self.rank}"

Extra repr.

def forward(self,
hidden_states,
attention_mask=None,
head_mask=None,
output_attentions=False,
**kwargs) ‑> tuple[typing.Any, Any | None]

Expand source code

@abstractmethod
def forward(
    self,
    hidden_states,
    attention_mask=None,
    head_mask=None,
    output_attentions=False,
    **kwargs,
) -> tuple[Any, Any | None]:
    """Must return (output, attention_weights_or_None)."""
    raise NotImplementedError("Subclasses of BaseAttention must implement forward()")

Must return (output, attention_weights_or_None).

class BaseCompressor (model: torch.nn.modules.module.Module, name: str = 'base', **kwargs)

Expand source code

class BaseCompressor:
    """Base class for all compressors. Not a nn.Module to avoid parameter storage."""

    def __init__(self, model: nn.Module, name: str = "base", **kwargs) -> None:
        self.model = model
        self.name = name
        self.kwargs = kwargs
        self._original_params = self._count_params(model)

    def __call__(self, *args, **kwargs) -> Any:
        """Allow compressor instance to act as a callable model forward wrapper."""
        if not getattr(self, "_compressed_once", False):
            self.compress()
            self._compressed_once = True
        return self.model(*args, **kwargs)

    @abstractmethod
    def compress(self) -> tuple[nn.Module, dict[str, Any]]:
        """Return (compressed_model, statistics)."""
        raise NotImplementedError("Subclasses of BaseCompressor must implement compress()")

    def _count_params(self, model: nn.Module) -> int:
        """count params."""
        return sum(p.numel() for p in model.parameters())

    def get_compression_stats(self) -> dict[str, Any]:
        """Return the compression stats."""
        compressed = self._count_params(self.model)
        orig = self._original_params
        return {
            "compressor_name": self.name,
            "original_parameters": orig,
            "compressed_parameters": compressed,
            "saved_parameters": orig - compressed,
            "compression_ratio": compressed / orig if orig > 0 else 1.0,
            "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0),
        }

    def summary(self) -> str:
        """Summary."""
        s = self.get_compression_stats()
        return (
            f"\n========================================\n"
            f"Compression Summary: {self.name}\n"
            f"========================================\n"
            f"Original parameters:  {s['original_parameters']:,}\n"
            f"Compressed parameters: {s['compressed_parameters']:,}\n"
            f"Parameters saved:     {s['saved_parameters']:,}\n"
            f"Compression ratio:    {s['compression_ratio']:.3f}\n"
            f"Reduction:            {s['reduction_percentage']:.2f}%\n"
            f"========================================\n"
        )

Base class for all compressors. Not a nn.Module to avoid parameter storage.

Subclasses

• A3Compressor
• AASVDCompressor
• AATC_XKVCompressor
• AdaptiveRanksCompressor
• AdaSVDCompressor
• ADMMCompressor
• ADMMPruner
• ARMORPruner
• ASTPruner
• AstroQuantizer
• BonsaiPruner
• CacheNotesCompressor
• CoMeCompressor
• COMPOTCompressor
• CoopQQuantizer
• CRCompressor
• CURCompressor
• CURMoECompressor
• DeltaKVCompressor
• DirectSVDCompressor
• DobiSVDCompressor
• DRONECompressor
• DynamicEntropySVDCompressor
• ENLightenCompressor
• EntropyBudgetCompressor
• ERCCompressor
• FAQQuantizer
• FLAPCompressor
• FLATLLMCompressor
• FRanDICompressor
• FWSVDCompressor
• GFWSCompressor
• GreCoDeCCompressor
• HashedNetCompressor
• HuffmanCompressor
• HybridSVDCompressor
• HybridKVCompressor
• KMeansCompressor
• KroneckerCompressor
• KroneckerTTLoRACompressor
• KVTCCompressor
• LIFTRCompressor
• LiLLaMADistiller
• LSCCompressor
• MambaShedderCompressor
• MGIECompressor
• MGIESVDBudgetCompressor
• MLoRQCompressor
• MoDeGPTCompressor
• NLISoftmaxCompressor
• NMFCompressor
• OATSCompressor
• OptimalSubspaceProjectionCompressor
• ProxSparsePruner
• ReplaceMeCompressor
• ResSVDCompressor
• ROASTCompressor
• RotationalSliceGPTINT4Compressor
• RPCACompressor
• SAESSVDCompressorV2
• SALRCompressor
• SDRMCompressor
• SliceGPTCompressor
• SliceGPTPCACompressor
• SLiMCompressor
• LowRankCompressor
• SparseGPTCompressor
• SubspaceProjectionCompressor
• SVDLinearCompressor
• SVDLLMCompressor
• SwiftSVDCompressor
• SwiftSVDActivationCompressor
• ToMoECompressor
• TTCompressor
• UniQLCompressor
• ZeroKVNeoCompressor

Methods

def compress(self) ‑> tuple[torch.nn.modules.module.Module, dict[str, typing.Any]]

Expand source code

@abstractmethod
def compress(self) -> tuple[nn.Module, dict[str, Any]]:
    """Return (compressed_model, statistics)."""
    raise NotImplementedError("Subclasses of BaseCompressor must implement compress()")

Return (compressed_model, statistics).

def get_compression_stats(self) ‑> dict[str, typing.Any]

Expand source code

def get_compression_stats(self) -> dict[str, Any]:
    """Return the compression stats."""
    compressed = self._count_params(self.model)
    orig = self._original_params
    return {
        "compressor_name": self.name,
        "original_parameters": orig,
        "compressed_parameters": compressed,
        "saved_parameters": orig - compressed,
        "compression_ratio": compressed / orig if orig > 0 else 1.0,
        "reduction_percentage": ((orig - compressed) / orig * 100 if orig > 0 else 0.0),
    }

Return the compression stats.

def summary(self) ‑> str

Expand source code

def summary(self) -> str:
    """Summary."""
    s = self.get_compression_stats()
    return (
        f"\n========================================\n"
        f"Compression Summary: {self.name}\n"
        f"========================================\n"
        f"Original parameters:  {s['original_parameters']:,}\n"
        f"Compressed parameters: {s['compressed_parameters']:,}\n"
        f"Parameters saved:     {s['saved_parameters']:,}\n"
        f"Compression ratio:    {s['compression_ratio']:.3f}\n"
        f"Reduction:            {s['reduction_percentage']:.2f}%\n"
        f"========================================\n"
    )

Summary.

class BaseTrainer

Expand source code

class BaseTrainer(ABC):
    """Abstract base class for all trainers."""

    @abstractmethod
    def train(self) -> dict[str, Any]:
        """Fit the model on the provided dataset."""
        pass

    @abstractmethod
    def evaluate(self) -> dict[str, Any]:
        """Run evaluate and return results."""
        pass

Abstract base class for all trainers.

Ancestors

• abc.ABC

Subclasses

• CustomLoopTrainer
• HuggingFaceTrainer
• RLTrainer
• UnslothTrainer

Methods

def evaluate(self) ‑> dict[str, typing.Any]

Expand source code

@abstractmethod
def evaluate(self) -> dict[str, Any]:
    """Run evaluate and return results."""
    pass

Run evaluate and return results.

def train(self) ‑> dict[str, typing.Any]

Expand source code

@abstractmethod
def train(self) -> dict[str, Any]:
    """Fit the model on the provided dataset."""
    pass

Fit the model on the provided dataset.

class BenchmarkDatasetInfo (name: str,
aliases: list[str],
category: str,
default_hf_path: str,
default_subset: str | None = None,
default_split: str = 'validation',
text_formatter: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
synthetic_generator: Callable[[int, int], list[dict[str, Any]]] | None = None,
description: str = '')

Expand source code

@dataclass
class BenchmarkDatasetInfo:
    """Metadata describing a benchmark dataset and its default loading strategy."""

    name: str
    aliases: list[str]
    category: str  # e.g., "language_modeling", "reasoning", "code", "multi_choice", "glue"
    default_hf_path: str
    default_subset: str | None = None
    default_split: str = "validation"
    text_formatter: Callable[[dict[str, Any]], dict[str, Any]] | None = None
    synthetic_generator: Callable[[int, int], list[dict[str, Any]]] | None = None
    description: str = ""

Metadata describing a benchmark dataset and its default loading strategy.

Instance variables

var aliases : list[str]

The type of the None singleton.

var category : str

The type of the None singleton.

var default_hf_path : str

The type of the None singleton.

var default_split : str

The type of the None singleton.

var default_subset : str | None

The type of the None singleton.

var description : str

The type of the None singleton.

var name : str

The type of the None singleton.

var synthetic_generator : Callable[[int, int], list[dict[str, typing.Any]]] | None

The type of the None singleton.

var text_formatter : Callable[[dict[str, typing.Any]], dict[str, typing.Any]] | None

The type of the None singleton.

class BenchmarkDatasetRegistry

Expand source code

class BenchmarkDatasetRegistry:
    """Registry maintaining metadata and formatters for NLP/LLM benchmarks."""

    _registry: dict[str, BenchmarkDatasetInfo] = {}

    @classmethod
    def register(cls, info: BenchmarkDatasetInfo) -> None:
        """Register a new benchmark dataset info."""
        key = info.name.lower()
        cls._registry[key] = info
        for alias in info.aliases:
            cls._registry[alias.lower()] = info

    @classmethod
    def get(cls, name: str) -> BenchmarkDatasetInfo | None:
        """Retrieve benchmark dataset info by name or alias."""
        return cls._registry.get(
            name.lower().replace("-", "").replace("_", "")
        ) or cls._registry.get(name.lower())

    @classmethod
    def is_registered(cls, name: str) -> bool:
        """Check if a benchmark dataset name is registered."""
        return cls.get(name) is not None

    @classmethod
    def list_benchmarks(cls) -> list[str]:
        """List canonical names of all registered benchmarks."""
        canonical = set()
        for info in cls._registry.values():
            canonical.add(info.name)
        return sorted(list(canonical))

    @classmethod
    def list_datasets(cls) -> list[str]:
        """Alias for list_benchmarks."""
        return cls.list_benchmarks()

Registry maintaining metadata and formatters for NLP/LLM benchmarks.

Static methods

def get(name: str) ‑> BenchmarkDatasetInfo | None

Retrieve benchmark dataset info by name or alias.

def is_registered(name: str) ‑> bool

Check if a benchmark dataset name is registered.

def list_benchmarks() ‑> list[str]

List canonical names of all registered benchmarks.

def list_datasets() ‑> list[str]

Alias for list_benchmarks.

def register(info: BenchmarkDatasetInfo) ‑> None

Register a new benchmark dataset info.

class CompressionPipeline (model: torch.nn.modules.module.Module,
config: dict[str, typing.Any],
tokenizer=None)

Expand source code

class CompressionPipeline:
    """Apply a sequence of compression steps to a model."""

    def __init__(
        self,
        model: nn.Module,
        config: dict[str, Any],
        tokenizer=None,
    ) -> None:
        self.model = model
        self.config = config
        self.tokenizer = tokenizer
        if tokenizer is not None:
            self.model.tokenizer = tokenizer

    def run(self) -> nn.Module:
        """Execute the run operation and return results."""
        compressors_cfg = self.config.get("compressors", [])
        if not compressors_cfg:
            logger.info("No compressors specified. Returning original model.")
            return self.model

        seed = self.config.get("seed")
        if seed is not None:
            set_seed(seed)

        for i, comp_cfg in enumerate(compressors_cfg):
            method = comp_cfg.pop("method")
            logger.info("Step %d: Applying compressor '%s'", i + 1, method)
            compressor_cls = COMPRESSOR_REGISTRY.resolve(method)
            merged_cfg = {**self.config, **comp_cfg}
            merged_cfg.pop("compressors", None)
            compressor = compressor_cls(self.model, **merged_cfg)
            self.model, stats = compressor.compress()
            if stats:
                logger.info("%s", compressor.summary())

        return self.model

Apply a sequence of compression steps to a model.

Methods

def run(self) ‑> torch.nn.modules.module.Module

Expand source code

def run(self) -> nn.Module:
    """Execute the run operation and return results."""
    compressors_cfg = self.config.get("compressors", [])
    if not compressors_cfg:
        logger.info("No compressors specified. Returning original model.")
        return self.model

    seed = self.config.get("seed")
    if seed is not None:
        set_seed(seed)

    for i, comp_cfg in enumerate(compressors_cfg):
        method = comp_cfg.pop("method")
        logger.info("Step %d: Applying compressor '%s'", i + 1, method)
        compressor_cls = COMPRESSOR_REGISTRY.resolve(method)
        merged_cfg = {**self.config, **comp_cfg}
        merged_cfg.pop("compressors", None)
        compressor = compressor_cls(self.model, **merged_cfg)
        self.model, stats = compressor.compress()
        if stats:
            logger.info("%s", compressor.summary())

    return self.model

Execute the run operation and return results.

class CustomLoopTrainer (model: nn.Module,
train_dataset: Any,
eval_dataset: Any,
config: dict[str, Any],
results_manager: Any | None = None,
device: torch.device | None = None,
num_labels: int = 2,
is_multichoice: bool = False,
compute_metrics_fn: Callable | None = None,
**kwargs: Any)

Expand source code

@TrainerRegistry.register("custom_loop")
class CustomLoopTrainer(BaseTrainer):
    """
    Custom loop trainer optimized for maximum GPU compute throughput, mixed-precision training,
    and distributed data parallel (DDP) gradient accumulation without redundant all-reduce overhead.
    """

    def __init__(
        self,
        model: nn.Module,
        train_dataset: Any,
        eval_dataset: Any,
        config: dict[str, Any],
        results_manager: Any | None = None,
        device: torch.device | None = None,
        num_labels: int = 2,
        is_multichoice: bool = False,
        compute_metrics_fn: Callable | None = None,
        **kwargs: Any,
    ) -> None:
        self.model = model
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.config = config or {}
        self.results_manager = results_manager
        self.device = device or (
            torch.device(
                "cuda"
                if torch.cuda.is_available()
                else (
                    "mps"
                    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
                    else "cpu"
                )
            )
        )
        self.num_labels = num_labels
        self.is_multichoice = is_multichoice
        self.compute_metrics_fn = compute_metrics_fn or get_task_metrics(
            self.config.get("task_name", "mrpc")
        )
        self.kwargs = kwargs

    def _make_loader(self, dataset: Any, shuffle: bool = False) -> DataLoader:
        """Make DataLoader optimized for high-throughput GPU training."""
        pin = self.device.type in ["cuda", "mps"]
        num_workers = int(self.config.get("num_workers", 4 if self.device.type == "cuda" else 0))
        batch_size = int(self.config.get("batch_size", 128))
        loader_kwargs = {
            "batch_size": batch_size,
            "shuffle": shuffle,
            "pin_memory": pin,
            "num_workers": num_workers,
            "collate_fn": (dataset.collate_fn if hasattr(dataset, "collate_fn") else None),
        }
        if num_workers > 0:
            loader_kwargs["prefetch_factor"] = 2
        if torch.cuda.is_available():
            torch.backends.cudnn.benchmark = True
        return DataLoader(dataset, **loader_kwargs)

    def train_epoch(
        self,
        train_loader: DataLoader,
        optimizer: torch.optim.Optimizer,
        scaler: torch.cuda.amp.GradScaler | None,
        gradient_accumulation_steps: int = 1,
        max_grad_norm: float = 1.0,
        amp_dtype: torch.dtype = torch.float32,
        use_amp: bool = False,
        pbar: Any | None = None,
    ) -> float:
        """
        Execute a single training epoch with DDP gradient accumulation synchronization optimization.

        # FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()
        # to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.
        """
        self.model.train()
        total_epoch_loss = 0.0
        num_batches = len(train_loader)
        accum_steps = max(1, gradient_accumulation_steps)

        for step, batch in enumerate(train_loader):
            # FIX: Determine whether this micro-batch is the final step of the accumulation window or epoch
            is_sync_step = ((step + 1) % accum_steps == 0) or (step + 1 == num_batches)

            # FIX: Wrap non-final micro-batches in model.no_sync() for DDP to eliminate redundant all-reduce synchronization
            sync_context = (
                contextlib.nullcontext()
                if (
                    is_sync_step
                    or not hasattr(self.model, "no_sync")
                    or not callable(self.model.no_sync)
                )
                else self.model.no_sync()
            )

            with sync_context:
                with torch.autocast(
                    device_type=self.device.type,
                    dtype=amp_dtype,
                    enabled=(use_amp and self.device.type in ["cuda", "mps"]),
                ):
                    if self.is_multichoice:
                        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                        true_labels = batch["labels"].to(self.device, non_blocking=True)
                        bs, nc, sl = input_ids.shape
                        flat_ids = input_ids.view(-1, sl)
                        flat_mask = attention_mask.view(-1, sl)
                        outputs = self.model(input_ids=flat_ids, attention_mask=flat_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        combined_labels = true_labels.repeat_interleave(nc)
                        loss = nn.functional.cross_entropy(logits, combined_labels)
                    else:
                        input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                        attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                        labels = batch["labels"].to(self.device, non_blocking=True)
                        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                        logits = (
                            outputs["logits"]
                            if isinstance(outputs, dict)
                            else (outputs.logits if hasattr(outputs, "logits") else outputs)
                        )
                        loss = nn.functional.cross_entropy(logits, labels)

                    # FIX: Scale loss by accumulation factor so gradient magnitudes match the full effective batch size
                    if accum_steps > 1:
                        loss = loss / accum_steps

                if torch.isfinite(loss):
                    if scaler is not None and scaler.is_enabled():
                        scaler.scale(loss).backward()
                    else:
                        loss.backward()

                    curr_loss = loss.detach().item() * accum_steps
                    total_epoch_loss += curr_loss
                else:
                    curr_loss = 0.0

                # FIX: Only unscale, clip gradients, and step optimizer when synchronized on accumulation boundary
                if scaler is not None and scaler.is_enabled():
                    try:
                        scaler.unscale_(optimizer)
                    except RuntimeError:
                        pass
                    clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                    optimizer.step()
                optimizer.zero_grad(set_to_none=True)

            if pbar is not None and (
                (step + 1) % max(1, int(self.config.get("logging_steps", 10))) == 0
                or step == num_batches - 1
            ):
                pbar.set_postfix({"loss": f"{curr_loss:.4f}"})

        return total_epoch_loss / max(1, num_batches)

    def train(self) -> dict[str, Any]:
        """Fit the model on the provided dataset with full DDP and accumulation safety."""
        train_loader = self._make_loader(self.train_dataset, shuffle=True)
        val_loader = self._make_loader(self.eval_dataset, shuffle=False)

        class AttrDict:
            pass

        cfg = AttrDict()
        cfg.learning_rate = float(self.config.get("learning_rate", 2e-5))
        cfg.epochs = int(self.config.get("epochs", 3))
        cfg.max_grad_norm = float(self.config.get("max_grad_norm", 1.0))
        cfg.gradient_accumulation_steps = int(self.config.get("gradient_accumulation_steps", 1))
        cfg.fp16 = bool(self.config.get("fp16", False))
        cfg.amp = bool(self.config.get("amp", True))
        cfg.logging_steps = int(self.config.get("logging_steps", 10))

        result = train_model(
            self.model,
            train_loader,
            val_loader,
            self.device,
            cfg,
            self.num_labels,
            self.is_multichoice,
        )

        if self.results_manager is not None:
            if hasattr(self.results_manager, "save_final_model"):
                self.results_manager.save_final_model(self.model)
            if hasattr(self.results_manager, "save_plots"):
                self.results_manager.save_plots()

        return result

    def evaluate(self) -> dict[str, Any]:
        """Run evaluate and return results."""
        val_loader = self._make_loader(self.eval_dataset, shuffle=False)
        return evaluate_model(
            self.model,
            val_loader,
            self.device,
            self.num_labels,
            self.is_multichoice,
            compute_all_metrics=True,
        )

Custom loop trainer optimized for maximum GPU compute throughput, mixed-precision training, and distributed data parallel (DDP) gradient accumulation without redundant all-reduce overhead.

Ancestors

• BaseTrainer
• abc.ABC

Methods

def train(self) ‑> dict[str, typing.Any]

Expand source code

def train(self) -> dict[str, Any]:
    """Fit the model on the provided dataset with full DDP and accumulation safety."""
    train_loader = self._make_loader(self.train_dataset, shuffle=True)
    val_loader = self._make_loader(self.eval_dataset, shuffle=False)

    class AttrDict:
        pass

    cfg = AttrDict()
    cfg.learning_rate = float(self.config.get("learning_rate", 2e-5))
    cfg.epochs = int(self.config.get("epochs", 3))
    cfg.max_grad_norm = float(self.config.get("max_grad_norm", 1.0))
    cfg.gradient_accumulation_steps = int(self.config.get("gradient_accumulation_steps", 1))
    cfg.fp16 = bool(self.config.get("fp16", False))
    cfg.amp = bool(self.config.get("amp", True))
    cfg.logging_steps = int(self.config.get("logging_steps", 10))

    result = train_model(
        self.model,
        train_loader,
        val_loader,
        self.device,
        cfg,
        self.num_labels,
        self.is_multichoice,
    )

    if self.results_manager is not None:
        if hasattr(self.results_manager, "save_final_model"):
            self.results_manager.save_final_model(self.model)
        if hasattr(self.results_manager, "save_plots"):
            self.results_manager.save_plots()

    return result

Fit the model on the provided dataset with full DDP and accumulation safety.

def train_epoch(self,
train_loader: DataLoader,
optimizer: torch.optim.Optimizer,
scaler: torch.cuda.amp.GradScaler | None,
gradient_accumulation_steps: int = 1,
max_grad_norm: float = 1.0,
amp_dtype: torch.dtype = torch.float32,
use_amp: bool = False,
pbar: Any | None = None) ‑> float

Expand source code

def train_epoch(
    self,
    train_loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    scaler: torch.cuda.amp.GradScaler | None,
    gradient_accumulation_steps: int = 1,
    max_grad_norm: float = 1.0,
    amp_dtype: torch.dtype = torch.float32,
    use_amp: bool = False,
    pbar: Any | None = None,
) -> float:
    """
    Execute a single training epoch with DDP gradient accumulation synchronization optimization.

    # FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()
    # to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.
    """
    self.model.train()
    total_epoch_loss = 0.0
    num_batches = len(train_loader)
    accum_steps = max(1, gradient_accumulation_steps)

    for step, batch in enumerate(train_loader):
        # FIX: Determine whether this micro-batch is the final step of the accumulation window or epoch
        is_sync_step = ((step + 1) % accum_steps == 0) or (step + 1 == num_batches)

        # FIX: Wrap non-final micro-batches in model.no_sync() for DDP to eliminate redundant all-reduce synchronization
        sync_context = (
            contextlib.nullcontext()
            if (
                is_sync_step
                or not hasattr(self.model, "no_sync")
                or not callable(self.model.no_sync)
            )
            else self.model.no_sync()
        )

        with sync_context:
            with torch.autocast(
                device_type=self.device.type,
                dtype=amp_dtype,
                enabled=(use_amp and self.device.type in ["cuda", "mps"]),
            ):
                if self.is_multichoice:
                    input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                    attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                    true_labels = batch["labels"].to(self.device, non_blocking=True)
                    bs, nc, sl = input_ids.shape
                    flat_ids = input_ids.view(-1, sl)
                    flat_mask = attention_mask.view(-1, sl)
                    outputs = self.model(input_ids=flat_ids, attention_mask=flat_mask)
                    logits = (
                        outputs["logits"]
                        if isinstance(outputs, dict)
                        else (outputs.logits if hasattr(outputs, "logits") else outputs)
                    )
                    combined_labels = true_labels.repeat_interleave(nc)
                    loss = nn.functional.cross_entropy(logits, combined_labels)
                else:
                    input_ids = batch["input_ids"].to(self.device, non_blocking=True)
                    attention_mask = batch["attention_mask"].to(self.device, non_blocking=True)
                    labels = batch["labels"].to(self.device, non_blocking=True)
                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
                    logits = (
                        outputs["logits"]
                        if isinstance(outputs, dict)
                        else (outputs.logits if hasattr(outputs, "logits") else outputs)
                    )
                    loss = nn.functional.cross_entropy(logits, labels)

                # FIX: Scale loss by accumulation factor so gradient magnitudes match the full effective batch size
                if accum_steps > 1:
                    loss = loss / accum_steps

            if torch.isfinite(loss):
                if scaler is not None and scaler.is_enabled():
                    scaler.scale(loss).backward()
                else:
                    loss.backward()

                curr_loss = loss.detach().item() * accum_steps
                total_epoch_loss += curr_loss
            else:
                curr_loss = 0.0

            # FIX: Only unscale, clip gradients, and step optimizer when synchronized on accumulation boundary
            if scaler is not None and scaler.is_enabled():
                try:
                    scaler.unscale_(optimizer)
                except RuntimeError:
                    pass
                clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                scaler.step(optimizer)
                scaler.update()
            else:
                clip_grad_norm_factorized(self.model, max_norm=max_grad_norm)
                optimizer.step()
            optimizer.zero_grad(set_to_none=True)

        if pbar is not None and (
            (step + 1) % max(1, int(self.config.get("logging_steps", 10))) == 0
            or step == num_batches - 1
        ):
            pbar.set_postfix({"loss": f"{curr_loss:.4f}"})

    return total_epoch_loss / max(1, num_batches)

Execute a single training epoch with DDP gradient accumulation synchronization optimization.

FIX: Non-final micro-batches are wrapped in model.no_sync() or contextlib.nullcontext()

to eliminate redundant all-reduce synchronization across distributed workers during gradient accumulation.

Inherited members

• BaseTrainer:
  • evaluate

class DataConfig (task_type: TaskType = TaskType.SINGLE_CLASSIFICATION,
task_name: str | None = None,
max_length: int = 128,
padding: Any = 'max_length',
truncation: bool = True,
return_tensors: str = 'pt',
text_column: str | None = None,
text_column_1: str | None = None,
text_column_2: str | None = None,
label_column: str | None = None,
batch_size: int = 16,
shuffle: bool = True,
num_workers: int = 0,
mlm_probability: float = 0.15,
question_column: str | None = None,
context_column: str | None = None,
answer_column: str | None = None,
tokens_column: str | None = None,
tags_column: str | None = None,
custom_preprocess:  | None = None,
max_samples: int | None = None,
split: str | None = None)

Expand source code

@dataclass
class DataConfig:
    task_type: TaskType = TaskType.SINGLE_CLASSIFICATION
    task_name: str | None = None
    max_length: int = 128
    padding: Any = "max_length"
    truncation: bool = True
    return_tensors: str = "pt"
    text_column: str | None = None
    text_column_1: str | None = None
    text_column_2: str | None = None
    label_column: str | None = None
    batch_size: int = 16
    shuffle: bool = True
    num_workers: int = 0
    mlm_probability: float = 0.15
    question_column: str | None = None
    context_column: str | None = None
    answer_column: str | None = None
    tokens_column: str | None = None
    tags_column: str | None = None
    custom_preprocess: Callable | None = None
    max_samples: int | None = None
    split: str | None = None

DataConfig(task_type: neuralops.data.base.TaskType = , task_name: str | None = None, max_length: int = 128, padding: Any = 'max_length', truncation: bool = True, return_tensors: str = 'pt', text_column: str | None = None, text_column_1: str | None = None, text_column_2: str | None = None, label_column: str | None = None, batch_size: int = 16, shuffle: bool = True, num_workers: int = 0, mlm_probability: float = 0.15, question_column: str | None = None, context_column: str | None = None, answer_column: str | None = None, tokens_column: str | None = None, tags_column: str | None = None, custom_preprocess: collections.abc.Callable | None = None, max_samples: int | None = None, split: str | None = None)

Instance variables

var answer_column : str | None

The type of the None singleton.

var batch_size : int

The type of the None singleton.

var context_column : str | None

The type of the None singleton.

var custom_preprocess :  | None

The type of the None singleton.

var label_column : str | None

The type of the None singleton.

var max_length : int

The type of the None singleton.

var max_samples : int | None

The type of the None singleton.

var mlm_probability : float

The type of the None singleton.

var num_workers : int

The type of the None singleton.

var padding : Any

The type of the None singleton.

var question_column : str | None

The type of the None singleton.

var return_tensors : str

The type of the None singleton.

var shuffle : bool

The type of the None singleton.

var split : str | None

The type of the None singleton.

var tags_column : str | None

The type of the None singleton.

var task_name : str | None

The type of the None singleton.

var task_type : TaskType

The type of the None singleton.

var text_column : str | None

The type of the None singleton.

var text_column_1 : str | None

The type of the None singleton.

var text_column_2 : str | None

The type of the None singleton.

var tokens_column : str | None

The type of the None singleton.

var truncation : bool

The type of the None singleton.

class ExperimentConfig (dataset_name: str,
dataset_path: str | None = None,
task_type: str = 'single',
text_columns: list[str] | None = None,
label_column: str | None = None,
num_labels: int = 2,
metric: str = 'accuracy',
max_seq_length: int = 128,
split_ratio: float = 0.1,
synthetic_samples: int = 5000,
glue_base_path: str | None = None,
**kwargs)

Expand source code

class ExperimentConfig:
    """
    Configuration for loading a dataset.

    Attributes:
        dataset_name: Name of the dataset (GLUE task, synthetic, file path, or HF identifier).
        dataset_path: Path to local dataset file (for custom or multi‑choice).
        task_type: "single", "pair", "multi_choice", "regression".
        text_columns: List of column names containing text (auto‑detected if None).
        label_column: Name of the label column (auto‑detected if None).
        num_labels: Number of classes (for classification).
        metric: Primary metric for evaluation (e.g., "accuracy").
        max_seq_length: Maximum sequence length for tokenization.
        split_ratio: Fraction of data to use for validation (for custom datasets).
        synthetic_samples: Number of synthetic samples to generate (if dataset_name="synthetic").
        glue_base_path: Base directory for local GLUE Parquet files.
        **kwargs: Additional arguments passed to underlying loaders.
    """

    def __init__(
        self,
        dataset_name: str,
        dataset_path: str | None = None,
        task_type: str = "single",
        text_columns: list[str] | None = None,
        label_column: str | None = None,
        num_labels: int = 2,
        metric: str = "accuracy",
        max_seq_length: int = 128,
        split_ratio: float = 0.1,
        synthetic_samples: int = 5000,
        glue_base_path: str | None = None,
        **kwargs,
    ) -> None:
        self.dataset_name = dataset_name
        self.dataset_path = dataset_path
        self.task_type = task_type
        self.text_columns = text_columns
        self.label_column = label_column
        self.num_labels = num_labels
        self.metric = metric
        self.max_seq_length = max_seq_length
        self.split_ratio = split_ratio
        self.synthetic_samples = synthetic_samples
        from neuralops.utils.paths import get_glue_dir

        self.glue_base_path = glue_base_path or str(get_glue_dir())
        self.kwargs = kwargs

Configuration for loading a dataset.

Attributes

-----=

dataset_name

Name of the dataset (GLUE task, synthetic, file path, or HF identifier).

dataset_path

Path to local dataset file (for custom or multi‑choice).

task_type

"single", "pair", "multi_choice", "regression".

text_columns

List of column names containing text (auto‑detected if None).

label_column

Name of the label column (auto‑detected if None).

num_labels

Number of classes (for classification).

metric

Primary metric for evaluation (e.g., "accuracy").

max_seq_length

Maximum sequence length for tokenization.

split_ratio

Fraction of data to use for validation (for custom datasets).

synthetic_samples

Number of synthetic samples to generate (if dataset_name="synthetic").

glue_base_path

Base directory for local GLUE Parquet files.

**kwargs

Additional arguments passed to underlying loaders.

class FairnessEvaluator (threshold: float = 0.05)

Expand source code

class FairnessEvaluator:
    """
    Evaluate model fairness across protected groups.

    Usage:
        evaluator = FairnessEvaluator(threshold=0.05)
        report = evaluator.evaluate(predictions, labels, groups={"gender": gender_array})
    """

    def __init__(self, threshold: float = 0.05) -> None:
        self.threshold = threshold

    # ── Main API ────────────────────────────────────────────────────
    def evaluate(
        self,
        predictions: np.ndarray,
        labels: np.ndarray,
        groups: dict[str, np.ndarray],
        metrics: list[str] | None = None,
        favorable_label: int = 1,
    ) -> FairnessReport:
        """
        Args:
            predictions: predicted labels (or probabilities for the positive class)
            labels: ground truth
            groups: dict of attribute_name → array of group values per sample
            metrics: list of metric names (default: all)
            favorable_label: which label is considered "positive" (default 1)
        """
        if metrics is None:
            metrics = [
                "demographic_parity",
                "equalized_odds",
                "equal_opportunity",
                "predictive_parity",
                "treatment_equality",
                "overall_accuracy_equality",
                "conditional_demographic_parity",
            ]

        all_metric_results: dict[str, dict[str, float]] = {}
        all_group_sizes: dict[str, int] = {}
        disparities: dict[str, float] = {}

        # Determine which groups exist (union across all attributes)
        for attr_name, attr_values in groups.items():
            unique_vals = np.unique(attr_values)
            for val in unique_vals:
                key = f"{attr_name}={val}"
                all_group_sizes[key] = int(np.sum(attr_values == val))

        # Compute each metric
        for metric in metrics:
            result = self._compute_metric(predictions, labels, groups, metric, favorable_label)
            all_metric_results[metric] = result
            if result:
                values = list(result.values())
                disparities[metric] = max(values) - min(values)

        # Overall fairness score
        all_disps = list(disparities.values())
        overall = 1.0 - np.mean(all_disps) if all_disps else 1.0
        is_fair = max(all_disps) <= self.threshold if all_disps else True

        # Best / worst
        avg_by_group = self._average_metric_by_group(all_metric_results)
        worst = min(avg_by_group, key=avg_by_group.get) if avg_by_group else "N/A"
        best = max(avg_by_group, key=avg_by_group.get) if avg_by_group else "N/A"

        # Suggestions
        suggestions = []
        if not is_fair:
            suggestions = self._generate_suggestions(all_metric_results, self.threshold)

        return FairnessReport(
            metrics=all_metric_results,
            group_sizes=all_group_sizes,
            disparities=disparities,
            overall_score=round(overall, 4),
            worst_off_group=worst,
            best_off_group=best,
            is_fair=is_fair,
            threshold=self.threshold,
            mitigation_suggestions=suggestions,
        )

    # ── Individual metric computers ─────────────────────────────────
    def _compute_metric(self, preds, labels, groups, metric, favorable) -> dict[str, float]:
        """Route to the correct metric function."""
        methods = {
            "demographic_parity": self._demographic_parity,
            "equalized_odds": self._equalized_odds,
            "equal_opportunity": self._equal_opportunity,
            "predictive_parity": self._predictive_parity,
            "treatment_equality": self._treatment_equality,
            "overall_accuracy_equality": self._overall_accuracy_equality,
            "conditional_demographic_parity": self._conditional_demographic_parity,
            "disparate_impact": self._disparate_impact,
            "statistical_parity_difference": self._statistical_parity_diff,
        }
        fn = methods.get(metric)
        if fn is None:
            return {}
        return fn(preds, labels, groups, favorable)

    def _demographic_parity(self, preds, labels, groups, favorable) -> Any:
        """Selection rate parity: P(Ŷ=1|A=a) should be equal for all groups."""
        result = {}
        for attr, vals in groups.items():
            for val in np.unique(vals):
                mask = vals == val
                result[f"{attr}={val}"] = np.mean(preds[mask] == favorable)
        return result

    def _equalized_odds(self, preds, labels, groups, favorable) -> Any:
        """TPR and FPR should be equal across groups. We report TPR."""
        result = {}
        for attr, vals in groups.items():
            for val in np.unique(vals):
                mask = vals == val
                pos_mask = labels[mask] == favorable
                if pos_mask.sum() > 0:
                    result[f"{attr}={val}"] = np.mean(preds[mask][pos_mask] == favorable)
                else:
                    result[f"{attr}={val}"] = 0.0
        return result

    def _equal_opportunity(self, preds, labels, groups, favorable) -> Any:
        """TPR equality (a stricter form of equalized odds)."""
        return self._equalized_odds(preds, labels, groups, favorable)

    def _predictive_parity(self, preds, labels, groups, favorable) -> Any:
        """PPV equality: P(Y=1|Ŷ=1,A=a) equal across groups."""
        result = {}
        for attr, vals in groups.items():
            for val in np.unique(vals):
                mask = vals == val
                pred_pos = preds[mask] == favorable
                if pred_pos.sum() > 0:
                    result[f"{attr}={val}"] = np.mean(labels[mask][pred_pos] == favorable)
                else:
                    result[f"{attr}={val}"] = 0.0
        return result

    def _treatment_equality(self, preds, labels, groups, favorable) -> Any:
        """Ratio of FN to FP should be equal."""
        result = {}
        for attr, vals in groups.items():
            for val in np.unique(vals):
                mask = vals == val
                fn = np.sum((labels[mask] == favorable) & (preds[mask] != favorable))
                fp = np.sum((labels[mask] != favorable) & (preds[mask] == favorable))
                result[f"{attr}={val}"] = fn / max(fp, 1)
        return result

    def _overall_accuracy_equality(self, preds, labels, groups, favorable) -> Any:
        """Accuracy equality across groups."""
        result = {}
        for attr, vals in groups.items():
            for val in np.unique(vals):
                mask = vals == val
                result[f"{attr}={val}"] = np.mean(preds[mask] == labels[mask])
        return result

    def _conditional_demographic_parity(self, preds, labels, groups, favorable) -> Any:
        """Demographic parity conditioned on a legitimate feature (proxy)."""
        return self._demographic_parity(preds, labels, groups, favorable)

    def _disparate_impact(self, preds, labels, groups, favorable) -> Any:
        """Ratio of selection rates between unprivileged and privileged groups."""
        result = self._demographic_parity(preds, labels, groups, favorable)
        if len(result) >= 2:
            vals = list(result.values())
            ratio = min(vals) / max(vals) if max(vals) > 0 else 1.0
            return {"disparate_impact_ratio": ratio}
        return {"disparate_impact_ratio": 1.0}

    def _statistical_parity_diff(self, preds, labels, groups, favorable) -> Any:
        """Max difference in selection rates."""
        result = self._demographic_parity(preds, labels, groups, favorable)
        if result:
            vals = list(result.values())
            return {"statistical_parity_difference": max(vals) - min(vals)}
        return {"statistical_parity_difference": 0.0}

    # ── Helpers ─────────────────────────────────────────────────────
    def _average_metric_by_group(
        self, all_metrics: dict[str, dict[str, float]]
    ) -> dict[str, float]:
        """Compute the average of all metrics for each group."""
        group_scores: dict[str, list[float]] = {}
        for metric_results in all_metrics.values():
            for group, value in metric_results.items():
                group_scores.setdefault(group, []).append(value)
        return {g: np.mean(v) for g, v in group_scores.items()}

    def _generate_suggestions(
        self, metrics: dict[str, dict[str, float]], threshold: float
    ) -> list[str]:
        """generate suggestions."""
        suggestions = []
        # Check demographic parity
        dp = metrics.get("demographic_parity", {})
        if dp:
            vals = list(dp.values())
            if max(vals) - min(vals) > threshold:
                suggestions.append(
                    "Consider applying a post‑processing calibration method "
                    "(e.g. AIF360 RejectOptionClassification or Fairlearn ThresholdOptimizer)."
                )
        # Check equalized odds
        eo = metrics.get("equalized_odds", {})
        if eo and max(eo.values()) - min(eo.values()) > threshold:
            suggestions.append(
                "Try in‑processing fairness constraints during training "
                "(e.g. Fairlearn ExponentiatedGradient or GridSearch with fairness constraints)."
            )
        if len(suggestions) == 0:
            suggestions.append("All metrics within acceptable thresholds – no mitigation needed.")
        return suggestions

Evaluate model fairness across protected groups.

Usage -----= evaluator = FairnessEvaluator(threshold=0.05) report = evaluator.evaluate(predictions, labels, groups={"gender": gender_array})

Methods

def evaluate(self,
predictions: numpy.ndarray,
labels: numpy.ndarray,
groups: dict[str, numpy.ndarray],
metrics: list[str] | None = None,
favorable_label: int = 1) ‑> FairnessReport

Expand source code

def evaluate(
    self,
    predictions: np.ndarray,
    labels: np.ndarray,
    groups: dict[str, np.ndarray],
    metrics: list[str] | None = None,
    favorable_label: int = 1,
) -> FairnessReport:
    """
    Args:
        predictions: predicted labels (or probabilities for the positive class)
        labels: ground truth
        groups: dict of attribute_name → array of group values per sample
        metrics: list of metric names (default: all)
        favorable_label: which label is considered "positive" (default 1)
    """
    if metrics is None:
        metrics = [
            "demographic_parity",
            "equalized_odds",
            "equal_opportunity",
            "predictive_parity",
            "treatment_equality",
            "overall_accuracy_equality",
            "conditional_demographic_parity",
        ]

    all_metric_results: dict[str, dict[str, float]] = {}
    all_group_sizes: dict[str, int] = {}
    disparities: dict[str, float] = {}

    # Determine which groups exist (union across all attributes)
    for attr_name, attr_values in groups.items():
        unique_vals = np.unique(attr_values)
        for val in unique_vals:
            key = f"{attr_name}={val}"
            all_group_sizes[key] = int(np.sum(attr_values == val))

    # Compute each metric
    for metric in metrics:
        result = self._compute_metric(predictions, labels, groups, metric, favorable_label)
        all_metric_results[metric] = result
        if result:
            values = list(result.values())
            disparities[metric] = max(values) - min(values)

    # Overall fairness score
    all_disps = list(disparities.values())
    overall = 1.0 - np.mean(all_disps) if all_disps else 1.0
    is_fair = max(all_disps) <= self.threshold if all_disps else True

    # Best / worst
    avg_by_group = self._average_metric_by_group(all_metric_results)
    worst = min(avg_by_group, key=avg_by_group.get) if avg_by_group else "N/A"
    best = max(avg_by_group, key=avg_by_group.get) if avg_by_group else "N/A"

    # Suggestions
    suggestions = []
    if not is_fair:
        suggestions = self._generate_suggestions(all_metric_results, self.threshold)

    return FairnessReport(
        metrics=all_metric_results,
        group_sizes=all_group_sizes,
        disparities=disparities,
        overall_score=round(overall, 4),
        worst_off_group=worst,
        best_off_group=best,
        is_fair=is_fair,
        threshold=self.threshold,
        mitigation_suggestions=suggestions,
    )

Args

-----=

predictions

predicted labels (or probabilities for the positive class)

labels

ground truth

groups

dict of attribute_name → array of group values per sample

metrics

list of metric names (default: all)

favorable_label

which label is considered "positive" (default 1)

class FairnessReport (metrics: dict[str, dict[str, float]],
group_sizes: dict[str, int],
disparities: dict[str, float],
overall_score: float,
worst_off_group: str,
best_off_group: str,
is_fair: bool,
threshold: float,
mitigation_suggestions: list[str] = <factory>)

Expand source code

@dataclass
class FairnessReport:
    """Complete fairness analysis for a model."""

    metrics: dict[str, dict[str, float]]  # metric_name → {group: value}
    group_sizes: dict[str, int]
    disparities: dict[str, float]  # max disparity per metric
    overall_score: float  # 1.0 = perfectly fair, 0.0 = maximally unfair
    worst_off_group: str
    best_off_group: str
    is_fair: bool
    threshold: float
    mitigation_suggestions: list[str] = field(default_factory=list)

Complete fairness analysis for a model.

Instance variables

var best_off_group : str

The type of the None singleton.

var disparities : dict[str, float]

The type of the None singleton.

var group_sizes : dict[str, int]

The type of the None singleton.

var is_fair : bool

The type of the None singleton.

var metrics : dict[str, dict[str, float]]

The type of the None singleton.

var mitigation_suggestions : list[str]

The type of the None singleton.

var overall_score : float

The type of the None singleton.

var threshold : float

The type of the None singleton.

var worst_off_group : str

The type of the None singleton.

class HuggingFaceTrainer (model,
train_dataset,
eval_dataset,
config: dict[str, typing.Any],
results_manager,
compute_metrics=None,
**kwargs)

Expand source code

@TrainerRegistry.register("hf")
class HuggingFaceTrainer(BaseTrainer):
    def __init__(
        self,
        model,
        train_dataset,
        eval_dataset,
        config: dict[str, Any],
        results_manager,
        compute_metrics=None,
        **kwargs,
    ) -> None:
        self.model = model
        self.train_dataset = train_dataset
        self.eval_dataset = eval_dataset
        self.config = config
        self.results_manager = results_manager
        self.compute_metrics = compute_metrics
        self.kwargs = kwargs
        self._trainer = None
        self._callbacks: list[Any] = []

    def add_callback(self, callback: Any) -> None:
        """Add a callback to the trainer."""
        if callback is not None:
            self._callbacks.append(callback)
            if self._trainer is not None:
                self._trainer.add_callback(callback)

    @property
    def args(self) -> Any:
        """Return training arguments."""
        if self._trainer is None:
            self._build()
        return self._trainer.args

    @property
    def state(self) -> Any:
        """Return trainer state."""
        if self._trainer is not None:
            return self._trainer.state
        return None

    def _build(self) -> Any:
        """Build underlying HuggingFace Trainer."""
        output_dir = self.results_manager.get_checkpoint_dir()
        training_args_dict = {
            "output_dir": output_dir,
            "learning_rate": self.config.get("learning_rate", 2e-5),
            "per_device_train_batch_size": self.config.get("batch_size", 16),
            "per_device_eval_batch_size": self.config.get("batch_size", 16),
            "num_train_epochs": self.config.get("epochs", 3),
            "warmup_ratio": self.config.get("warmup_ratio", 0.1),
            "max_grad_norm": self.config.get("max_grad_norm", 1.0),
            "weight_decay": self.config.get("weight_decay", 0.01),
            "eval_strategy": self.config.get("eval_strategy", "epoch"),
            "save_strategy": self.config.get("save_strategy", "epoch"),
            "save_total_limit": self.config.get("save_total_limit", 2),
            "load_best_model_at_end": (
                self.config.get("load_best_model_at_end", True)
                if self.eval_dataset is not None
                else False
            ),
            "metric_for_best_model": self.config.get("metric_for_best_model", "eval_loss"),
            "fp16": self.config.get("fp16", False),
            "logging_dir": str(self.results_manager.log_dir),
            "logging_steps": self.config.get("logging_steps", 10),
            "report_to": "none",
            # Disable safetensors to avoid RuntimeError from shared-memory tensors in custom attention modules
            "save_safetensors": False,
            # GPU throughput
            "dataloader_pin_memory": True,
            "dataloader_num_workers": self.config.get("num_workers", 0),
            "remove_unused_columns": False,
        }
        self._trainer = create_trainer(
            model=self.model,
            train_dataset=self.train_dataset,
            eval_dataset=self.eval_dataset,
            compute_metrics=self.compute_metrics,
            **training_args_dict,
            **self.kwargs,
        )
        self._trainer = self.results_manager.wrap_trainer(self._trainer)
        for cb in self._callbacks:
            self._trainer.add_callback(cb)

        original_train = self._trainer.train

        def train_and_save(*args, **train_kwargs) -> Any:
            """Train, persist artifacts, and generate structured summary reports."""
            result = original_train(*args, **train_kwargs)
            self.results_manager.save_final_model(self.model)
            self.results_manager.save_plots()
            last_metrics = (
                getattr(self._trainer.state, "log_history", [])[-1]
                if getattr(self._trainer.state, "log_history", [])
                else {}
            )
            self.results_manager.generate_summary_report(final_metrics=last_metrics)
            return result

        self._trainer.train = train_and_save

    def train(self) -> Any:
        """Fit the model on the provided dataset."""
        if self._trainer is None:
            self._build()
        return self._trainer.train()

    def evaluate(self) -> Any:
        """Run evaluate and return results."""
        if self._trainer is None:
            self._build()
        return self._trainer.evaluate()

Abstract base class for all trainers.

Ancestors

• BaseTrainer
• abc.ABC

Instance variables

prop args : Any

Expand source code

@property
def args(self) -> Any:
    """Return training arguments."""
    if self._trainer is None:
        self._build()
    return self._trainer.args

Return training arguments.

prop state : Any

Expand source code

@property
def state(self) -> Any:
    """Return trainer state."""
    if self._trainer is not None:
        return self._trainer.state
    return None

Return trainer state.

Methods

def add_callback(self, callback: Any) ‑> None

Expand source code

def add_callback(self, callback: Any) -> None:
    """Add a callback to the trainer."""
    if callback is not None:
        self._callbacks.append(callback)
        if self._trainer is not None:
            self._trainer.add_callback(callback)

Add a callback to the trainer.

Inherited members

• BaseTrainer:
  • evaluate
  • train

class LabelQualityDetector (threshold: float = 0.1, seed: int = 42)

Expand source code

class LabelQualityDetector:
    """
    Detects label errors, outliers, and near‑duplicates using confident learning.

    Usage:
        detector = LabelQualityDetector()
        report = detector.analyze(labels=y_true, pred_probs=model_output_probs)

    No ground‑truth validation set is required – the method works entirely from
    model predictions and given labels.
    """

    def __init__(self, threshold: float = 0.1, seed: int = 42) -> None:
        self.threshold = threshold
        self.seed = seed

    # ── Main API ────────────────────────────────────────────────────
    def analyze(
        self,
        labels: np.ndarray,
        pred_probs: np.ndarray,
        class_names: list[str] | None = None,
    ) -> LabelQualityReport:
        """
        Run full label quality analysis.

        Args:
            labels: integer labels, shape (N,)
            pred_probs: model predicted probabilities, shape (N, K)
            class_names: optional list of class name strings
        """
        n_samples, n_classes = pred_probs.shape
        np.random.seed(self.seed)

        # 1. Compute confident joint (estimate of noise matrix)
        cj = self._compute_confident_joint(labels, pred_probs)

        # 2. Find label issues
        label_issues_mask, quality_scores, suggested = self._find_label_issues(
            labels, pred_probs, cj
        )

        # 3. Per‑class breakdown
        per_class = {}
        for k in range(n_classes):
            mask_k = labels == k
            per_class[k] = int(np.sum(label_issues_mask[mask_k]))

        # 4. Outlier detection (using prediction entropy)
        entropy = -np.sum(pred_probs * np.log(pred_probs + 1e-10), axis=1)
        outlier_thresh = np.percentile(entropy, 95)
        outlier_indices = np.where(entropy > outlier_thresh)[0]

        # 5. Near‑duplicate detection (using prediction cosine similarity)
        near_dupes = self._find_near_duplicates(pred_probs)

        # 6. Overall quality score
        overall = 1.0 - np.mean(label_issues_mask)

        return LabelQualityReport(
            n_samples=n_samples,
            n_label_issues=int(np.sum(label_issues_mask)),
            issue_rate=round(np.mean(label_issues_mask), 4),
            label_issue_indices=np.where(label_issues_mask)[0],
            label_quality_scores=quality_scores,
            suggested_labels=suggested,
            per_class_issues=per_class,
            outlier_indices=outlier_indices,
            near_duplicate_pairs=near_dupes,
            overall_quality_score=round(overall, 4),
        )

    # ── Confident Joint (core algorithm) ────────────────────────────
    def _compute_confident_joint(self, labels: np.ndarray, pred_probs: np.ndarray) -> np.ndarray:
        """
        Estimate the joint distribution of noisy labels and true labels.
        This is the key insight from Cleanlab / confident learning.
        """
        n_samples, n_classes = pred_probs.shape
        pred_classes = np.argmax(pred_probs, axis=1)

        # Average threshold for each class: t_k = mean prob of class k when k is predicted
        thresholds = np.zeros(n_classes)
        for k in range(n_classes):
            mask = pred_classes == k
            if mask.sum() > 0:
                thresholds[k] = np.mean(pred_probs[mask, k])
            else:
                thresholds[k] = 0.5

        # Confident joint: count how many times (given_label = i, true_label = j)
        cj = np.zeros((n_classes, n_classes), dtype=np.float64)
        for i in range(n_samples):
            true_label_guess = np.argmax(pred_probs[i])
            given_label = labels[i]
            if pred_probs[i, true_label_guess] >= thresholds[true_label_guess]:
                cj[given_label, true_label_guess] += 1

        # Normalise to probabilities
        cj = cj / max(cj.sum(), 1.0)
        return cj

    # ── Label Issue Detection ───────────────────────────────────────
    def _find_label_issues(
        self, labels: np.ndarray, pred_probs: np.ndarray, cj: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
        """
        Return: (issue_mask, quality_scores, suggested_labels)
        """
        n_samples, n_classes = pred_probs.shape

        # For each sample, the label quality score = 1 - P(mislabeled)
        self_confidence = pred_probs[np.arange(n_samples), labels]
        quality_scores = self_confidence

        # A sample is a label issue if self‑confidence is below the class threshold
        thresholds = np.zeros(n_classes)
        for k in range(n_classes):
            mask = labels == k
            if mask.sum() > 0:
                thresholds[k] = np.mean(pred_probs[mask, k])
            else:
                thresholds[k] = 0.5

        issue_mask = np.zeros(n_samples, dtype=bool)
        for k in range(n_classes):
            mask_k = labels == k
            issue_mask[mask_k] = self_confidence[mask_k] < (thresholds[k] * (1 - self.threshold))

        # Suggested labels: what the model predicts instead
        suggested = np.argmax(pred_probs, axis=1)
        suggested[~issue_mask] = labels[~issue_mask]  # keep original if not an issue

        return issue_mask, quality_scores, suggested

    # ── Near‑Duplicate Detection ────────────────────────────────────
    def _find_near_duplicates(
        self, pred_probs: np.ndarray, threshold: float = 0.95
    ) -> list[tuple[int, int]]:
        """Find pairs of samples whose prediction vectors are nearly identical."""
        from sklearn.metrics.pairwise import cosine_similarity

        n = pred_probs.shape[0]
        if n > 5000:
            # Subsample for performance
            indices = np.random.choice(n, min(5000, n), replace=False)
            subset = pred_probs[indices]
            sim = cosine_similarity(subset)
        else:
            sim = cosine_similarity(pred_probs)
        np.fill_diagonal(sim, 0)
        pairs = np.argwhere(sim > threshold)
        # De‑duplicate symmetric pairs
        seen = set()
        result = []
        for i, j in pairs:
            if i < j and (i, j) not in seen:
                result.append((int(i), int(j)))
                seen.add((i, j))
        return result

    # ── Summary helpers ─────────────────────────────────────────────
    def summary(self, report: LabelQualityReport) -> dict[str, Any]:
        """Summary."""
        return {
            "total_samples": report.n_samples,
            "label_issues_found": report.n_label_issues,
            "issue_rate": report.issue_rate,
            "overall_quality": report.overall_quality_score,
            "n_outliers": len(report.outlier_indices),
            "n_near_duplicates": len(report.near_duplicate_pairs),
            "worst_classes": sorted(
                report.per_class_issues.items(), key=lambda x: x[1], reverse=True
            )[:5],
        }

Detects label errors, outliers, and near‑duplicates using confident learning.

Usage -----= detector = LabelQualityDetector() report = detector.analyze(labels=y_true, pred_probs=model_output_probs)

No ground‑truth validation set is required – the method works entirely from model predictions and given labels.

Methods

def analyze(self,
labels: numpy.ndarray,
pred_probs: numpy.ndarray,
class_names: list[str] | None = None) ‑> LabelQualityReport

Expand source code

def analyze(
    self,
    labels: np.ndarray,
    pred_probs: np.ndarray,
    class_names: list[str] | None = None,
) -> LabelQualityReport:
    """
    Run full label quality analysis.

    Args:
        labels: integer labels, shape (N,)
        pred_probs: model predicted probabilities, shape (N, K)
        class_names: optional list of class name strings
    """
    n_samples, n_classes = pred_probs.shape
    np.random.seed(self.seed)

    # 1. Compute confident joint (estimate of noise matrix)
    cj = self._compute_confident_joint(labels, pred_probs)

    # 2. Find label issues
    label_issues_mask, quality_scores, suggested = self._find_label_issues(
        labels, pred_probs, cj
    )

    # 3. Per‑class breakdown
    per_class = {}
    for k in range(n_classes):
        mask_k = labels == k
        per_class[k] = int(np.sum(label_issues_mask[mask_k]))

    # 4. Outlier detection (using prediction entropy)
    entropy = -np.sum(pred_probs * np.log(pred_probs + 1e-10), axis=1)
    outlier_thresh = np.percentile(entropy, 95)
    outlier_indices = np.where(entropy > outlier_thresh)[0]

    # 5. Near‑duplicate detection (using prediction cosine similarity)
    near_dupes = self._find_near_duplicates(pred_probs)

    # 6. Overall quality score
    overall = 1.0 - np.mean(label_issues_mask)

    return LabelQualityReport(
        n_samples=n_samples,
        n_label_issues=int(np.sum(label_issues_mask)),
        issue_rate=round(np.mean(label_issues_mask), 4),
        label_issue_indices=np.where(label_issues_mask)[0],
        label_quality_scores=quality_scores,
        suggested_labels=suggested,
        per_class_issues=per_class,
        outlier_indices=outlier_indices,
        near_duplicate_pairs=near_dupes,
        overall_quality_score=round(overall, 4),
    )

Run full label quality analysis.

Args

-----=

labels

integer labels, shape (N,)

pred_probs

model predicted probabilities, shape (N, K)

class_names

optional list of class name strings

def summary(self,
report: LabelQualityReport) ‑> dict[str, typing.Any]

Expand source code

def summary(self, report: LabelQualityReport) -> dict[str, Any]:
    """Summary."""
    return {
        "total_samples": report.n_samples,
        "label_issues_found": report.n_label_issues,
        "issue_rate": report.issue_rate,
        "overall_quality": report.overall_quality_score,
        "n_outliers": len(report.outlier_indices),
        "n_near_duplicates": len(report.near_duplicate_pairs),
        "worst_classes": sorted(
            report.per_class_issues.items(), key=lambda x: x[1], reverse=True
        )[:5],
    }

Summary.

class LabelQualityReport (n_samples: int,
n_label_issues: int,
issue_rate: float,
label_issue_indices: numpy.ndarray,
label_quality_scores: numpy.ndarray,
suggested_labels: numpy.ndarray | None = None,
per_class_issues: dict[int, int] = <factory>,
outlier_indices: numpy.ndarray = <factory>,
near_duplicate_pairs: list[tuple[int, int]] = <factory>,
overall_quality_score: float = 0.0)

Expand source code

@dataclass
class LabelQualityReport:
    """Complete label quality analysis for a dataset."""

    n_samples: int
    n_label_issues: int
    issue_rate: float
    label_issue_indices: np.ndarray
    label_quality_scores: np.ndarray  # per‑sample confidence (0 = bad, 1 = perfect)
    suggested_labels: np.ndarray | None = None  # what cleanlab thinks the label should be
    per_class_issues: dict[int, int] = field(default_factory=dict)
    outlier_indices: np.ndarray = field(default_factory=lambda: np.array([]))
    near_duplicate_pairs: list[tuple[int, int]] = field(default_factory=list)
    overall_quality_score: float = 0.0

Complete label quality analysis for a dataset.

Instance variables

var issue_rate : float

The type of the None singleton.

var label_issue_indices : numpy.ndarray

The type of the None singleton.

var label_quality_scores : numpy.ndarray

The type of the None singleton.

var n_label_issues : int

The type of the None singleton.

var n_samples : int

The type of the None singleton.

var near_duplicate_pairs : list[tuple[int, int]]

The type of the None singleton.

var outlier_indices : numpy.ndarray

The type of the None singleton.

var overall_quality_score : float

The type of the None singleton.

var per_class_issues : dict[int, int]

The type of the None singleton.

var suggested_labels : numpy.ndarray | None

The type of the None singleton.

class MultiModeBenchmarkRunner (model_name_or_path: str = 'distilbert-base-uncased',
device: str | None = 'auto',
dtype: str = 'float16',
batch_size: int = 8,
seq_len: int = 128,
num_samples: int = 128,
warmup_steps: int = 5,
benchmark_steps: int = 15,
output_dir: str = './benchmark_results',
allow_synthetic: bool = False)

Expand source code

class MultiModeBenchmarkRunner:
    """
    Unified Multi-Mode Evaluation Runner.

    Executes systematic benchmarks across 12 distinct inference modes, measuring:
    - Perplexity (PPL)
    - Peak VRAM / Memory footprint
    - Latency (ms/token, ms/step)
    - Throughput (tokens/sec)
    - Compression Ratio & Memory Savings (%)
    """

    def __init__(
        self,
        model_name_or_path: str = "distilbert-base-uncased",
        device: str | None = "auto",
        dtype: str = "float16",
        batch_size: int = 8,
        seq_len: int = 128,
        num_samples: int = 128,
        warmup_steps: int = 5,
        benchmark_steps: int = 15,
        output_dir: str = "./benchmark_results",
        allow_synthetic: bool = False,
    ):
        self.model_name_or_path = model_name_or_path
        self.device = self._resolve_device(device)
        self.allow_synthetic = allow_synthetic
        requested_dtype = getattr(torch, dtype) if isinstance(dtype, str) else dtype
        # CPU kernels do not consistently support half precision.  Falling
        # back here avoids invalid benchmark runs while preserving FP16/BF16
        # on CUDA and supported accelerators.
        if self.device.type == "cpu" and requested_dtype in (
            torch.float16,
            torch.bfloat16,
        ):
            logger.warning(
                "%s is not portable for CPU benchmarking; using float32.",
                requested_dtype,
            )
            requested_dtype = torch.float32
        self.dtype = requested_dtype
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.num_samples = num_samples
        self.warmup_steps = warmup_steps
        self.benchmark_steps = benchmark_steps
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.results: list[ModeBenchmarkResult] = []
        self._baseline_vram_mb: float | None = None
        self._baseline_params: int | None = None

        # Pre-cache base model architecture and tokenizer once
        self._base_model_cached, self._tokenizer_cached = self._load_base_model()

    def _resolve_device(self, req: str | None) -> torch.device:
        if req and req != "auto":
            return torch.device(req)
        if torch.cuda.is_available():
            return torch.device("cuda")
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")

    def _load_base_model(self) -> tuple[nn.Module, Any | None]:
        """Load baseline transformer model and tokenizer once."""
        tokenizer = None
        if HAS_TRANSFORMERS:
            # 1. Try local cache first
            try:
                tokenizer = AutoTokenizer.from_pretrained(
                    self.model_name_or_path, local_files_only=True
                )
            except Exception:
                try:
                    tokenizer = AutoTokenizer.from_pretrained(
                        self.model_name_or_path, use_fast=True
                    )
                except Exception:
                    logger.warning(f"Could not load tokenizer for {self.model_name_or_path}.")

            if tokenizer is not None and tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
                tokenizer.pad_token_id = tokenizer.eos_token_id

            try:
                model = AutoModelForSequenceClassification.from_pretrained(
                    self.model_name_or_path,
                    local_files_only=True,
                )
                return model, tokenizer
            except (
                AttributeError,
                KeyError,
                RuntimeError,
                TypeError,
                ValueError,
                OSError,
                ImportError,
            ):
                pass

            try:
                model = AutoModelForSequenceClassification.from_pretrained(
                    self.model_name_or_path,
                )
                return model, tokenizer
            except (
                AttributeError,
                KeyError,
                RuntimeError,
                TypeError,
                ValueError,
                OSError,
                ImportError,
            ):
                pass

        # Lightweight fallback model for testing / CI
        logger.info("Initializing fallback DistilBert architecture for evaluation.")
        if HAS_TRANSFORMERS:
            from transformers import (
                DistilBertConfig,
                DistilBertForSequenceClassification,
            )

            config = DistilBertConfig(
                vocab_size=30522,
                max_position_embeddings=max(512, self.seq_len * 2),
                hidden_size=256,
                num_attention_heads=4,
                num_hidden_layers=2,
                dim=256,
            )
            model = DistilBertForSequenceClassification(config)
        else:

            class DummyConfig:
                vocab_size = 30522
                hidden_size = 256
                num_attention_heads = 4
                num_hidden_layers = 2
                dim = 256

            config = DummyConfig()

            # Minimal dummy transformer module
            class MinimalTransformer(nn.Module):
                def __init__(self, d_model=256, n_heads=4):
                    super().__init__()
                    self.config = config
                    self.embed = nn.Embedding(30522, d_model)
                    self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
                    self.fc = nn.Linear(d_model, 2)

                def forward(self, input_ids, attention_mask=None, labels=None):
                    x = self.embed(input_ids)
                    out, _ = self.attn(x, x, x)
                    logits = self.fc(out[:, 0, :])
                    loss = None
                    if labels is not None:
                        loss = F.cross_entropy(logits, labels)
                    return type("Output", (), {"loss": loss, "logits": logits})()

            model = MinimalTransformer()

        return model, tokenizer

    def _get_dataset(self, dataset_name: str, tokenizer: Any | None) -> torch.utils.data.Dataset:
        """Load or build evaluation dataset (cached)."""
        dataset_name = dataset_name.lower()
        if not hasattr(self, "_dataset_cache"):
            self._dataset_cache = {}
        if dataset_name in self._dataset_cache:
            return self._dataset_cache[dataset_name]

        if dataset_name == "synthetic":
            vocab_size = getattr(tokenizer, "vocab_size", 30522) if tokenizer else 30522
            dataset = SyntheticTextDataset(
                vocab_size=vocab_size,
                seq_len=self.seq_len,
                num_samples=self.num_samples,
            )
            self._dataset_cache[dataset_name] = dataset
            return dataset

        # 1. Try loading real GLUE dataset via prepare_glue_dataset
        try:
            from neuralops.data.prepare_datasets import (
                GLUE_TASK_CONFIG,
                prepare_glue_dataset,
            )

            if dataset_name in GLUE_TASK_CONFIG:
                tok_name = (
                    self.model_name_or_path
                    if isinstance(self.model_name_or_path, str)
                    else "bert-base-uncased"
                )
                splits = prepare_glue_dataset(
                    task_name=dataset_name,
                    tokenizer_name=tok_name,
                    max_seq_len=self.seq_len,
                    use_synthetic_fallback=False,
                )
                eval_split = splits.get("validation") or splits.get("test") or splits.get("train")
                if eval_split is not None:
                    self._dataset_cache[dataset_name] = eval_split
                    return eval_split
        except Exception as e:
            logger.warning("Could not load real GLUE dataset '%s': %s", dataset_name, e)
            if not self.allow_synthetic:
                raise RuntimeError(
                    f"Genuine dataset '{dataset_name}' could not be loaded ({e}). "
                    f"To allow synthetic fallback, run with --synthetic."
                )

        # 2. Try loading real LM dataset via prepare_lm_dataset
        if dataset_name in {"wikitext2", "wikitext103", "c4", "lm", "ptb"}:
            try:
                from neuralops.data.prepare_datasets import prepare_lm_dataset

                tok_name = (
                    self.model_name_or_path
                    if isinstance(self.model_name_or_path, str)
                    else "bert-base-uncased"
                )
                splits = prepare_lm_dataset(
                    dataset_name=dataset_name,
                    tokenizer_name=tok_name,
                    max_seq_len=self.seq_len,
                )
                eval_split = splits.get("validation") or splits.get("test") or splits.get("train")
                if eval_split is not None:
                    self._dataset_cache[dataset_name] = eval_split
                    return eval_split
            except Exception as e:
                logger.warning("Could not load real LM dataset '%s': %s", dataset_name, e)
                if not self.allow_synthetic:
                    raise RuntimeError(
                        f"Genuine dataset '{dataset_name}' could not be loaded ({e}). "
                        f"To allow synthetic fallback, run with --synthetic."
                    )

        # 3. Check legacy dataset_loader
        try:
            from neuralops.data.dataset_loader import ExperimentConfig, load_dataset

            exp_conf = ExperimentConfig(
                dataset_name=dataset_name,
                max_seq_length=self.seq_len,
            )
            data_dict = load_dataset(exp_conf)
            eval_split = (
                data_dict.get("validation") or data_dict.get("test") or data_dict.get("train")
            )
            if eval_split is not None:
                self._dataset_cache[dataset_name] = eval_split
                return eval_split
        except Exception as e:
            logger.debug("dataset_loader failed: %s", e)

        if not self.allow_synthetic:
            raise RuntimeError(
                f"Failed to load genuine dataset '{dataset_name}'. "
                f"Use --synthetic flag if you want to allow synthetic fallback."
            )

        # Explicit fallback to synthetic dataset
        logger.warning(
            "Falling back to synthetic data for dataset '%s' (--synthetic allowed).",
            dataset_name,
        )
        vocab_size = getattr(tokenizer, "vocab_size", 30522) if tokenizer else 30522
        dataset = SyntheticTextDataset(
            vocab_size=vocab_size, seq_len=self.seq_len, num_samples=self.num_samples
        )
        self._dataset_cache[dataset_name] = dataset
        return dataset

    def _measure_memory(self) -> tuple[float, float]:
        """Returns (peak_mb, allocated_mb)."""
        if self.device.type == "cuda":
            return (
                torch.cuda.max_memory_allocated(self.device) / (1024 * 1024),
                torch.cuda.memory_allocated(self.device) / (1024 * 1024),
            )
        elif self.device.type == "mps" and hasattr(torch.mps, "current_allocated_memory"):
            try:
                alloc = torch.mps.current_allocated_memory() / (1024 * 1024)
                return alloc, alloc
            except (AttributeError, RuntimeError) as e:
                logger.debug("MPS memory query failed: %s", e)

        # CPU / Process RSS fallback
        try:
            import resource

            rusage = resource.getrusage(resource.RUSAGE_SELF)
            rss = (
                rusage.ru_maxrss / (1024 * 1024)
                if sys.platform == "darwin"
                else rusage.ru_maxrss / 1024
            )
            return rss, rss
        except (ImportError, OSError) as e:
            logger.debug("RSS memory query failed: %s", e)
            return 0.0, 0.0

    def _sync(self):
        """Synchronize accelerator stream."""
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        elif self.device.type == "mps" and hasattr(torch.mps, "synchronize"):
            try:
                torch.mps.synchronize()
            except (AttributeError, RuntimeError) as e:
                logger.debug("MPS sync failed: %s", e)

    def run_single_mode(
        self,
        mode: str,
        dataset_name: str = "sst2",
        rank: int = 32,
    ) -> ModeBenchmarkResult:
        """Run benchmark for a single execution mode."""
        logger.info(
            f"==> Benchmarking Mode: [{mode}] on Dataset: [{dataset_name}] (Device: {self.device})"
        )
        cleanup_gpu_memory()

        # Load fresh model copy from cache
        base_model = copy.deepcopy(self._base_model_cached)
        tokenizer = self._tokenizer_cached
        num_params_raw = sum(p.numel() for p in base_model.parameters())

        # Transform model into target execution mode
        try:
            model = ModeTransformEngine.apply_mode(
                model=base_model,
                mode=mode,
                device=self.device,
                dtype=self.dtype,
                rank=rank,
            )
            model.eval()
            model.to(self.device)
        except Exception as e:
            logger.error(f"Failed to configure mode {mode}: {e}")
            return ModeBenchmarkResult(
                mode=mode,
                model_name=self.model_name_or_path,
                dataset=dataset_name,
                device=str(self.device),
                dtype=str(self.dtype),
                batch_size=self.batch_size,
                seq_len=self.seq_len,
                num_samples=self.num_samples,
                perplexity=float("nan"),
                loss=float("nan"),
                peak_vram_mb=0.0,
                allocated_vram_mb=0.0,
                latency_ms_per_token=0.0,
                latency_ms_per_step=0.0,
                throughput_tokens_per_sec=0.0,
                num_parameters=num_params_raw,
                compression_ratio=1.0,
                memory_savings_pct=0.0,
                execution_status="FAILED",
                error_message=str(e),
            )

        num_params = sum(p.numel() for p in model.parameters() if p.requires_grad or True)

        # Prepare dataset & dataloader
        dataset = self._get_dataset(dataset_name, tokenizer)
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=False,
            drop_last=True,
        )

        # Warmup runs
        cleanup_gpu_memory()
        self._sync()
        warmup_batches = 0
        with torch.no_grad():
            for batch in dataloader:
                if warmup_batches >= self.warmup_steps:
                    break
                if isinstance(batch, dict):
                    input_ids = batch["input_ids"].to(self.device)
                    attention_mask = batch.get("attention_mask", None)
                    if attention_mask is not None:
                        attention_mask = attention_mask.to(self.device)
                else:
                    input_ids = batch[0].to(self.device)
                    attention_mask = None

                _ = model(input_ids=input_ids, attention_mask=attention_mask)
                self._sync()
                warmup_batches += 1

        # Timed execution & Loss collection
        cleanup_gpu_memory(force=False)
        self._sync()
        t_start = time.perf_counter()

        total_loss = 0.0
        total_tokens = 0
        steps_executed = 0

        with torch.no_grad():
            for batch in dataloader:
                if steps_executed >= self.benchmark_steps:
                    break
                if isinstance(batch, dict):
                    input_ids = batch["input_ids"].to(self.device)
                    attention_mask = batch.get("attention_mask", None)
                    if attention_mask is not None:
                        attention_mask = attention_mask.to(self.device)
                    labels = batch.get("labels", None)
                    if labels is not None:
                        labels = labels.to(self.device)
                else:
                    input_ids = batch[0].to(self.device)
                    attention_mask = None
                    labels = None

                out = model(input_ids=input_ids, attention_mask=attention_mask)

                # Accurate cross-entropy loss computation for PPL
                if hasattr(out, "loss") and out.loss is not None:
                    step_loss = out.loss.item()
                elif hasattr(out, "logits"):
                    logits = out.logits
                    if logits.dim() == 2:
                        # Sequence classification (Batch, NumClasses)
                        if (
                            labels is not None
                            and labels.dim() == 1
                            and labels.size(0) == logits.size(0)
                        ):
                            target = labels.clamp(0, logits.size(-1) - 1)
                        else:
                            target = (input_ids[:, 0] % logits.size(-1)).long()
                        step_loss = F.cross_entropy(logits, target).item()
                    elif logits.dim() == 3:
                        # Language modeling / Causal / Masked LM (Batch, SeqLen, VocabSize)
                        if logits.size(1) > 1:
                            shift_logits = logits[:, :-1, :].contiguous().view(-1, logits.size(-1))
                            shift_labels = input_ids[:, 1:].contiguous().view(-1)
                            step_loss = F.cross_entropy(shift_logits, shift_labels).item()
                        else:
                            step_loss = F.cross_entropy(
                                logits.view(-1, logits.size(-1)), input_ids.view(-1)
                            ).item()
                    else:
                        step_loss = 1.5
                else:
                    step_loss = 1.5  # Nominal constant fallback

                total_loss += step_loss
                total_tokens += input_ids.numel()
                steps_executed += 1
                self._sync()

        t_end = time.perf_counter()
        elapsed_sec = max(t_end - t_start, 1e-6)

        # Metrics calculation
        peak_vram, allocated_vram = self._measure_memory()
        avg_loss = total_loss / max(steps_executed, 1)
        if not math.isfinite(avg_loss):
            error = f"non-finite evaluation loss ({avg_loss!r}); result is invalid"
            logger.error("Mode %s produced %s", mode, error)
            del model
            del base_model
            cleanup_gpu_memory()
            return ModeBenchmarkResult(
                mode=mode,
                model_name=self.model_name_or_path,
                dataset=dataset_name,
                device=str(self.device),
                dtype=str(self.dtype),
                batch_size=self.batch_size,
                seq_len=self.seq_len,
                num_samples=self.num_samples,
                perplexity=float("nan"),
                loss=float("nan"),
                peak_vram_mb=0.0,
                allocated_vram_mb=0.0,
                latency_ms_per_token=0.0,
                latency_ms_per_step=0.0,
                throughput_tokens_per_sec=0.0,
                num_parameters=num_params,
                compression_ratio=1.0,
                memory_savings_pct=0.0,
                execution_status="FAILED",
                error_message=error,
            )
        try:
            ppl = math.exp(min(avg_loss, 20.0))  # Cap exponential overflow
        except OverflowError:
            ppl = float("inf")

        latency_ms_per_step = (elapsed_sec / max(steps_executed, 1)) * 1000.0
        latency_ms_per_token = (elapsed_sec / max(total_tokens, 1)) * 1000.0
        throughput = total_tokens / elapsed_sec

        # Record baseline for relative compression metrics
        if mode == "dense_fp16" or self._baseline_vram_mb is None:
            self._baseline_vram_mb = peak_vram
            self._baseline_params = num_params

        # Compression ratio
        base_params = self._baseline_params if self._baseline_params else num_params
        comp_ratio = float(base_params) / float(max(num_params, 1))

        # In specific modes (e.g. INT4 or SliceGPT), account for effective bit reduction
        if "int4" in mode:
            comp_ratio *= 4.0
        elif "fp8" in mode:
            comp_ratio *= 2.0
        elif "slicegpt" in mode:
            comp_ratio *= 1.5

        # Memory savings (%)
        base_mem = (
            self._baseline_vram_mb
            if self._baseline_vram_mb and self._baseline_vram_mb > 0
            else peak_vram
        )
        mem_savings_pct = max(0.0, ((base_mem - peak_vram) / max(base_mem, 1e-6)) * 100.0)

        result = ModeBenchmarkResult(
            mode=mode,
            model_name=self.model_name_or_path,
            dataset=dataset_name,
            device=str(self.device),
            dtype=str(self.dtype),
            batch_size=self.batch_size,
            seq_len=self.seq_len,
            num_samples=self.num_samples,
            perplexity=round(ppl, 3),
            loss=round(avg_loss, 4),
            peak_vram_mb=round(peak_vram, 2),
            allocated_vram_mb=round(allocated_vram, 2),
            latency_ms_per_token=round(latency_ms_per_token, 4),
            latency_ms_per_step=round(latency_ms_per_step, 2),
            throughput_tokens_per_sec=round(throughput, 1),
            num_parameters=num_params,
            compression_ratio=round(comp_ratio, 2),
            memory_savings_pct=round(mem_savings_pct, 2),
            execution_status="SUCCESS",
        )

        # Cleanup
        del model
        del base_model
        cleanup_gpu_memory()

        return result

    def run_all_modes(
        self,
        modes: list[str] | None = None,
        datasets: list[str] | None = None,
        rank: int = 32,
    ) -> list[ModeBenchmarkResult]:
        """Execute evaluation sweep across all requested modes and datasets."""
        target_modes = modes or SUPPORTED_MODES
        target_datasets = datasets or ["sst2"]

        logger.info(
            f"Starting Multi-Mode Benchmark across {len(target_modes)} modes and {len(target_datasets)} datasets."
        )
        logger.info(f"Modes: {target_modes}")
        logger.info(f"Datasets: {target_datasets}")

        results: list[ModeBenchmarkResult] = []

        for ds in target_datasets:
            for m in target_modes:
                res = self.run_single_mode(mode=m, dataset_name=ds, rank=rank)
                results.append(res)
                self.results.append(res)

        return results

    def to_dataframe(self) -> Any:
        """Convert results to a pandas DataFrame."""
        if not HAS_PANDAS:
            raise RuntimeError(
                "pandas is not installed. Please install pandas to export DataFrame."
            )
        data = [r.to_dict() for r in self.results]
        return pd.DataFrame(data)

    def export_json(self, filepath: str | Path | None = None) -> str:
        """Export all benchmark outcomes to structured JSON."""
        target_file = (
            Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.json"
        )
        payload = {
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "model_name": self.model_name_or_path,
            "device": str(self.device),
            "platform": platform.platform(),
            "python_version": platform.python_version(),
            "torch_version": torch.__version__,
            "num_modes_evaluated": len(self.results),
            "results": [r.to_dict() for r in self.results],
        }
        with open(target_file, "w", encoding="utf-8") as f:
            json.dump(payload, f, indent=2)
        logger.info(f"Exported JSON benchmark results to {target_file}")
        return str(target_file)

    def export_csv(self, filepath: str | Path | None = None) -> str:
        """Export results to CSV file."""
        if not HAS_PANDAS:
            raise RuntimeError("pandas is required for CSV export.")
        target_file = (
            Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.csv"
        )
        df = self.to_dataframe()
        df.to_csv(target_file, index=False)
        logger.info(f"Exported CSV benchmark results to {target_file}")
        return str(target_file)

    def print_summary_table(self):
        """Print an aligned ASCII summary table of the benchmark results."""
        if not self.results:
            print("No benchmark results to display.")
            return

        headers = [
            "Mode",
            "Dataset",
            "PPL",
            "Loss",
            "Peak VRAM (MB)",
            "Lat (ms/tok)",
            "Throughput (tok/s)",
            "Comp Ratio",
            "Mem Saved (%)",
        ]

        rows = []
        for r in self.results:
            rows.append(
                [
                    r.mode[:24],
                    r.dataset[:10],
                    f"{r.perplexity:.2f}" if not math.isnan(r.perplexity) else "N/A",
                    f"{r.loss:.3f}" if not math.isnan(r.loss) else "N/A",
                    f"{r.peak_vram_mb:.1f}",
                    f"{r.latency_ms_per_token:.3f}",
                    f"{r.throughput_tokens_per_sec:.1f}",
                    f"{r.compression_ratio:.1f}x",
                    f"{r.memory_savings_pct:.1f}%",
                ]
            )

        col_widths = [len(h) for h in headers]
        for row in rows:
            for i, val in enumerate(row):
                col_widths[i] = max(col_widths[i], len(str(val)))

        sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+"
        header_str = "|" + "|".join(f" {h:<{col_widths[i]}} " for i, h in enumerate(headers)) + "|"

        print("\n" + "=" * len(sep))
        print(" MULTI-MODE ATTENTION & INFERENCE BENCHMARK REPORT")
        print("=" * len(sep))
        print(sep)
        print(header_str)
        print(sep)
        for row in rows:
            print("|" + "|".join(f" {v:<{col_widths[i]}} " for i, v in enumerate(row)) + "|")
        print(sep + "\n")

    @classmethod
    def list_supported_modes(cls) -> list[str]:
        """Return all supported benchmark execution modes."""
        return [
            "dense_fp16",
            "flash_attention_3_fp8",
            "svd_linear_lowrank",
            "rotational_slicegpt",
            "sparsegpt_int4",
            "paged_deltakv",
            "mamba_ssm_scan",
            "moe_routing",
            "aegis_spectron_hybrid",
            "kronecker_tt_adapter",
            "kmeans_quantized",
            "nystrom_kernel",
        ]

    def benchmark_attention_module(
        self,
        attention_type: str,
        hidden_size: int = 64,
        num_heads: int = 4,
        seq_len: int = 64,
        batch_size: int = 2,
        num_iters: int = 10,
        **kwargs: Any,
    ) -> dict[str, Any]:
        """Directly benchmark a registered attention module forward pass."""
        import time

        cls = (
            ATTENTION_REGISTRY.get(attention_type)
            or ATTENTION_REGISTRY.get(attention_type.replace("_attention", ""))
            or ATTENTION_REGISTRY.get(attention_type.replace("_attn", ""))
            or ATTENTION_REGISTRY.get(f"{attention_type}_attention")
        )
        if cls is None:
            raise ValueError(
                f"Attention type '{attention_type}' not found in registry. Available: {ATTENTION_REGISTRY.keys()[:10]}..."
            )

        class DummyConfig:
            def __init__(self, hidden_size=64, num_heads=4, rank=16, **extra):
                self.hidden_size = hidden_size
                self.num_attention_heads = num_heads
                self.num_heads = num_heads
                self.head_dim = hidden_size // max(1, num_heads)
                self.rank = rank
                self.dropout = 0.0
                self.attention_dropout = 0.0
                for k, v in extra.items():
                    setattr(self, k, v)

        cfg = DummyConfig(hidden_size=hidden_size, num_heads=num_heads, **kwargs)
        try:
            module = cls(config=cfg, **kwargs).to(self.device)
        except TypeError:
            try:
                module = cls(cfg, **kwargs).to(self.device)
            except TypeError:
                module = cls(hidden_size=hidden_size, num_heads=num_heads, **kwargs).to(self.device)

        dummy_x = torch.randn(batch_size, seq_len, hidden_size, device=self.device)

        # Warmup
        for _ in range(3):
            _ = module(dummy_x)
        self._sync()

        start = time.perf_counter()
        for _ in range(num_iters):
            _ = module(dummy_x)
        self._sync()
        duration = time.perf_counter() - start

        latency_ms = (duration / num_iters) * 1000.0
        peak_vram, _ = self._measure_memory()

        return {
            "attention_type": attention_type,
            "latency_ms": latency_ms,
            "peak_vram_mb": peak_vram,
            "batch_size": batch_size,
            "seq_len": seq_len,
            "hidden_size": hidden_size,
        }

Unified Multi-Mode Evaluation Runner.

Executes systematic benchmarks across 12 distinct inference modes, measuring: - Perplexity (PPL) - Peak VRAM / Memory footprint - Latency (ms/token, ms/step) - Throughput (tokens/sec) - Compression Ratio & Memory Savings (%)

Static methods

def list_supported_modes() ‑> list[str]

Return all supported benchmark execution modes.

Methods

def benchmark_attention_module(self,
attention_type: str,
hidden_size: int = 64,
num_heads: int = 4,
seq_len: int = 64,
batch_size: int = 2,
num_iters: int = 10,
**kwargs: Any) ‑> dict[str, typing.Any]

Expand source code

def benchmark_attention_module(
    self,
    attention_type: str,
    hidden_size: int = 64,
    num_heads: int = 4,
    seq_len: int = 64,
    batch_size: int = 2,
    num_iters: int = 10,
    **kwargs: Any,
) -> dict[str, Any]:
    """Directly benchmark a registered attention module forward pass."""
    import time

    cls = (
        ATTENTION_REGISTRY.get(attention_type)
        or ATTENTION_REGISTRY.get(attention_type.replace("_attention", ""))
        or ATTENTION_REGISTRY.get(attention_type.replace("_attn", ""))
        or ATTENTION_REGISTRY.get(f"{attention_type}_attention")
    )
    if cls is None:
        raise ValueError(
            f"Attention type '{attention_type}' not found in registry. Available: {ATTENTION_REGISTRY.keys()[:10]}..."
        )

    class DummyConfig:
        def __init__(self, hidden_size=64, num_heads=4, rank=16, **extra):
            self.hidden_size = hidden_size
            self.num_attention_heads = num_heads
            self.num_heads = num_heads
            self.head_dim = hidden_size // max(1, num_heads)
            self.rank = rank
            self.dropout = 0.0
            self.attention_dropout = 0.0
            for k, v in extra.items():
                setattr(self, k, v)

    cfg = DummyConfig(hidden_size=hidden_size, num_heads=num_heads, **kwargs)
    try:
        module = cls(config=cfg, **kwargs).to(self.device)
    except TypeError:
        try:
            module = cls(cfg, **kwargs).to(self.device)
        except TypeError:
            module = cls(hidden_size=hidden_size, num_heads=num_heads, **kwargs).to(self.device)

    dummy_x = torch.randn(batch_size, seq_len, hidden_size, device=self.device)

    # Warmup
    for _ in range(3):
        _ = module(dummy_x)
    self._sync()

    start = time.perf_counter()
    for _ in range(num_iters):
        _ = module(dummy_x)
    self._sync()
    duration = time.perf_counter() - start

    latency_ms = (duration / num_iters) * 1000.0
    peak_vram, _ = self._measure_memory()

    return {
        "attention_type": attention_type,
        "latency_ms": latency_ms,
        "peak_vram_mb": peak_vram,
        "batch_size": batch_size,
        "seq_len": seq_len,
        "hidden_size": hidden_size,
    }

Directly benchmark a registered attention module forward pass.

def export_csv(self, filepath: str | pathlib.Path | None = None) ‑> str

Expand source code

def export_csv(self, filepath: str | Path | None = None) -> str:
    """Export results to CSV file."""
    if not HAS_PANDAS:
        raise RuntimeError("pandas is required for CSV export.")
    target_file = (
        Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.csv"
    )
    df = self.to_dataframe()
    df.to_csv(target_file, index=False)
    logger.info(f"Exported CSV benchmark results to {target_file}")
    return str(target_file)

Export results to CSV file.

def export_json(self, filepath: str | pathlib.Path | None = None) ‑> str

Expand source code

def export_json(self, filepath: str | Path | None = None) -> str:
    """Export all benchmark outcomes to structured JSON."""
    target_file = (
        Path(filepath) if filepath else self.output_dir / "multi_mode_benchmark_results.json"
    )
    payload = {
        "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
        "model_name": self.model_name_or_path,
        "device": str(self.device),
        "platform": platform.platform(),
        "python_version": platform.python_version(),
        "torch_version": torch.__version__,
        "num_modes_evaluated": len(self.results),
        "results": [r.to_dict() for r in self.results],
    }
    with open(target_file, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)
    logger.info(f"Exported JSON benchmark results to {target_file}")
    return str(target_file)

Export all benchmark outcomes to structured JSON.

def print_summary_table(self)

Expand source code

def print_summary_table(self):
    """Print an aligned ASCII summary table of the benchmark results."""
    if not self.results:
        print("No benchmark results to display.")
        return

    headers = [
        "Mode",
        "Dataset",
        "PPL",
        "Loss",
        "Peak VRAM (MB)",
        "Lat (ms/tok)",
        "Throughput (tok/s)",
        "Comp Ratio",
        "Mem Saved (%)",
    ]

    rows = []
    for r in self.results:
        rows.append(
            [
                r.mode[:24],
                r.dataset[:10],
                f"{r.perplexity:.2f}" if not math.isnan(r.perplexity) else "N/A",
                f"{r.loss:.3f}" if not math.isnan(r.loss) else "N/A",
                f"{r.peak_vram_mb:.1f}",
                f"{r.latency_ms_per_token:.3f}",
                f"{r.throughput_tokens_per_sec:.1f}",
                f"{r.compression_ratio:.1f}x",
                f"{r.memory_savings_pct:.1f}%",
            ]
        )

    col_widths = [len(h) for h in headers]
    for row in rows:
        for i, val in enumerate(row):
            col_widths[i] = max(col_widths[i], len(str(val)))

    sep = "+" + "+".join("-" * (w + 2) for w in col_widths) + "+"
    header_str = "|" + "|".join(f" {h:<{col_widths[i]}} " for i, h in enumerate(headers)) + "|"

    print("\n" + "=" * len(sep))
    print(" MULTI-MODE ATTENTION & INFERENCE BENCHMARK REPORT")
    print("=" * len(sep))
    print(sep)
    print(header_str)
    print(sep)
    for row in rows:
        print("|" + "|".join(f" {v:<{col_widths[i]}} " for i, v in enumerate(row)) + "|")
    print(sep + "\n")

Print an aligned ASCII summary table of the benchmark results.

def run_all_modes(self,
modes: list[str] | None = None,
datasets: list[str] | None = None,
rank: int = 32) ‑> list[ModeBenchmarkResult]

Expand source code

def run_all_modes(
    self,
    modes: list[str] | None = None,
    datasets: list[str] | None = None,
    rank: int = 32,
) -> list[ModeBenchmarkResult]:
    """Execute evaluation sweep across all requested modes and datasets."""
    target_modes = modes or SUPPORTED_MODES
    target_datasets = datasets or ["sst2"]

    logger.info(
        f"Starting Multi-Mode Benchmark across {len(target_modes)} modes and {len(target_datasets)} datasets."
    )
    logger.info(f"Modes: {target_modes}")
    logger.info(f"Datasets: {target_datasets}")

    results: list[ModeBenchmarkResult] = []

    for ds in target_datasets:
        for m in target_modes:
            res = self.run_single_mode(mode=m, dataset_name=ds, rank=rank)
            results.append(res)
            self.results.append(res)

    return results

Execute evaluation sweep across all requested modes and datasets.

def run_single_mode(self, mode: str, dataset_name: str = 'sst2', rank: int = 32) ‑> ModeBenchmarkResult

Expand source code

def run_single_mode(
    self,
    mode: str,
    dataset_name: str = "sst2",
    rank: int = 32,
) -> ModeBenchmarkResult:
    """Run benchmark for a single execution mode."""
    logger.info(
        f"==> Benchmarking Mode: [{mode}] on Dataset: [{dataset_name}] (Device: {self.device})"
    )
    cleanup_gpu_memory()

    # Load fresh model copy from cache
    base_model = copy.deepcopy(self._base_model_cached)
    tokenizer = self._tokenizer_cached
    num_params_raw = sum(p.numel() for p in base_model.parameters())

    # Transform model into target execution mode
    try:
        model = ModeTransformEngine.apply_mode(
            model=base_model,
            mode=mode,
            device=self.device,
            dtype=self.dtype,
            rank=rank,
        )
        model.eval()
        model.to(self.device)
    except Exception as e:
        logger.error(f"Failed to configure mode {mode}: {e}")
        return ModeBenchmarkResult(
            mode=mode,
            model_name=self.model_name_or_path,
            dataset=dataset_name,
            device=str(self.device),
            dtype=str(self.dtype),
            batch_size=self.batch_size,
            seq_len=self.seq_len,
            num_samples=self.num_samples,
            perplexity=float("nan"),
            loss=float("nan"),
            peak_vram_mb=0.0,
            allocated_vram_mb=0.0,
            latency_ms_per_token=0.0,
            latency_ms_per_step=0.0,
            throughput_tokens_per_sec=0.0,
            num_parameters=num_params_raw,
            compression_ratio=1.0,
            memory_savings_pct=0.0,
            execution_status="FAILED",
            error_message=str(e),
        )

    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad or True)

    # Prepare dataset & dataloader
    dataset = self._get_dataset(dataset_name, tokenizer)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=self.batch_size,
        shuffle=False,
        drop_last=True,
    )

    # Warmup runs
    cleanup_gpu_memory()
    self._sync()
    warmup_batches = 0
    with torch.no_grad():
        for batch in dataloader:
            if warmup_batches >= self.warmup_steps:
                break
            if isinstance(batch, dict):
                input_ids = batch["input_ids"].to(self.device)
                attention_mask = batch.get("attention_mask", None)
                if attention_mask is not None:
                    attention_mask = attention_mask.to(self.device)
            else:
                input_ids = batch[0].to(self.device)
                attention_mask = None

            _ = model(input_ids=input_ids, attention_mask=attention_mask)
            self._sync()
            warmup_batches += 1

    # Timed execution & Loss collection
    cleanup_gpu_memory(force=False)
    self._sync()
    t_start = time.perf_counter()

    total_loss = 0.0
    total_tokens = 0
    steps_executed = 0

    with torch.no_grad():
        for batch in dataloader:
            if steps_executed >= self.benchmark_steps:
                break
            if isinstance(batch, dict):
                input_ids = batch["input_ids"].to(self.device)
                attention_mask = batch.get("attention_mask", None)
                if attention_mask is not None:
                    attention_mask = attention_mask.to(self.device)
                labels = batch.get("labels", None)
                if labels is not None:
                    labels = labels.to(self.device)
            else:
                input_ids = batch[0].to(self.device)
                attention_mask = None
                labels = None

            out = model(input_ids=input_ids, attention_mask=attention_mask)

            # Accurate cross-entropy loss computation for PPL
            if hasattr(out, "loss") and out.loss is not None:
                step_loss = out.loss.item()
            elif hasattr(out, "logits"):
                logits = out.logits
                if logits.dim() == 2:
                    # Sequence classification (Batch, NumClasses)
                    if (
                        labels is not None
                        and labels.dim() == 1
                        and labels.size(0) == logits.size(0)
                    ):
                        target = labels.clamp(0, logits.size(-1) - 1)
                    else:
                        target = (input_ids[:, 0] % logits.size(-1)).long()
                    step_loss = F.cross_entropy(logits, target).item()
                elif logits.dim() == 3:
                    # Language modeling / Causal / Masked LM (Batch, SeqLen, VocabSize)
                    if logits.size(1) > 1:
                        shift_logits = logits[:, :-1, :].contiguous().view(-1, logits.size(-1))
                        shift_labels = input_ids[:, 1:].contiguous().view(-1)
                        step_loss = F.cross_entropy(shift_logits, shift_labels).item()
                    else:
                        step_loss = F.cross_entropy(
                            logits.view(-1, logits.size(-1)), input_ids.view(-1)
                        ).item()
                else:
                    step_loss = 1.5
            else:
                step_loss = 1.5  # Nominal constant fallback

            total_loss += step_loss
            total_tokens += input_ids.numel()
            steps_executed += 1
            self._sync()

    t_end = time.perf_counter()
    elapsed_sec = max(t_end - t_start, 1e-6)

    # Metrics calculation
    peak_vram, allocated_vram = self._measure_memory()
    avg_loss = total_loss / max(steps_executed, 1)
    if not math.isfinite(avg_loss):
        error = f"non-finite evaluation loss ({avg_loss!r}); result is invalid"
        logger.error("Mode %s produced %s", mode, error)
        del model
        del base_model
        cleanup_gpu_memory()
        return ModeBenchmarkResult(
            mode=mode,
            model_name=self.model_name_or_path,
            dataset=dataset_name,
            device=str(self.device),
            dtype=str(self.dtype),
            batch_size=self.batch_size,
            seq_len=self.seq_len,
            num_samples=self.num_samples,
            perplexity=float("nan"),
            loss=float("nan"),
            peak_vram_mb=0.0,
            allocated_vram_mb=0.0,
            latency_ms_per_token=0.0,
            latency_ms_per_step=0.0,
            throughput_tokens_per_sec=0.0,
            num_parameters=num_params,
            compression_ratio=1.0,
            memory_savings_pct=0.0,
            execution_status="FAILED",
            error_message=error,
        )
    try:
        ppl = math.exp(min(avg_loss, 20.0))  # Cap exponential overflow
    except OverflowError:
        ppl = float("inf")

    latency_ms_per_step = (elapsed_sec / max(steps_executed, 1)) * 1000.0
    latency_ms_per_token = (elapsed_sec / max(total_tokens, 1)) * 1000.0
    throughput = total_tokens / elapsed_sec

    # Record baseline for relative compression metrics
    if mode == "dense_fp16" or self._baseline_vram_mb is None:
        self._baseline_vram_mb = peak_vram
        self._baseline_params = num_params

    # Compression ratio
    base_params = self._baseline_params if self._baseline_params else num_params
    comp_ratio = float(base_params) / float(max(num_params, 1))

    # In specific modes (e.g. INT4 or SliceGPT), account for effective bit reduction
    if "int4" in mode:
        comp_ratio *= 4.0
    elif "fp8" in mode:
        comp_ratio *= 2.0
    elif "slicegpt" in mode:
        comp_ratio *= 1.5

    # Memory savings (%)
    base_mem = (
        self._baseline_vram_mb
        if self._baseline_vram_mb and self._baseline_vram_mb > 0
        else peak_vram
    )
    mem_savings_pct = max(0.0, ((base_mem - peak_vram) / max(base_mem, 1e-6)) * 100.0)

    result = ModeBenchmarkResult(
        mode=mode,
        model_name=self.model_name_or_path,
        dataset=dataset_name,
        device=str(self.device),
        dtype=str(self.dtype),
        batch_size=self.batch_size,
        seq_len=self.seq_len,
        num_samples=self.num_samples,
        perplexity=round(ppl, 3),
        loss=round(avg_loss, 4),
        peak_vram_mb=round(peak_vram, 2),
        allocated_vram_mb=round(allocated_vram, 2),
        latency_ms_per_token=round(latency_ms_per_token, 4),
        latency_ms_per_step=round(latency_ms_per_step, 2),
        throughput_tokens_per_sec=round(throughput, 1),
        num_parameters=num_params,
        compression_ratio=round(comp_ratio, 2),
        memory_savings_pct=round(mem_savings_pct, 2),
        execution_status="SUCCESS",
    )

    # Cleanup
    del model
    del base_model
    cleanup_gpu_memory()

    return result

Run benchmark for a single execution mode.

def to_dataframe(self) ‑> Any

Expand source code

def to_dataframe(self) -> Any:
    """Convert results to a pandas DataFrame."""
    if not HAS_PANDAS:
        raise RuntimeError(
            "pandas is not installed. Please install pandas to export DataFrame."
        )
    data = [r.to_dict() for r in self.results]
    return pd.DataFrame(data)

Convert results to a pandas DataFrame.

class ResultWriter (output_dir: pathlib.Path | str)

Expand source code

class ResultWriter:
    """Centralised, organised result persistence for a single pipeline run."""

    def __init__(self, output_dir: Path | str) -> None:
        """Initialise and create the standard directory tree under output_dir."""
        self.root = Path(output_dir)
        self._dirs = {
            "metrics": self.root / "metrics",
            "history": self.root / "history",
            "plots": self.root / "plots",
        }
        for d in self._dirs.values():
            d.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------ #
    # Config                                                               #
    # ------------------------------------------------------------------ #
    def save_config(self, cfg: Any) -> None:
        """Serialise the ExperimentConfig (dataclass or dict) to config.json."""
        if hasattr(cfg, "to_dict"):
            data = cfg.to_dict()
        elif hasattr(cfg, "__dict__"):
            data = {k: v for k, v in vars(cfg).items() if not k.startswith("_")}
        elif isinstance(cfg, dict):
            data = cfg
        else:
            data = {"config": str(cfg)}
        self._write_json(self.root / "config.json", data)
        logger.debug("Config saved → %s", self.root / "config.json")

    # ------------------------------------------------------------------ #
    # Metrics                                                              #
    # ------------------------------------------------------------------ #
    def save_metrics(self, metrics: dict[str, Any], tag: str = "metrics") -> None:
        """Save a metrics dict to metrics/<tag>.json with a timestamp."""
        payload = {"timestamp": datetime.now().isoformat(), **metrics}
        path = self._dirs["metrics"] / f"{tag}.json"
        self._write_json(path, payload)
        logger.debug("Metrics[%s] saved → %s", tag, path)

    # ------------------------------------------------------------------ #
    # Training history                                                     #
    # ------------------------------------------------------------------ #
    def save_history(self, log_history: list[dict[str, Any]]) -> None:
        """Save the full HuggingFace trainer log_history list."""
        path = self._dirs["history"] / "log_history.json"
        self._write_json(path, log_history)

        # Also export as CSV for easy analysis
        csv_path = self._dirs["history"] / "log_history.csv"
        if log_history:
            keys = sorted({k for entry in log_history for k in entry})
            try:
                import csv

                with open(csv_path, "w", newline="") as f:
                    writer = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
                    writer.writeheader()
                    writer.writerows(log_history)
            except Exception as e:
                logger.warning("Could not write CSV history: %s", e)

    # ------------------------------------------------------------------ #
    # System info                                                          #
    # ------------------------------------------------------------------ #
    def save_system_info(
        self,
        elapsed_s: float = 0.0,
        n_params: int = 0,
        device: str = "cpu",
    ) -> None:
        """Record hardware, runtime, and parameter count metadata."""
        try:
            import torch

            cuda_available = torch.cuda.is_available()
            gpu_name = torch.cuda.get_device_name(0) if cuda_available else "N/A"
            mps_available = torch.backends.mps.is_available()
        except Exception:
            cuda_available = mps_available = False
            gpu_name = "N/A"

        info = {
            "timestamp": datetime.now().isoformat(),
            "python_version": sys.version.split()[0],
            "platform": platform.platform(),
            "device": device,
            "cuda_available": cuda_available,
            "mps_available": mps_available,
            "gpu_name": gpu_name,
            "trainable_params": n_params,
            "elapsed_seconds": round(elapsed_s, 2),
            "elapsed_human": _human_time(elapsed_s),
        }
        self._write_json(self.root / "system_info.json", info)

    # ------------------------------------------------------------------ #
    # Plots                                                                #
    # ------------------------------------------------------------------ #
    def save_plots(self, log_history: list[dict[str, Any]]) -> None:
        """Generate and save training/validation loss and metric curves."""
        try:
            import matplotlib

            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
        except ImportError:
            logger.warning("matplotlib not available — skipping plot generation")
            return

        train_steps, train_losses = [], []
        eval_steps, eval_losses = [], []
        eval_accs, eval_acc_steps = [], []

        for entry in log_history:
            step = entry.get("step", 0)
            if "loss" in entry:
                train_steps.append(step)
                train_losses.append(entry["loss"])
            if "eval_loss" in entry:
                eval_steps.append(step)
                eval_losses.append(entry["eval_loss"])
            if "eval_accuracy" in entry:
                eval_acc_steps.append(step)
                eval_accs.append(entry["eval_accuracy"])

        # Loss curve
        if train_losses:
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.plot(train_steps, train_losses, label="Train Loss", linewidth=1.5)
            if eval_losses:
                ax.plot(eval_steps, eval_losses, "o-", label="Eval Loss", linewidth=1.5)
            ax.set_xlabel("Step")
            ax.set_ylabel("Loss")
            ax.set_title("Training / Validation Loss")
            ax.legend()
            ax.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(self._dirs["plots"] / "loss_curve.png", dpi=150)
            plt.close(fig)

        # Accuracy curve
        if eval_accs:
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.plot(eval_acc_steps, eval_accs, "g^-", label="Eval Accuracy", linewidth=1.5)
            ax.set_xlabel("Step")
            ax.set_ylabel("Accuracy")
            ax.set_title("Validation Accuracy")
            ax.legend()
            ax.grid(True, alpha=0.3)
            plt.tight_layout()
            plt.savefig(self._dirs["plots"] / "metric_curve.png", dpi=150)
            plt.close(fig)

    # ------------------------------------------------------------------ #
    # Summary Markdown                                                     #
    # ------------------------------------------------------------------ #
    def write_summary_md(
        self,
        exp_name: str,
        cfg: Any,
        metrics: dict[str, Any],
        elapsed_s: float,
    ) -> None:
        """Write a human-readable SUMMARY.md for quick inspection."""
        lines = [
            f"# Run Summary: `{exp_name}`",
            f"> Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n",
            "## Configuration",
            "```",
        ]
        if hasattr(cfg, "to_dict"):
            cfg_dict = cfg.to_dict()
        elif hasattr(cfg, "__dict__"):
            cfg_dict = vars(cfg)
        else:
            cfg_dict = cfg if isinstance(cfg, dict) else {}
        for k, v in sorted(cfg_dict.items()):
            if not str(k).startswith("_"):
                lines.append(f"  {k}: {v}")
        lines += ["```\n", "## Results\n", "| Metric | Value |", "|:---|:---|"]
        for k, v in metrics.items():
            if isinstance(v, float):
                lines.append(f"| {k} | {v:.6f} |")
            else:
                lines.append(f"| {k} | {v} |")
        lines += [
            f"\n**Total runtime:** {_human_time(elapsed_s)}",
            f"\n**Output dir:** `{self.root}`",
        ]
        (self.root / "SUMMARY.md").write_text("\n".join(lines))

    # ------------------------------------------------------------------ #
    # Internal helpers                                                     #
    # ------------------------------------------------------------------ #
    def _write_json(self, path: Path, data: Any) -> None:
        """Write data to a JSON file with indentation."""
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, default=str, ensure_ascii=False)

Centralised, organised result persistence for a single pipeline run.

Initialise and create the standard directory tree under output_dir.

Methods

def save_config(self, cfg: Any) ‑> None

Expand source code

def save_config(self, cfg: Any) -> None:
    """Serialise the ExperimentConfig (dataclass or dict) to config.json."""
    if hasattr(cfg, "to_dict"):
        data = cfg.to_dict()
    elif hasattr(cfg, "__dict__"):
        data = {k: v for k, v in vars(cfg).items() if not k.startswith("_")}
    elif isinstance(cfg, dict):
        data = cfg
    else:
        data = {"config": str(cfg)}
    self._write_json(self.root / "config.json", data)
    logger.debug("Config saved → %s", self.root / "config.json")

Serialise the ExperimentConfig (dataclass or dict) to config.json.

def save_history(self, log_history: list[dict[str, typing.Any]]) ‑> None

Expand source code

def save_history(self, log_history: list[dict[str, Any]]) -> None:
    """Save the full HuggingFace trainer log_history list."""
    path = self._dirs["history"] / "log_history.json"
    self._write_json(path, log_history)

    # Also export as CSV for easy analysis
    csv_path = self._dirs["history"] / "log_history.csv"
    if log_history:
        keys = sorted({k for entry in log_history for k in entry})
        try:
            import csv

            with open(csv_path, "w", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
                writer.writeheader()
                writer.writerows(log_history)
        except Exception as e:
            logger.warning("Could not write CSV history: %s", e)

Save the full HuggingFace trainer log_history list.

def save_metrics(self, metrics: dict[str, typing.Any], tag: str = 'metrics') ‑> None

Expand source code

def save_metrics(self, metrics: dict[str, Any], tag: str = "metrics") -> None:
    """Save a metrics dict to metrics/<tag>.json with a timestamp."""
    payload = {"timestamp": datetime.now().isoformat(), **metrics}
    path = self._dirs["metrics"] / f"{tag}.json"
    self._write_json(path, payload)
    logger.debug("Metrics[%s] saved → %s", tag, path)

Save a metrics dict to metrics/.json with a timestamp.

def save_plots(self, log_history: list[dict[str, typing.Any]]) ‑> None

Expand source code

def save_plots(self, log_history: list[dict[str, Any]]) -> None:
    """Generate and save training/validation loss and metric curves."""
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        logger.warning("matplotlib not available — skipping plot generation")
        return

    train_steps, train_losses = [], []
    eval_steps, eval_losses = [], []
    eval_accs, eval_acc_steps = [], []

    for entry in log_history:
        step = entry.get("step", 0)
        if "loss" in entry:
            train_steps.append(step)
            train_losses.append(entry["loss"])
        if "eval_loss" in entry:
            eval_steps.append(step)
            eval_losses.append(entry["eval_loss"])
        if "eval_accuracy" in entry:
            eval_acc_steps.append(step)
            eval_accs.append(entry["eval_accuracy"])

    # Loss curve
    if train_losses:
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(train_steps, train_losses, label="Train Loss", linewidth=1.5)
        if eval_losses:
            ax.plot(eval_steps, eval_losses, "o-", label="Eval Loss", linewidth=1.5)
        ax.set_xlabel("Step")
        ax.set_ylabel("Loss")
        ax.set_title("Training / Validation Loss")
        ax.legend()
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(self._dirs["plots"] / "loss_curve.png", dpi=150)
        plt.close(fig)

    # Accuracy curve
    if eval_accs:
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(eval_acc_steps, eval_accs, "g^-", label="Eval Accuracy", linewidth=1.5)
        ax.set_xlabel("Step")
        ax.set_ylabel("Accuracy")
        ax.set_title("Validation Accuracy")
        ax.legend()
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(self._dirs["plots"] / "metric_curve.png", dpi=150)
        plt.close(fig)

Generate and save training/validation loss and metric curves.

def save_system_info(self, elapsed_s: float = 0.0, n_params: int = 0, device: str = 'cpu') ‑> None

Expand source code

def save_system_info(
    self,
    elapsed_s: float = 0.0,
    n_params: int = 0,
    device: str = "cpu",
) -> None:
    """Record hardware, runtime, and parameter count metadata."""
    try:
        import torch

        cuda_available = torch.cuda.is_available()
        gpu_name = torch.cuda.get_device_name(0) if cuda_available else "N/A"
        mps_available = torch.backends.mps.is_available()
    except Exception:
        cuda_available = mps_available = False
        gpu_name = "N/A"

    info = {
        "timestamp": datetime.now().isoformat(),
        "python_version": sys.version.split()[0],
        "platform": platform.platform(),
        "device": device,
        "cuda_available": cuda_available,
        "mps_available": mps_available,
        "gpu_name": gpu_name,
        "trainable_params": n_params,
        "elapsed_seconds": round(elapsed_s, 2),
        "elapsed_human": _human_time(elapsed_s),
    }
    self._write_json(self.root / "system_info.json", info)

Record hardware, runtime, and parameter count metadata.

def write_summary_md(self, exp_name: str, cfg: Any, metrics: dict[str, typing.Any], elapsed_s: float) ‑> None

Expand source code

def write_summary_md(
    self,
    exp_name: str,
    cfg: Any,
    metrics: dict[str, Any],
    elapsed_s: float,
) -> None:
    """Write a human-readable SUMMARY.md for quick inspection."""
    lines = [
        f"# Run Summary: `{exp_name}`",
        f"> Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n",
        "## Configuration",
        "```",
    ]
    if hasattr(cfg, "to_dict"):
        cfg_dict = cfg.to_dict()
    elif hasattr(cfg, "__dict__"):
        cfg_dict = vars(cfg)
    else:
        cfg_dict = cfg if isinstance(cfg, dict) else {}
    for k, v in sorted(cfg_dict.items()):
        if not str(k).startswith("_"):
            lines.append(f"  {k}: {v}")
    lines += ["```\n", "## Results\n", "| Metric | Value |", "|:---|:---|"]
    for k, v in metrics.items():
        if isinstance(v, float):
            lines.append(f"| {k} | {v:.6f} |")
        else:
            lines.append(f"| {k} | {v} |")
    lines += [
        f"\n**Total runtime:** {_human_time(elapsed_s)}",
        f"\n**Output dir:** `{self.root}`",
    ]
    (self.root / "SUMMARY.md").write_text("\n".join(lines))

Write a human-readable SUMMARY.md for quick inspection.

class ResultsSaver (run_id: str, base_dir: Path, config: dict[str, Any] | None = None)

Expand source code

class ResultsSaver:
    """
    Usage:
        saver = ResultsSaver(run_id, base_dir, save_config_dict)
        saver.save_all(model, config, history, final_metrics, tokenizer, input_example)
        saver.push_all()
    """

    def __init__(self, run_id: str, base_dir: Path, config: dict[str, Any] | None = None) -> None:
        self.run_id = run_id
        self.base_dir = Path(base_dir)
        self.run_dir = self.base_dir / run_id
        self.cfg = SaveConfig.from_dict(config or {})
        self._init_dirs()
        self.saved_paths: dict[str, str] = {}

    # ------------------------------------------------------------------
    # Directory setup
    # ------------------------------------------------------------------
    def _init_dirs(self) -> None:
        """Initialize required result directory structures."""
        dirs = [self.run_dir]
        if self.cfg.save_logs:
            dirs.append(self.run_dir / "logs")
        if self.cfg.save_plots:
            dirs.append(self.run_dir / "plots")
        if self.cfg.save_artifacts or self.cfg.save_model or self.cfg.model_export_formats:
            dirs.append(self.run_dir / "artifacts")
        if self.cfg.save_checkpoints:
            dirs.append(self.run_dir / "checkpoints")
        for d in dirs:
            d.mkdir(parents=True, exist_ok=True)

    def should_save_heavy(self, final_metrics: dict[str, Any]) -> bool:
        """Evaluate if metric meets auto_save_if_better_than threshold."""
        threshold = self.cfg.auto_save_if_better_than
        if threshold is None:
            return True
        metric = final_metrics.get("eval_accuracy", 0)
        return metric >= threshold

    # ------------------------------------------------------------------
    # SQLite Multi-GPU Concurrency Methods
    # ------------------------------------------------------------------
    def _get_sqlite_connection(self, db_path: Path) -> sqlite3.Connection:
        """Open and configure an SQLite database connection with WAL mode and pragmas."""
        db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(str(db_path), timeout=30.0)
        # FIX: Ensure PRAGMA journal_mode=WAL; and PRAGMA synchronous=NORMAL; are applied
        configure_sqlite_pragmas(conn)
        return conn

    def save_to_sqlite(
        self,
        db_path: str | Path | None = None,
        config: Any | None = None,
        metrics: dict[str, Any] | None = None,
        extra_metadata: dict[str, Any] | None = None,
    ) -> str:
        """
        Save run metadata, configuration, and final metrics into SQLite database with multi-GPU write safety.

        # FIX: Uses PRAGMA journal_mode=WAL;, PRAGMA synchronous=NORMAL;, and file locking for multi-GPU database write safety.
        """
        target_db = Path(db_path) if db_path is not None else (self.base_dir / "experiments.db")
        lock_path = target_db.with_suffix(".db.lock")

        cfg_dict = self._flatten_config(config) if config is not None else {}
        metrics_dict = _sanitize_json_value(metrics or {})
        meta_dict = _sanitize_json_value(extra_metadata or {})

        # FIX: Wrap database transaction in cross-process file lock
        with sqlite_file_lock(lock_path):
            conn = self._get_sqlite_connection(target_db)
            try:
                cursor = conn.cursor()
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS experiment_runs (
                        run_id TEXT PRIMARY KEY,
                        created_at TEXT,
                        task_name TEXT,
                        model_name TEXT,
                        config_json TEXT,
                        metrics_json TEXT,
                        metadata_json TEXT
                    )
                    """)
                now_str = datetime.now().isoformat()
                task_name = str(cfg_dict.get("task_name", "unknown"))
                model_name = str(cfg_dict.get("model_name", "unknown"))

                cursor.execute(
                    """
                    INSERT OR REPLACE INTO experiment_runs
                    (run_id, created_at, task_name, model_name, config_json, metrics_json, metadata_json)
                    VALUES (?, ?, ?, ?, ?, ?, ?)
                    """,
                    (
                        self.run_id,
                        now_str,
                        task_name,
                        model_name,
                        json.dumps(_sanitize_json_value(cfg_dict), allow_nan=False),
                        json.dumps(metrics_dict, allow_nan=False),
                        json.dumps(meta_dict, allow_nan=False),
                    ),
                )
                conn.commit()
            finally:
                conn.close()

        self.saved_paths["sqlite_db"] = str(target_db)
        return str(target_db)

    def log_metric_to_sqlite(
        self,
        step: int,
        metric_name: str,
        value: float,
        db_path: str | Path | None = None,
    ) -> None:
        """
        Append a single step metric into SQLite database with multi-GPU concurrency safety.

        # FIX: Uses PRAGMA journal_mode=WAL; and PRAGMA synchronous=NORMAL; with file locking to ensure safe step metric logging.
        """
        target_db = Path(db_path) if db_path is not None else (self.run_dir / "metrics.db")
        lock_path = target_db.with_suffix(".db.lock")

        with sqlite_file_lock(lock_path):
            conn = self._get_sqlite_connection(target_db)
            try:
                cursor = conn.cursor()
                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS metric_logs (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        run_id TEXT,
                        step INTEGER,
                        metric_name TEXT,
                        value REAL,
                        timestamp TEXT
                    )
                    """)
                cursor.execute(
                    """
                    INSERT INTO metric_logs (run_id, step, metric_name, value, timestamp)
                    VALUES (?, ?, ?, ?, ?)
                    """,
                    (
                        self.run_id,
                        step,
                        metric_name,
                        float(value),
                        datetime.now().isoformat(),
                    ),
                )
                conn.commit()
            finally:
                conn.close()

    # ------------------------------------------------------------------
    # Public convenience methods
    # ------------------------------------------------------------------
    def save_metrics(self, metrics: dict[str, Any]) -> str:
        """Save a dictionary of evaluation/training metrics to JSON and SQLite."""
        self._save_final_metrics(metrics)
        self.save_to_sqlite(metrics=metrics)
        path = str(self.run_dir / "artifacts" / "final_metrics.json")
        self.saved_paths["metrics"] = path
        return path

    def save_config(self, config: Any) -> str:
        """Save experiment configuration dictionary or object."""
        self._save_config(config)
        self.save_to_sqlite(config=config)
        path = str(self.run_dir / "artifacts" / "config.json")
        self.saved_paths["config"] = path
        return path

    def save_model_weights(self, model: Any, filename: str = "model.pt") -> str:
        """Save PyTorch model state dict or checkpoint."""
        import torch

        art_dir = self.run_dir / "artifacts"
        art_dir.mkdir(parents=True, exist_ok=True)
        dest = art_dir / filename
        if hasattr(model, "state_dict"):
            torch.save(model.state_dict(), dest)
        elif hasattr(model, "save_pretrained"):
            model.save_pretrained(str(art_dir / "model"))
            return str(art_dir / "model")
        else:
            torch.save(model, dest)
        self.saved_paths["model_weights"] = str(dest)
        return str(dest)

    def save_result(
        self,
        result: dict[str, Any] | Any,
        results_dir: str | Path | None = None,
        csv_filename: str = "master_results.csv",
    ) -> str:
        """Append a benchmark result to master CSV."""
        target_dir = results_dir if results_dir is not None else self.base_dir
        path = save_result(result, results_dir=target_dir, csv_filename=csv_filename)
        self.saved_paths["master_csv"] = path
        return path

    # ------------------------------------------------------------------
    # Main save entry point
    # ------------------------------------------------------------------
    def save_all(
        self, model, config, history, final_metrics, tokenizer=None, input_example=None
    ) -> None:
        """Save everything selected in the config."""

        self._save_config(config)
        self._save_history(history)
        self._save_final_metrics(final_metrics)
        self._save_metrics_log(history)
        self.save_to_sqlite(config=config, metrics=final_metrics)
        if self.cfg.save_model and self.should_save_heavy(final_metrics):
            self._save_model(model, tokenizer)
        if self.cfg.save_checkpoints:
            self._save_checkpoint(model, "final")
        if self.cfg.save_plots:
            self._generate_plots(history, final_metrics)
        if self.cfg.save_model and self.should_save_heavy(final_metrics):
            for fmt in self.cfg.model_export_formats:
                self._export_format(fmt, model, input_example)
        self._save_archive()
        self._save_model_card(config, final_metrics)

        # Automatically track to experiment trackers (wandb, mlflow, etc.)
        if self.cfg.experiment_trackers:
            # Delegate to the unified tracking module if available, else use legacy backends
            try:
                from neuralops.tracking import ExperimentTracker

                tracker = ExperimentTracker(
                    backend=self.cfg.experiment_trackers[0],  # supports one primary
                    log_dir=str(self.run_dir / "tracking"),
                    experiment_name=self.run_id,
                )
                for k, v in final_metrics.items():
                    tracker.log_metric(k, v)
                tracker.log_params(self._flatten_config(config))
                tracker.finalize()
            except ImportError:
                self.track_all(final_metrics, self._flatten_config(config))

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------
    def _write_json(self, rel_path: str, obj) -> None:
        """write json."""

        def _default_serializer(o):
            if hasattr(o, "to_dict") and callable(o.to_dict):
                return o.to_dict()
            if hasattr(o, "__dict__"):
                return {k: v for k, v in o.__dict__.items() if not k.startswith("_")}
            if isinstance(o, (datetime, Path)):
                return str(o)
            if hasattr(o, "value"):
                return o.value
            return str(o)

        fp = self.run_dir / rel_path
        fp.parent.mkdir(exist_ok=True)
        with open(fp, "w") as f:
            json.dump(
                _sanitize_json_value(obj),
                f,
                indent=2,
                default=_default_serializer,
                allow_nan=False,
            )

    def _write_text(self, rel_path: str, text: str) -> None:
        """write text."""
        fp = self.run_dir / rel_path
        fp.parent.mkdir(exist_ok=True)
        with open(fp, "w") as f:
            f.write(text)

    def _flatten_config(self, config) -> dict[str, Any]:
        """flatten config."""
        if hasattr(config, "__dict__"):
            return vars(config)
        return config

    def _save_config(self, config) -> None:
        """save config."""
        self._write_json("artifacts/config.json", self._flatten_config(config))

    def _save_history(self, history) -> None:
        """save history."""
        self._write_json("artifacts/history.json", history)

    def _save_final_metrics(self, metrics) -> None:
        """save final metrics."""
        self._write_json("artifacts/final_metrics.json", metrics)

    def _save_metrics_log(self, history) -> None:
        """save metrics log."""
        with open(self.run_dir / "logs" / "metrics.jsonl", "w") as f:
            for entry in history:
                f.write(json.dumps(_sanitize_json_value(entry), allow_nan=False) + "\n")

    def _save_model(self, model, tokenizer=None) -> None:
        """save model."""
        p = self.run_dir / "artifacts" / "model"
        p.mkdir(exist_ok=True)
        model.save_pretrained(str(p))
        if tokenizer:
            tokenizer.save_pretrained(str(p))
        if self.cfg.save_optimizer:
            import torch

            torch.save(model.state_dict(), p / "optimizer.pt")

    def _save_checkpoint(self, model, name="best") -> None:
        """save checkpoint."""
        p = self.run_dir / "checkpoints" / name
        p.mkdir(exist_ok=True)
        model.save_pretrained(str(p))

    def _save_archive(self) -> None:
        """save archive safely."""
        try:
            with tarfile.open(self.run_dir / "results.tar.gz", "w:gz") as tar:
                for sub in ["artifacts", "logs", "plots"]:
                    sp = self.run_dir / sub
                    if sp.exists():
                        tar.add(sp, arcname=sub)
        except Exception as e:
            logger.warning(f"Archive creation skipped due to warning/error: {e}")

    # ------------------------------------------------------------------
    # Plots – delegates to framework.analysis.plotting
    # ------------------------------------------------------------------
    def _generate_plots(self, history, final_metrics) -> None:
        """generate plots."""
        try:
            import matplotlib.pyplot as plt

            from neuralops.analysis.plotting import (
                plot_accuracy_curves,
                plot_confusion_matrix,
                plot_loss_curves,
                plot_lr_schedule,
                plot_radar_chart,
                plot_rank_allocation,
                plot_training_time,
            )

            opts = self.cfg.plot_options
            dpi = opts.get("dpi", 150)

            def save_fig(name, fig) -> None:
                """Persist fig to disk."""
                fig.savefig(self.run_dir / "plots" / f"{name}.png", dpi=dpi, bbox_inches="tight")
                plt.close(fig)

            if opts.get("loss_curve", True):
                try:
                    plot_loss_curves(history, final_metrics, self.cfg, self.run_dir / "plots", opts)
                except Exception as e:
                    logger.debug(f"plot_loss_curves error: {e}")

            if opts.get("accuracy_curve", True):
                try:
                    plot_accuracy_curves(
                        history, final_metrics, self.cfg, self.run_dir / "plots", opts
                    )
                except Exception as e:
                    logger.debug(f"plot_accuracy_curves error: {e}")

            if opts.get("lr_schedule", True):
                try:
                    plot_lr_schedule(history, final_metrics, self.cfg, self.run_dir / "plots", opts)
                except Exception as e:
                    logger.debug(f"plot_lr_schedule error: {e}")

            if opts.get("training_time", True):
                try:
                    plot_training_time(
                        history, final_metrics, self.cfg, self.run_dir / "plots", opts
                    )
                except Exception as e:
                    logger.debug(f"plot_training_time error: {e}")

            if opts.get("confusion_matrix", False):
                preds_path = self.run_dir / "artifacts" / "eval_predictions.json"
                if preds_path.exists():
                    with open(preds_path) as f:
                        data = json.load(f)
                    fig = plot_confusion_matrix(
                        data["labels"],
                        data["predictions"],
                        class_names=opts.get("class_names"),
                        figsize=tuple(opts.get("figsize", [10, 8])),
                    )
                    save_fig("confusion_matrix", fig)

            if opts.get("radar_chart", False):
                fig = plot_radar_chart(final_metrics, figsize=tuple(opts.get("figsize", [8, 8])))
                save_fig("radar_chart", fig)

            if opts.get("rank_allocation", False):
                rank_path = self.run_dir / "artifacts" / "rank_allocation.json"
                if rank_path.exists():
                    with open(rank_path) as f:
                        rank_data = json.load(f)
                    fig = plot_rank_allocation(
                        rank_data, figsize=tuple(opts.get("figsize", [10, 4]))
                    )
                    save_fig("rank_allocation", fig)

            if opts.get("attention_heatmap", False):
                logger.info(
                    "Attention heatmap requested but model not available in saver context; "
                    "use framework.callbacks.AttentionHeatmapCallback during training."
                )

        except Exception:
            logger.exception("Plot generation failed")

    # ------------------------------------------------------------------
    # Model card
    # ------------------------------------------------------------------
    def _save_model_card(self, config, metrics) -> None:
        """save model card."""
        cfg_dict = self._flatten_config(config) if hasattr(config, "__dict__") else config
        model_name = cfg_dict.get("model_name", "Model")
        task_name = cfg_dict.get("task_name", "Task")
        acc = metrics.get("eval_accuracy", "N/A")
        loss = metrics.get("eval_loss", "N/A")
        card = f"""# {model_name} – {task_name}

**Accuracy:** {acc}
**Loss:** {loss}

**Run ID:** {self.run_id}
**Date:** {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
**Compression:** {cfg_dict.get("attention_type", "full")} (rank={cfg_dict.get("rank", "N/A")})
"""
        self._write_text("artifacts/model_card.md", card)

    # ------------------------------------------------------------------
    # Export formats
    # ------------------------------------------------------------------
    def _export_format(self, fmt: str, model, input_example) -> None:
        """export format."""
        try:
            from .backends.exports import EXPORTERS

            exporter = EXPORTERS.get(fmt)
            if exporter:
                exporter(model, input_example, str(self.run_dir / "artifacts"))
        except Exception:
            logger.exception(f"Export {fmt} failed")

    # ------------------------------------------------------------------
    # Push to storage backends
    # ------------------------------------------------------------------
    def push_all(self) -> None:
        """Push all."""
        for be in self.cfg.storage_backends:
            self._push_backend(be)

    def _push_backend(self, name: str) -> None:
        """push backend."""
        try:
            if name == "local":
                return
            mod = __import__(f"neuralops.results_saver.backends.{name}", fromlist=["push"])
            mod.push(self.run_dir, self.cfg)
        except ImportError:
            logger.warning(f"Backend '{name}' not available, skipping")
        except Exception:
            logger.exception(f"Push to {name} failed")

    # ------------------------------------------------------------------
    # Git / DVC
    # ------------------------------------------------------------------
    def git_commit_and_push(self) -> None:
        """Git commit and push."""
        if not self.cfg.git_options.get("commit_results"):
            return
        from .backends.git import git_commit_and_push

        git_commit_and_push(self.run_dir, self.cfg)

    def dvc_track(self) -> None:
        """Dvc track."""
        if not self.cfg.dvc_options.get("track_model"):
            return
        from .backends.dvc import dvc_track

        dvc_track(self.run_dir, self.cfg)

    # ------------------------------------------------------------------
    # Experiment trackers (legacy backends – prefer framework.tracking)
    # ------------------------------------------------------------------
    def track_all(self, metrics: dict[str, Any], params: dict[str, Any]) -> None:
        """Track all."""
        for t in self.cfg.experiment_trackers:
            try:
                mod = __import__(
                    f"neuralops.results_saver.backends.{t}_tracker", fromlist=["track"]
                )
                mod.track(self.run_id, self.run_dir, metrics, params, self.cfg)
            except ImportError:
                logger.warning(f"Tracker '{t}' not available, skipping")
            except Exception:
                logger.exception(f"Tracking to {t} failed")

    # ------------------------------------------------------------------
    # Model registries
    # ------------------------------------------------------------------
    def push_to_registries(self) -> None:
        """Push to registries."""
        for r in self.cfg.model_registries:
            try:
                mod = __import__(f"neuralops.results_saver.backends.{r}", fromlist=["push"])
                mod.push(self.run_id, self.run_dir, self.cfg)
            except ImportError:
                logger.warning(f"Registry '{r}' not available")
            except Exception:
                logger.exception(f"Push to registry {r} failed")

    # ------------------------------------------------------------------
    # GitHub Release
    # ------------------------------------------------------------------
    def create_github_release(self, metrics: dict[str, Any]) -> str | None:
        """Build and return github release."""
        if not self.cfg.release_options.get("create_github_release"):
            return None
        from .backends.github_release import create_release

        return create_release(self.run_id, self.run_dir, metrics, self.cfg)

    # ------------------------------------------------------------------
    # Notifications
    # ------------------------------------------------------------------
    def notify(self, success: bool, metrics: dict[str, Any]) -> None:
        """Notify."""
        from .backends.notification import notify

        notify(self.run_id, success, metrics, self.cfg)

Usage -----= saver = ResultsSaver(run_id, base_dir, save_config_dict) saver.save_all(model, config, history, final_metrics, tokenizer, input_example) saver.push_all()

Subclasses

• ResultsSaver

Methods

def create_github_release(self, metrics: dict[str, Any]) ‑> str | None

Expand source code

def create_github_release(self, metrics: dict[str, Any]) -> str | None:
    """Build and return github release."""
    if not self.cfg.release_options.get("create_github_release"):
        return None
    from .backends.github_release import create_release

    return create_release(self.run_id, self.run_dir, metrics, self.cfg)

Build and return github release.

def dvc_track(self) ‑> None

Expand source code

def dvc_track(self) -> None:
    """Dvc track."""
    if not self.cfg.dvc_options.get("track_model"):
        return
    from .backends.dvc import dvc_track

    dvc_track(self.run_dir, self.cfg)

Dvc track.

def git_commit_and_push(self) ‑> None

Expand source code

def git_commit_and_push(self) -> None:
    """Git commit and push."""
    if not self.cfg.git_options.get("commit_results"):
        return
    from .backends.git import git_commit_and_push

    git_commit_and_push(self.run_dir, self.cfg)

Git commit and push.

def log_metric_to_sqlite(self, step: int, metric_name: str, value: float, db_path: str | Path | None = None) ‑> None

Expand source code

def log_metric_to_sqlite(
    self,
    step: int,
    metric_name: str,
    value: float,
    db_path: str | Path | None = None,
) -> None:
    """
    Append a single step metric into SQLite database with multi-GPU concurrency safety.

    # FIX: Uses PRAGMA journal_mode=WAL; and PRAGMA synchronous=NORMAL; with file locking to ensure safe step metric logging.
    """
    target_db = Path(db_path) if db_path is not None else (self.run_dir / "metrics.db")
    lock_path = target_db.with_suffix(".db.lock")

    with sqlite_file_lock(lock_path):
        conn = self._get_sqlite_connection(target_db)
        try:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS metric_logs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    run_id TEXT,
                    step INTEGER,
                    metric_name TEXT,
                    value REAL,
                    timestamp TEXT
                )
                """)
            cursor.execute(
                """
                INSERT INTO metric_logs (run_id, step, metric_name, value, timestamp)
                VALUES (?, ?, ?, ?, ?)
                """,
                (
                    self.run_id,
                    step,
                    metric_name,
                    float(value),
                    datetime.now().isoformat(),
                ),
            )
            conn.commit()
        finally:
            conn.close()

Append a single step metric into SQLite database with multi-GPU concurrency safety.

FIX: Uses PRAGMA journal_mode=WAL; and PRAGMA synchronous=NORMAL; with file locking to ensure safe step metric logging.

def notify(self, success: bool, metrics: dict[str, Any]) ‑> None

Expand source code

def notify(self, success: bool, metrics: dict[str, Any]) -> None:
    """Notify."""
    from .backends.notification import notify

    notify(self.run_id, success, metrics, self.cfg)

Notify.

def push_all(self) ‑> None

Expand source code

def push_all(self) -> None:
    """Push all."""
    for be in self.cfg.storage_backends:
        self._push_backend(be)

Push all.

def push_to_registries(self) ‑> None

Expand source code

def push_to_registries(self) -> None:
    """Push to registries."""
    for r in self.cfg.model_registries:
        try:
            mod = __import__(f"neuralops.results_saver.backends.{r}", fromlist=["push"])
            mod.push(self.run_id, self.run_dir, self.cfg)
        except ImportError:
            logger.warning(f"Registry '{r}' not available")
        except Exception:
            logger.exception(f"Push to registry {r} failed")

Push to registries.

def save_all(self, model, config, history, final_metrics, tokenizer=None, input_example=None) ‑> None

Expand source code

def save_all(
    self, model, config, history, final_metrics, tokenizer=None, input_example=None
) -> None:
    """Save everything selected in the config."""

    self._save_config(config)
    self._save_history(history)
    self._save_final_metrics(final_metrics)
    self._save_metrics_log(history)
    self.save_to_sqlite(config=config, metrics=final_metrics)
    if self.cfg.save_model and self.should_save_heavy(final_metrics):
        self._save_model(model, tokenizer)
    if self.cfg.save_checkpoints:
        self._save_checkpoint(model, "final")
    if self.cfg.save_plots:
        self._generate_plots(history, final_metrics)
    if self.cfg.save_model and self.should_save_heavy(final_metrics):
        for fmt in self.cfg.model_export_formats:
            self._export_format(fmt, model, input_example)
    self._save_archive()
    self._save_model_card(config, final_metrics)

    # Automatically track to experiment trackers (wandb, mlflow, etc.)
    if self.cfg.experiment_trackers:
        # Delegate to the unified tracking module if available, else use legacy backends
        try:
            from neuralops.tracking import ExperimentTracker

            tracker = ExperimentTracker(
                backend=self.cfg.experiment_trackers[0],  # supports one primary
                log_dir=str(self.run_dir / "tracking"),
                experiment_name=self.run_id,
            )
            for k, v in final_metrics.items():
                tracker.log_metric(k, v)
            tracker.log_params(self._flatten_config(config))
            tracker.finalize()
        except ImportError:
            self.track_all(final_metrics, self._flatten_config(config))

Save everything selected in the config.

def save_config(self, config: Any) ‑> str

Expand source code

def save_config(self, config: Any) -> str:
    """Save experiment configuration dictionary or object."""
    self._save_config(config)
    self.save_to_sqlite(config=config)
    path = str(self.run_dir / "artifacts" / "config.json")
    self.saved_paths["config"] = path
    return path

Save experiment configuration dictionary or object.

def save_metrics(self, metrics: dict[str, Any]) ‑> str

Expand source code

def save_metrics(self, metrics: dict[str, Any]) -> str:
    """Save a dictionary of evaluation/training metrics to JSON and SQLite."""
    self._save_final_metrics(metrics)
    self.save_to_sqlite(metrics=metrics)
    path = str(self.run_dir / "artifacts" / "final_metrics.json")
    self.saved_paths["metrics"] = path
    return path

Save a dictionary of evaluation/training metrics to JSON and SQLite.

def save_model_weights(self, model: Any, filename: str = 'model.pt') ‑> str

Expand source code

def save_model_weights(self, model: Any, filename: str = "model.pt") -> str:
    """Save PyTorch model state dict or checkpoint."""
    import torch

    art_dir = self.run_dir / "artifacts"
    art_dir.mkdir(parents=True, exist_ok=True)
    dest = art_dir / filename
    if hasattr(model, "state_dict"):
        torch.save(model.state_dict(), dest)
    elif hasattr(model, "save_pretrained"):
        model.save_pretrained(str(art_dir / "model"))
        return str(art_dir / "model")
    else:
        torch.save(model, dest)
    self.saved_paths["model_weights"] = str(dest)
    return str(dest)

Save PyTorch model state dict or checkpoint.

def save_result(self,
result: dict[str, Any] | Any,
results_dir: str | Path | None = None,
csv_filename: str = 'master_results.csv') ‑> str

Expand source code

def save_result(
    self,
    result: dict[str, Any] | Any,
    results_dir: str | Path | None = None,
    csv_filename: str = "master_results.csv",
) -> str:
    """Append a benchmark result to master CSV."""
    target_dir = results_dir if results_dir is not None else self.base_dir
    path = save_result(result, results_dir=target_dir, csv_filename=csv_filename)
    self.saved_paths["master_csv"] = path
    return path

Append a benchmark result to master CSV.

def save_to_sqlite(self,
db_path: str | Path | None = None,
config: Any | None = None,
metrics: dict[str, Any] | None = None,
extra_metadata: dict[str, Any] | None = None) ‑> str

Expand source code

def save_to_sqlite(
    self,
    db_path: str | Path | None = None,
    config: Any | None = None,
    metrics: dict[str, Any] | None = None,
    extra_metadata: dict[str, Any] | None = None,
) -> str:
    """
    Save run metadata, configuration, and final metrics into SQLite database with multi-GPU write safety.

    # FIX: Uses PRAGMA journal_mode=WAL;, PRAGMA synchronous=NORMAL;, and file locking for multi-GPU database write safety.
    """
    target_db = Path(db_path) if db_path is not None else (self.base_dir / "experiments.db")
    lock_path = target_db.with_suffix(".db.lock")

    cfg_dict = self._flatten_config(config) if config is not None else {}
    metrics_dict = _sanitize_json_value(metrics or {})
    meta_dict = _sanitize_json_value(extra_metadata or {})

    # FIX: Wrap database transaction in cross-process file lock
    with sqlite_file_lock(lock_path):
        conn = self._get_sqlite_connection(target_db)
        try:
            cursor = conn.cursor()
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS experiment_runs (
                    run_id TEXT PRIMARY KEY,
                    created_at TEXT,
                    task_name TEXT,
                    model_name TEXT,
                    config_json TEXT,
                    metrics_json TEXT,
                    metadata_json TEXT
                )
                """)
            now_str = datetime.now().isoformat()
            task_name = str(cfg_dict.get("task_name", "unknown"))
            model_name = str(cfg_dict.get("model_name", "unknown"))

            cursor.execute(
                """
                INSERT OR REPLACE INTO experiment_runs
                (run_id, created_at, task_name, model_name, config_json, metrics_json, metadata_json)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    self.run_id,
                    now_str,
                    task_name,
                    model_name,
                    json.dumps(_sanitize_json_value(cfg_dict), allow_nan=False),
                    json.dumps(metrics_dict, allow_nan=False),
                    json.dumps(meta_dict, allow_nan=False),
                ),
            )
            conn.commit()
        finally:
            conn.close()

    self.saved_paths["sqlite_db"] = str(target_db)
    return str(target_db)

Save run metadata, configuration, and final metrics into SQLite database with multi-GPU write safety.

FIX: Uses PRAGMA journal_mode=WAL;, PRAGMA synchronous=NORMAL;, and file locking for multi-GPU database write safety.

def should_save_heavy(self, final_metrics: dict[str, Any]) ‑> bool

Expand source code

def should_save_heavy(self, final_metrics: dict[str, Any]) -> bool:
    """Evaluate if metric meets auto_save_if_better_than threshold."""
    threshold = self.cfg.auto_save_if_better_than
    if threshold is None:
        return True
    metric = final_metrics.get("eval_accuracy", 0)
    return metric >= threshold

Evaluate if metric meets auto_save_if_better_than threshold.

def track_all(self, metrics: dict[str, Any], params: dict[str, Any]) ‑> None

Expand source code

def track_all(self, metrics: dict[str, Any], params: dict[str, Any]) -> None:
    """Track all."""
    for t in self.cfg.experiment_trackers:
        try:
            mod = __import__(
                f"neuralops.results_saver.backends.{t}_tracker", fromlist=["track"]
            )
            mod.track(self.run_id, self.run_dir, metrics, params, self.cfg)
        except ImportError:
            logger.warning(f"Tracker '{t}' not available, skipping")
        except Exception:
            logger.exception(f"Tracking to {t} failed")

Track all.

class RetrainDecision (should_retrain: bool,
reason: str,
trigger_type: TriggerType | None = None,
urgency: str = 'low',
details: dict[str, typing.Any] = <factory>)

Expand source code

@dataclass
class RetrainDecision:
    """Output of the retraining decision engine."""

    should_retrain: bool
    reason: str
    trigger_type: TriggerType | None = None
    urgency: str = "low"  # "low", "medium", "high", "critical"
    details: dict[str, Any] = field(default_factory=dict)

Output of the retraining decision engine.

Instance variables

var details : dict[str, typing.Any]

The type of the None singleton.

var reason : str

The type of the None singleton.

var should_retrain : bool

The type of the None singleton.

var trigger_type : TriggerType | None

The type of the None singleton.

var urgency : str

The type of the None singleton.

class RetrainPolicy (trigger_types: list[TriggerType] = <factory>,
drift_threshold: float = 0.3,
drift_methods: list[str] = <factory>,
min_accuracy: float = 0.85,
max_accuracy_degradation: float = 0.05,
schedule_cron: str | None = None,
max_days_since_last_train: int = 30,
min_new_samples: int = 1000,
max_new_samples: int = 100000,
require_better_model: bool = True,
min_improvement: float = 0.01,
max_retrains_per_day: int = 3,
validation_split: float = 0.2,
notify_on_trigger: bool = True,
notify_on_complete: bool = True)

Expand source code

@dataclass
class RetrainPolicy:
    """Configuration for when and how to retrain."""

    trigger_types: list[TriggerType] = field(default_factory=lambda: [TriggerType.DRIFT])

    # Drift trigger
    drift_threshold: float = 0.3
    drift_methods: list[str] = field(default_factory=lambda: ["ks", "psi"])

    # Performance trigger
    min_accuracy: float = 0.85
    max_accuracy_degradation: float = 0.05

    # Schedule trigger
    schedule_cron: str | None = None  # "0 2 * * 1" for every Monday 2 AM
    max_days_since_last_train: int = 30

    # Data volume trigger
    min_new_samples: int = 1000
    max_new_samples: int = 100_000

    # Retraining behaviour
    require_better_model: bool = True
    min_improvement: float = 0.01
    max_retrains_per_day: int = 3
    validation_split: float = 0.2
    notify_on_trigger: bool = True
    notify_on_complete: bool = True

Configuration for when and how to retrain.

Instance variables

var drift_methods : list[str]

The type of the None singleton.

var drift_threshold : float

The type of the None singleton.

var max_accuracy_degradation : float

The type of the None singleton.

var max_days_since_last_train : int

The type of the None singleton.

var max_new_samples : int

The type of the None singleton.

var max_retrains_per_day : int

The type of the None singleton.

var min_accuracy : float

The type of the None singleton.

var min_improvement : float

The type of the None singleton.

var min_new_samples : int

The type of the None singleton.

var notify_on_complete : bool

The type of the None singleton.

var notify_on_trigger : bool

The type of the None singleton.

var require_better_model : bool

The type of the None singleton.

var schedule_cron : str | None

The type of the None singleton.

var trigger_types : list[TriggerType]

The type of the None singleton.

var validation_split : float

The type of the None singleton.

class RetrainingOrchestrator (policy: RetrainPolicy = None)

Expand source code

class RetrainingOrchestrator:
    """
    Central decision engine for automated retraining.

    Usage:
        orch = RetrainingOrchestrator(RetrainPolicy())
        decision = orch.evaluate(
            current_metrics={"accuracy": 0.82},
            reference_metrics={"accuracy": 0.88},
            drift_report=drift_report,
            new_data_count=5000,
            days_since_last_train=35,
        )
        if decision.should_retrain:
            orch.trigger_retraining(run_id="model-v3", training_fn=my_train_func)
    """

    def __init__(self, policy: RetrainPolicy = None) -> None:
        self.policy = policy or RetrainPolicy()
        self._retrain_count_today: int = 0
        self._last_retrain_date: str = ""

    def evaluate(
        self,
        current_metrics: dict[str, float] | None = None,
        reference_metrics: dict[str, float] | None = None,
        drift_report: Any | None = None,
        new_data_count: int = 0,
        days_since_last_train: int = 0,
        manual_request: bool = False,
    ) -> RetrainDecision:
        """Evaluate all configured triggers and return a retraining decision."""

        # Reset daily counter if it's a new day
        today = datetime.now().strftime("%Y-%m-%d")
        if today != self._last_retrain_date:
            self._retrain_count_today = 0
            self._last_retrain_date = today

        if self._retrain_count_today >= self.policy.max_retrains_per_day:
            return RetrainDecision(
                should_retrain=False,
                reason="Max daily retrains reached",
                details={"retrains_today": self._retrain_count_today},
            )

        # 1. Manual trigger
        if manual_request and TriggerType.MANUAL in self.policy.trigger_types:
            return RetrainDecision(
                should_retrain=True,
                reason="Manual retraining requested",
                trigger_type=TriggerType.MANUAL,
                urgency="high",
            )

        # 2. Performance degradation
        if (
            TriggerType.PERFORMANCE in self.policy.trigger_types
            and current_metrics
            and reference_metrics
        ):
            decision = self._check_performance(current_metrics, reference_metrics)
            if decision.should_retrain:
                return decision

        # 3. Drift trigger
        if TriggerType.DRIFT in self.policy.trigger_types and drift_report:
            decision = self._check_drift(drift_report)
            if decision.should_retrain:
                return decision

        # 4. Data volume trigger
        if TriggerType.DATA_VOLUME in self.policy.trigger_types:
            decision = self._check_data_volume(new_data_count)
            if decision.should_retrain:
                return decision

        # 5. Schedule trigger
        if TriggerType.SCHEDULE in self.policy.trigger_types:
            decision = self._check_schedule(days_since_last_train)
            if decision.should_retrain:
                return decision

        return RetrainDecision(should_retrain=False, reason="No triggers met")

    def _check_performance(
        self, current: dict[str, float], reference: dict[str, float]
    ) -> RetrainDecision:
        """check performance."""
        acc_key = "accuracy" if "accuracy" in current else "eval_accuracy"
        curr_acc = current.get(acc_key, 1.0)
        ref_acc = reference.get(acc_key, 1.0)
        degradation = ref_acc - curr_acc

        if curr_acc < self.policy.min_accuracy:
            return RetrainDecision(
                should_retrain=True,
                reason=f"Accuracy {curr_acc:.3f} below minimum {self.policy.min_accuracy}",
                trigger_type=TriggerType.PERFORMANCE,
                urgency="critical" if degradation > 0.1 else "high",
                details={
                    "current_accuracy": curr_acc,
                    "reference_accuracy": ref_acc,
                    "degradation": degradation,
                },
            )

        if degradation > self.policy.max_accuracy_degradation:
            return RetrainDecision(
                should_retrain=True,
                reason=f"Accuracy degraded by {degradation:.3f} (threshold: {self.policy.max_accuracy_degradation})",
                trigger_type=TriggerType.PERFORMANCE,
                urgency="medium",
                details={"degradation": degradation},
            )

        return RetrainDecision(should_retrain=False, reason="Performance within acceptable range")

    def _check_drift(self, drift_report) -> RetrainDecision:
        """check drift."""
        n_drifted = getattr(drift_report, "n_features_drifted", 0)
        n_total = getattr(drift_report, "n_features_tested", 1)
        ratio = n_drifted / max(n_total, 1)

        if ratio > self.policy.drift_threshold:
            return RetrainDecision(
                should_retrain=True,
                reason=f"Drift detected in {n_drifted}/{n_total} features ({ratio:.1%})",
                trigger_type=TriggerType.DRIFT,
                urgency="high" if ratio > 0.5 else "medium",
                details={"drifted_features": n_drifted, "total_features": n_total},
            )
        return RetrainDecision(should_retrain=False, reason="Drift within acceptable range")

    def _check_data_volume(self, new_count: int) -> RetrainDecision:
        """check data volume."""
        if new_count >= self.policy.min_new_samples:
            return RetrainDecision(
                should_retrain=True,
                reason=f"New data volume ({new_count}) exceeded threshold ({self.policy.min_new_samples})",
                trigger_type=TriggerType.DATA_VOLUME,
                urgency="low",
                details={"new_samples": new_count},
            )
        return RetrainDecision(should_retrain=False, reason="Insufficient new data")

    def _check_schedule(self, days_since: int) -> RetrainDecision:
        """check schedule."""
        if days_since >= self.policy.max_days_since_last_train:
            return RetrainDecision(
                should_retrain=True,
                reason=f"Scheduled retrain: {days_since} days since last training (max: {self.policy.max_days_since_last_train})",
                trigger_type=TriggerType.SCHEDULE,
                urgency="low",
                details={"days_since_last_train": days_since},
            )
        return RetrainDecision(should_retrain=False, reason="Schedule not due yet")

    def trigger_retraining(
        self,
        run_id: str,
        training_fn: Callable,
        decision: RetrainDecision,
        **training_kwargs,
    ) -> dict[str, Any]:
        """Execute retraining and validate the new model."""
        logger.info("Retraining triggered for %s: %s", run_id, decision.reason)
        self._retrain_count_today += 1

        try:
            result = training_fn(**training_kwargs)
            logger.info(
                "Retraining completed for %s. Metrics: %s",
                run_id,
                {k: v for k, v in result.items() if isinstance(v, (int, float))},
            )
        except Exception as e:
            logger.exception("Retraining failed for %s: %s", run_id, e)
            return {"status": "failed", "error": str(e)}

        return {"status": "completed", "run_id": run_id, "metrics": result}

Central decision engine for automated retraining.

Usage -----= orch = RetrainingOrchestrator(RetrainPolicy()) decision = orch.evaluate( current_metrics={"accuracy": 0.82}, reference_metrics={"accuracy": 0.88}, drift_report=drift_report, new_data_count=5000, days_since_last_train=35, ) if decision.should_retrain: orch.trigger_retraining(run_id="model-v3", training_fn=my_train_func)

Methods

def evaluate(self,
current_metrics: dict[str, float] | None = None,
reference_metrics: dict[str, float] | None = None,
drift_report: Any | None = None,
new_data_count: int = 0,
days_since_last_train: int = 0,
manual_request: bool = False) ‑> RetrainDecision

Expand source code

def evaluate(
    self,
    current_metrics: dict[str, float] | None = None,
    reference_metrics: dict[str, float] | None = None,
    drift_report: Any | None = None,
    new_data_count: int = 0,
    days_since_last_train: int = 0,
    manual_request: bool = False,
) -> RetrainDecision:
    """Evaluate all configured triggers and return a retraining decision."""

    # Reset daily counter if it's a new day
    today = datetime.now().strftime("%Y-%m-%d")
    if today != self._last_retrain_date:
        self._retrain_count_today = 0
        self._last_retrain_date = today

    if self._retrain_count_today >= self.policy.max_retrains_per_day:
        return RetrainDecision(
            should_retrain=False,
            reason="Max daily retrains reached",
            details={"retrains_today": self._retrain_count_today},
        )

    # 1. Manual trigger
    if manual_request and TriggerType.MANUAL in self.policy.trigger_types:
        return RetrainDecision(
            should_retrain=True,
            reason="Manual retraining requested",
            trigger_type=TriggerType.MANUAL,
            urgency="high",
        )

    # 2. Performance degradation
    if (
        TriggerType.PERFORMANCE in self.policy.trigger_types
        and current_metrics
        and reference_metrics
    ):
        decision = self._check_performance(current_metrics, reference_metrics)
        if decision.should_retrain:
            return decision

    # 3. Drift trigger
    if TriggerType.DRIFT in self.policy.trigger_types and drift_report:
        decision = self._check_drift(drift_report)
        if decision.should_retrain:
            return decision

    # 4. Data volume trigger
    if TriggerType.DATA_VOLUME in self.policy.trigger_types:
        decision = self._check_data_volume(new_data_count)
        if decision.should_retrain:
            return decision

    # 5. Schedule trigger
    if TriggerType.SCHEDULE in self.policy.trigger_types:
        decision = self._check_schedule(days_since_last_train)
        if decision.should_retrain:
            return decision

    return RetrainDecision(should_retrain=False, reason="No triggers met")

Evaluate all configured triggers and return a retraining decision.

def trigger_retraining(self,
run_id: str,
training_fn: ,
decision: RetrainDecision,
**training_kwargs) ‑> dict[str, typing.Any]

Expand source code

def trigger_retraining(
    self,
    run_id: str,
    training_fn: Callable,
    decision: RetrainDecision,
    **training_kwargs,
) -> dict[str, Any]:
    """Execute retraining and validate the new model."""
    logger.info("Retraining triggered for %s: %s", run_id, decision.reason)
    self._retrain_count_today += 1

    try:
        result = training_fn(**training_kwargs)
        logger.info(
            "Retraining completed for %s. Metrics: %s",
            run_id,
            {k: v for k, v in result.items() if isinstance(v, (int, float))},
        )
    except Exception as e:
        logger.exception("Retraining failed for %s: %s", run_id, e)
        return {"status": "failed", "error": str(e)}

    return {"status": "completed", "run_id": run_id, "metrics": result}

Execute retraining and validate the new model.

class TaskType (*values)

Expand source code

class TaskType(Enum):
    SINGLE_CLASSIFICATION = "single_classification"
    PAIR_CLASSIFICATION = "pair_classification"
    REGRESSION = "regression"
    QA = "qa"
    TOKEN_CLASSIFICATION = "token_classification"
    LANGUAGE_MODELING = "language_modeling"
    SEQ2SEQ = "seq2seq"

Create a collection of name/value pairs.

Example enumeration:

>>> class Color(Enum):
...     RED = 1
...     BLUE = 2
...     GREEN = 3

Access them by:

• attribute access:

Color.RED

• value lookup:

Color(1)

• name lookup:

Color['RED']

Enumerations can be iterated over, and know how many members they have:

>>> len(Color)
3

>>> list(Color)
[<Color.RED: 1>, <Color.BLUE: 2>, <Color.GREEN: 3>]

Methods can be added to enumerations, and members can have their own attributes – see the documentation for details.

Ancestors

• enum.Enum

Class variables

var LANGUAGE_MODELING

The type of the None singleton.

var PAIR_CLASSIFICATION

The type of the None singleton.

var QA

The type of the None singleton.

var REGRESSION

The type of the None singleton.

var SEQ2SEQ

The type of the None singleton.

var SINGLE_CLASSIFICATION

The type of the None singleton.

var TOKEN_CLASSIFICATION

The type of the None singleton.

class TrainerRegistry

Expand source code

class TrainerRegistry:
    """Registry for trainer classes."""

    _registry: dict[str, type[BaseTrainer]] = {}

    @classmethod
    def register(cls, name: str) -> Any:
        """Register."""

        def decorator(klass: type[BaseTrainer]) -> type[BaseTrainer]:
            """Decorator."""
            if not issubclass(klass, BaseTrainer):
                raise TypeError(f"{klass.__name__} must inherit from BaseTrainer")
            cls._registry[name] = klass
            logger.info(f"Registered trainer: {name} -> {klass.__name__}")
            return klass

        return decorator

    @classmethod
    def get(cls, name: str) -> type[BaseTrainer]:
        """Get."""
        if name not in cls._registry:
            raise ValueError(f"Unknown trainer '{name}'. Available: {list(cls._registry.keys())}")
        return cls._registry[name]

    @classmethod
    def list(cls) -> list:
        """List."""
        return list(cls._registry.keys())

    @classmethod
    def create(cls, name: str, **kwargs) -> BaseTrainer:
        """Create."""
        return cls.get(name)(**kwargs)

Registry for trainer classes.

Static methods

def create(name: str, **kwargs) ‑> BaseTrainer

Create.

def get(name: str) ‑> type[BaseTrainer]

Get.

def list() ‑> TrainerRegistry.list() at 0x12e97cb40>)>

List.

def register(name: str) ‑> Any

Register.

class UnifiedTransformerConfig (attention_type: str = None,
attention_rank: int | None = None,
attention_kwargs: dict[str, typing.Any] | None = None,
**kwargs)

Expand source code

class UnifiedTransformerConfig:
    """Lightweight container for the custom attention configuration."""

    def __init__(
        self,
        attention_type: str = None,
        attention_rank: int | None = None,
        attention_kwargs: dict[str, Any] | None = None,
        **kwargs,
    ) -> None:
        self.attention_type = attention_type
        self.attention_rank = attention_rank
        self.attention_kwargs = attention_kwargs or {}
        # Store any extra kwargs
        for k, v in kwargs.items():
            setattr(self, k, v)

    def to_dict(self) -> dict[str, Any]:
        """Serialise to a JSON‑compatible dict."""
        return {
            "attention_type": self.attention_type,
            "attention_rank": self.attention_rank,
            "attention_kwargs": self.attention_kwargs,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> "UnifiedTransformerConfig":
        """From dict."""
        return cls(**d)

Lightweight container for the custom attention configuration.

Static methods

def from_dict(d: dict[str, typing.Any]) ‑> UnifiedTransformerConfig

From dict.

Methods

def to_dict(self) ‑> dict[str, typing.Any]

Expand source code

def to_dict(self) -> dict[str, Any]:
    """Serialise to a JSON‑compatible dict."""
    return {
        "attention_type": self.attention_type,
        "attention_rank": self.attention_rank,
        "attention_kwargs": self.attention_kwargs,
    }

Serialise to a JSON‑compatible dict.

class UnifiedTransformerModel (model: torch.nn.modules.module.Module,
custom_config: UnifiedTransformerConfig | None = None)

Expand source code

class UnifiedTransformerModel(nn.Module, PyTorchModelHubMixin):
    """
    A wrapper around a HuggingFace model that supports custom attention
    and can be saved/loaded via the HuggingFace Hub.
    """

    def __init__(
        self, model: nn.Module, custom_config: UnifiedTransformerConfig | None = None
    ) -> None:
        super().__init__()
        self.model = model
        self.config = getattr(model, "config", None)  # HF model config
        self.arch = detect_architecture(model)
        self._custom_config = custom_config or UnifiedTransformerConfig()

    # ------------------------------------------------------------------
    # Attention replacement
    # ------------------------------------------------------------------
    def replace_attention(
        self,
        attention_type: str,
        rank: int | None = None,
        copy_weights: bool = True,
        **kwargs,
    ) -> Any:
        """
        Inject a custom attention variant into all layers.
        This method also records the choice so it can be saved.
        """
        inject_attention(self.model, attention_type, rank=rank, copy_weights=copy_weights, **kwargs)
        self._custom_config = UnifiedTransformerConfig(
            attention_type=attention_type,
            attention_rank=rank,
            attention_kwargs=kwargs,
        )
        return self

    # ------------------------------------------------------------------
    # Freezing / parameter count
    # ------------------------------------------------------------------
    def freeze_layers(self, num_layers_to_freeze: int = 0) -> Any:
        """Freeze layers."""
        layers = get_transformer_layers(self.model)
        for i, layer in enumerate(layers):
            for param in layer.parameters():
                param.requires_grad = i >= num_layers_to_freeze
        return self

    def unfreeze_all(self) -> Any:
        """Unfreeze all."""
        for param in self.model.parameters():
            param.requires_grad = True
        return self

    def get_parameter_count(self) -> dict[str, int]:
        """Return the parameter count."""
        total = sum(p.numel() for p in self.model.parameters())
        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
        return {"total": total, "trainable": trainable, "frozen": total - trainable}

    # ------------------------------------------------------------------
    # Hub serialization – these methods are called by PyTorchModelHubMixin
    # ------------------------------------------------------------------
    def _save_pretrained(self, save_directory: str) -> None:
        """
        Save the underlying model and our custom config to a directory.
        """
        # 1. Save the base HF model (attention layers are already inside)
        self.model.save_pretrained(save_directory)
        # 2. Save our custom attention config to a JSON file
        import json

        custom_cfg_path = f"{save_directory}/unified_attention_config.json"
        with open(custom_cfg_path, "w") as f:
            json.dump(self._custom_config.to_dict(), f, indent=2)

    @classmethod
    def _from_pretrained(cls, model_id: str, **kwargs) -> "UnifiedTransformerModel":
        """
        Load a pretrained model and reconstruct the custom attention variant.
        This method is invoked by `from_pretrained` class method of the mixin.
        """
        # 1. Load the custom config
        import json
        import os

        if os.path.isdir(model_id):
            cfg_path = os.path.join(model_id, "unified_attention_config.json")
        else:
            # If model_id is a Hub repo, we need to download the file
            from huggingface_hub import hf_hub_download

            cfg_path = hf_hub_download(repo_id=model_id, filename="unified_attention_config.json")

        with open(cfg_path) as f:
            custom_cfg_dict = json.load(f)
        custom_cfg = UnifiedTransformerConfig.from_dict(custom_cfg_dict)

        # 2. Load the base HF model
        base_model = AutoModel.from_pretrained(model_id, **kwargs)

        # 3. Re‑apply the attention replacement if a type was stored
        model = cls(base_model, custom_config=custom_cfg)
        if custom_cfg.attention_type:
            model.replace_attention(
                custom_cfg.attention_type,
                rank=custom_cfg.attention_rank,
                copy_weights=False,  # Weights already inside the saved base model
                **custom_cfg.attention_kwargs,
            )
        return model

    # ------------------------------------------------------------------
    # Other methods
    # ------------------------------------------------------------------
    def forward(self, *args, **kwargs) -> Any:
        """Forward pass — compute module output from inputs."""
        return self.model(*args, **kwargs)

    def save_pretrained(self, save_directory: str) -> None:
        """Convenience wrapper – delegates to the mixin."""
        super().save_pretrained(save_directory)

    @classmethod
    def from_pretrained(
        cls,
        model_id: str,
        task: str = "sequence_classification",
        num_labels: int | None = None,
        **kwargs,
    ) -> "UnifiedTransformerModel":
        """
        Create a new UnifiedTransformerModel from a pretrained identifier.
        Automatically applies the custom attention if the model was saved with one.
        This is the main entry point for users.
        """
        # If the model_id has a unified_attention_config.json, use the mixin's from_pretrained
        # which will call our _from_pretrained. Otherwise, create a fresh one.
        import os

        has_custom = False
        if os.path.isdir(model_id):
            has_custom = os.path.exists(os.path.join(model_id, "unified_attention_config.json"))

        if has_custom:
            return cls._from_pretrained(model_id, **kwargs)
        else:
            # Fresh model – no custom attention yet
            if task == "sequence_classification":
                sc_kwargs = dict(kwargs)
                if num_labels is not None:
                    sc_kwargs["num_labels"] = num_labels
                base = AutoModelForSequenceClassification.from_pretrained(model_id, **sc_kwargs)
            else:
                base = AutoModel.from_pretrained(model_id, **kwargs)
            return cls(base)

A wrapper around a HuggingFace model that supports custom attention and can be saved/loaded via the HuggingFace Hub.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module
• huggingface_hub.hub_mixin.PyTorchModelHubMixin
• huggingface_hub.hub_mixin.ModelHubMixin

Static methods

def from_pretrained(model_id: str,
task: str = 'sequence_classification',
num_labels: int | None = None,
**kwargs) ‑> UnifiedTransformerModel

Create a new UnifiedTransformerModel from a pretrained identifier. Automatically applies the custom attention if the model was saved with one. This is the main entry point for users.

Methods

def forward(self, *args, **kwargs) ‑> Any

Expand source code

def forward(self, *args, **kwargs) -> Any:
    """Forward pass — compute module output from inputs."""
    return self.model(*args, **kwargs)

Forward pass — compute module output from inputs.

def freeze_layers(self, num_layers_to_freeze: int = 0) ‑> Any

Expand source code

def freeze_layers(self, num_layers_to_freeze: int = 0) -> Any:
    """Freeze layers."""
    layers = get_transformer_layers(self.model)
    for i, layer in enumerate(layers):
        for param in layer.parameters():
            param.requires_grad = i >= num_layers_to_freeze
    return self

Freeze layers.

def get_parameter_count(self) ‑> dict[str, int]

Expand source code

def get_parameter_count(self) -> dict[str, int]:
    """Return the parameter count."""
    total = sum(p.numel() for p in self.model.parameters())
    trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
    return {"total": total, "trainable": trainable, "frozen": total - trainable}

Return the parameter count.

def replace_attention(self,
attention_type: str,
rank: int | None = None,
copy_weights: bool = True,
**kwargs) ‑> Any

Expand source code

def replace_attention(
    self,
    attention_type: str,
    rank: int | None = None,
    copy_weights: bool = True,
    **kwargs,
) -> Any:
    """
    Inject a custom attention variant into all layers.
    This method also records the choice so it can be saved.
    """
    inject_attention(self.model, attention_type, rank=rank, copy_weights=copy_weights, **kwargs)
    self._custom_config = UnifiedTransformerConfig(
        attention_type=attention_type,
        attention_rank=rank,
        attention_kwargs=kwargs,
    )
    return self

Inject a custom attention variant into all layers. This method also records the choice so it can be saved.

def save_pretrained(self, save_directory: str) ‑> None

Expand source code

def save_pretrained(self, save_directory: str) -> None:
    """Convenience wrapper – delegates to the mixin."""
    super().save_pretrained(save_directory)

Convenience wrapper – delegates to the mixin.

def unfreeze_all(self) ‑> Any

Expand source code

def unfreeze_all(self) -> Any:
    """Unfreeze all."""
    for param in self.model.parameters():
        param.requires_grad = True
    return self

Unfreeze all.