Module neuralops.experiment_utils

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

Sub-modules

neuralops.experiment_utils.comparison

Experiment Comparison Engine – side-by-side diff, parallel coordinates, metric delta tables, and run diff for training runs.

neuralops.experiment_utils.config

neuralops.experiment_utils.data

neuralops.experiment_utils.main

Efficient Attention Framework – GLUE Fine‑Tuning & Utility Functions

neuralops.experiment_utils.model

Module providing model implementations for NeuralOps.

neuralops.experiment_utils.orchestrator

MegaOrchestrator – single entry point that executes the full experiment pipeline. Uses the core framework for all ML logic.

neuralops.experiment_utils.runner

Deprecated shim for neuralops.experiment_utils.main …

neuralops.experiment_utils.tasks

Background training task (Django integration) with WebSocket broadcasting and AI assistant.

neuralops.experiment_utils.train

Functions

def broadcast_terminal(msg: str) ‑> None

Expand source code

def broadcast_terminal(msg: str) -> None:
    """Send a log message to the terminal WebSocket group."""
    if not HAS_CHANNELS:
        return
    try:
        channel_layer = get_channel_layer()
        if channel_layer and async_to_sync:
            async_to_sync(channel_layer.group_send)(
                "terminal", {"type": "terminal_message", "data": msg}
            )
    except Exception as e:
        print(f"WS broadcast error: {e}")

Send a log message to the terminal WebSocket group.

def compare_runs_api(run_ids: list[str], baseline_id: str | None = None) ‑> dict

Expand source code

def compare_runs_api(run_ids: list[str], baseline_id: str | None = None) -> dict:
    """API-friendly wrapper."""
    comparator = ExperimentComparator(run_ids, baseline_id)
    result = comparator.compare()
    return {
        "run_ids": result.run_ids,
        "metrics_table": result.metrics_table,
        "config_diff": result.config_diff,
        "hyperparameter_table": result.hyperparameter_table,
        "metric_deltas": result.metric_deltas,
        "timeline_data": result.timeline_data,
        "parallel_coords": result.parallel_coords,
        "best_run_id": result.best_run_id,
        "summary": result.summary,
    }

API-friendly wrapper.

def run_compression(config, output_dir) ‑> Any

Expand source code

def run_compression(config, output_dir) -> Any:
    """Run compression."""
    if isinstance(config, dict):
        cfg = data_ExperimentConfig()
        for k, v in config.items():
            setattr(cfg, k, v)
    else:
        cfg = config

    compressor_name = cfg.compressor or "svd"
    if compressor_name not in COMPRESSOR_REGISTRY:
        raise ValueError(f"Unknown compressor: {compressor_name}")

    set_seed(cfg.seed)
    device = get_device()

    model = AutoModelForSequenceClassification.from_pretrained(cfg.model_name)
    model.to(device)

    compressor_cls = COMPRESSOR_REGISTRY[compressor_name]
    compressor = compressor_cls(model, rank=cfg.compression_rank)
    compressed_model = compressor.compress()

    os.makedirs(output_dir, exist_ok=True)
    compressed_model.save_pretrained(output_dir)
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name)
    tokenizer.save_pretrained(output_dir)

    original_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    compressed_bytes = sum(p.numel() * p.element_size() for p in compressed_model.parameters())
    reduction_ratio = original_bytes / compressed_bytes if compressed_bytes > 0 else 1

    stats = {
        "compressor": compressor_name,
        "original_size_mb": round(original_bytes / (1024**2), 2),
        "compressed_size_mb": round(compressed_bytes / (1024**2), 2),
        "reduction_ratio": round(reduction_ratio, 2),
        "compression_rank": cfg.compression_rank,
    }
    logger.info(f"Compression complete: {stats}")
    return stats

Run compression.

def run_evaluation(config, model_path) ‑> Any

Expand source code

def run_evaluation(config, model_path) -> Any:
    """Run evaluation."""
    if isinstance(config, dict):
        cfg = data_ExperimentConfig()
        for k, v in config.items():
            setattr(cfg, k, v)
    else:
        cfg = config

    set_seed(cfg.seed)
    device = get_device()

    # Reload data (simplified)
    raw_datasets = load_glue_local(cfg.task, cfg.glue_base_path)
    num_labels, is_regression = get_num_labels_and_is_regression(cfg.task)
    sent1_key, sent2_key = GLUE_TASK_KEYS[cfg.task]
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, use_fast=True)

    preprocess_fn = build_preprocess(
        tokenizer, cfg.task, cfg.max_len, sent1_key, sent2_key, is_regression
    )
    dataset = raw_datasets.map(preprocess_fn, batched=True)
    cols = ["input_ids", "attention_mask", "labels"]
    if "token_type_ids" in dataset["validation"].column_names:
        cols.append("token_type_ids")
    dataset.set_format(type="torch", columns=cols)

    try:
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
    except Exception:
        logger.warning(f"Could not load from {model_path}, using base model")
        model = AutoModelForSequenceClassification.from_pretrained(
            cfg.model_name, num_labels=num_labels
        )
    model.to(device)

    training_args = TrainingArguments(
        output_dir="./tmp_eval",
        per_device_eval_batch_size=cfg.batch_size * 2,
        remove_unused_columns=False,
        report_to="none",
    )
    eval_trainer = Trainer(
        model=model,
        args=training_args,
        eval_dataset=dataset.get("validation", dataset.get("validation_matched")),
        data_collator=DataCollatorWithPadding(tokenizer),
    )
    eval_result = eval_trainer.evaluate()
    return {k.replace("eval_", ""): v for k, v in eval_result.items()}

Run evaluation.

def run_training(config,
stop_event=None,
progress_callback=None,
log_callback=None,
final_callback=None) ‑> Any

Expand source code

def run_training(
    config,
    stop_event=None,
    progress_callback=None,
    log_callback=None,
    final_callback=None,
) -> Any:
    """Run training."""
    if isinstance(config, dict):
        cfg = data_ExperimentConfig()
        for k, v in config.items():
            setattr(cfg, k, v)
    else:
        cfg = config

    set_seed(cfg.seed)
    device = get_device()
    logger.info(f"Using device: {device}")

    # Data
    raw_datasets = load_glue_local(cfg.task, cfg.glue_base_path)
    num_labels, is_regression = get_num_labels_and_is_regression(cfg.task)
    sent1_key, sent2_key = GLUE_TASK_KEYS[cfg.task]
    tokenizer = AutoTokenizer.from_pretrained(cfg.model_name, use_fast=True)

    preprocess_fn = build_preprocess(
        tokenizer, cfg.task, cfg.max_len, sent1_key, sent2_key, is_regression
    )
    dataset = raw_datasets.map(
        preprocess_fn, batched=True, remove_columns=raw_datasets["train"].column_names
    )
    cols = ["input_ids", "attention_mask", "labels"]
    if "token_type_ids" in dataset["train"].column_names:
        cols.append("token_type_ids")
    dataset.set_format(type="torch", columns=cols)

    # Model
    model = AutoModelForSequenceClassification.from_pretrained(
        cfg.model_name, num_labels=num_labels
    )
    fix_layernorm_naming(model)
    model.to(device)

    # Trainer
    trainer_type = getattr(cfg, "trainer_type", "hf")
    trainer_cls = get_trainer(trainer_type)

    callbacks = [
        DjangoCallback(progress_callback, log_callback, final_callback, stop_event),
        TrainerLoggingCallback(None, logger),
        MemoryCallback(),
        TimeCallback(),
    ]
    if getattr(cfg, "compressor", None):
        callbacks.append(SVDAttentionCallback())

    trainer = trainer_cls.create(
        model=model,
        dataset=dataset,
        config=cfg,
        tokenizer=tokenizer,
        logger=logger,
        callbacks=callbacks,
    )

    logger.info("Starting training...")
    trainer.train()
    eval_results = trainer.evaluate()
    logger.info("Training completed.")

    # Save
    output_dir = Path(cfg.output_dir) / (cfg.exp_name or "training")
    output_dir.mkdir(parents=True, exist_ok=True)
    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))

    try:
        create_comprehensive_plots(
            metrics=trainer.state.log_history,
            output_dir=str(output_dir / "plots"),
            prefix=cfg.exp_name or "training",
        )
    except Exception as e:
        logger.warning(f"Plot generation failed: {e}")

    if final_callback:
        final_callback(eval_results, trainer.state.log_history)

    return eval_results

Run training.

def run_training_async(experiment, stop_event: threading.Event) ‑> None

Expand source code

def run_training_async(experiment, stop_event: threading.Event) -> None:
    """
    Main entry point called by Django in a background thread.

    Args:
        experiment: Django model instance with fields:
            id, config (JSON), status, logs, current_step, total_steps, etc.
        stop_event: threading.Event to signal cancellation.
    """
    config = experiment.config
    run_id = str(experiment.id)
    base_dir = _get_media_root() / "experiments"

    def log(msg: str) -> None:
        """Log."""
        experiment.logs = (experiment.logs or "") + msg + "\n"
        experiment.save(update_fields=["logs"])
        WSMessageQueue.send(run_id, {"type": "log", "data": msg})
        broadcast_terminal(msg)

    def update_status(status: str) -> None:
        """Update status."""
        experiment.status = status
        experiment.save(update_fields=["status"])
        WSMessageQueue.send(run_id, {"type": "status", "status": status})

    def progress_callback(current: int, total: int, pct: float) -> None:
        """Progress callback."""
        experiment.current_step = current
        experiment.total_steps = total
        experiment.estimated_remaining_seconds = (total - current) * 0.1  # crude
        experiment.cpu_usage = psutil.cpu_percent()
        experiment.ram_usage_mb = psutil.virtual_memory().used / 1024**2
        experiment.save(
            update_fields=[
                "current_step",
                "total_steps",
                "estimated_remaining_seconds",
                "cpu_usage",
                "ram_usage_mb",
            ]
        )
        WSMessageQueue.send(
            run_id,
            {
                "type": "progress",
                "progress": pct,
                "current_step": current,
                "total_steps": total,
            },
        )

    def log_callback(msg: str):
        # Called by the orchestrator for log messages
        """Log callback."""
        log(msg)

    def final_callback(metrics: dict) -> None:
        """Final callback."""
        experiment.config["final_metrics"] = metrics
        experiment.save(update_fields=["config"])
        WSMessageQueue.send(run_id, {"type": "final", "data": metrics})
        log(f"Training finished. Final metrics: {json.dumps(metrics)}")
        save_artifacts(experiment, metrics)

    try:
        update_status("running")
        orchestrator = MegaOrchestrator(
            config=config,
            run_id=run_id,
            base_dir=base_dir,
            progress_callback=progress_callback,
            log_callback=log_callback,
            final_callback=final_callback,
            stop_event=stop_event,
        )
        orchestrator.run()

        if stop_event.is_set():
            update_status("stopped")
            log("Training cancelled by user.")
        else:
            update_status("completed")

    except Exception as e:
        error_msg = traceback.format_exc()
        log(f"ERROR: {str(e)}\n{error_msg}")
        update_status("failed")

Main entry point called by Django in a background thread.

Args

-----=

experiment

Django model instance with fields: id, config (JSON), status, logs, current_step, total_steps, etc.

stop_event

threading.Event to signal cancellation.

Classes

class AIAssistant (llm_function)

Expand source code

class AIAssistant:
    """
    Monitors training metrics and suggests actions (e.g., early stop).
    Uses a callable LLM function for decision making.
    """

    def __init__(self, llm_function) -> None:
        self.llm = llm_function

    def monitor(self, exp_id: str, metrics_history: list) -> None:
        """Monitor."""
        last_losses = [m.get("loss") for m in metrics_history[-5:] if "loss" in m]
        if len(last_losses) >= 5 and all(
            abs(last_losses[i] - last_losses[i - 1]) < 0.01 for i in range(1, 5)
        ):
            suggestion = "early_stop"
            reason = "Loss plateau detected"
            WSMessageQueue.send(
                exp_id,
                {"type": "ai_suggestion", "suggestion": suggestion, "reason": reason},
            )
            # Optionally trigger early stop
            # stop_event.set()

Monitors training metrics and suggests actions (e.g., early stop). Uses a callable LLM function for decision making.

Methods

def monitor(self, exp_id: str, metrics_history: list) ‑> None

Expand source code

def monitor(self, exp_id: str, metrics_history: list) -> None:
    """Monitor."""
    last_losses = [m.get("loss") for m in metrics_history[-5:] if "loss" in m]
    if len(last_losses) >= 5 and all(
        abs(last_losses[i] - last_losses[i - 1]) < 0.01 for i in range(1, 5)
    ):
        suggestion = "early_stop"
        reason = "Loss plateau detected"
        WSMessageQueue.send(
            exp_id,
            {"type": "ai_suggestion", "suggestion": suggestion, "reason": reason},
        )
        # Optionally trigger early stop
        # stop_event.set()

Monitor.

class ExperimentComparator (run_ids: list[str], baseline_id: str | None = None)

Expand source code

class ExperimentComparator:
    """
    Compare multiple TrainingRuns with rich diff output.
    Inspired by ZenML's comparison tool, MLflow compare, and RagaAI diff.
    """

    def __init__(self, run_ids: list[str], baseline_id: str | None = None) -> None:
        self.run_ids = run_ids
        self.baseline_id = baseline_id or (run_ids[0] if run_ids else None)
        self.runs: list[TrainingRun] = []

    def load_runs(self) -> None:
        """Load runs from disk or config."""
        self.runs = list(TrainingRun.objects.filter(id__in=self.run_ids).order_by("created_at"))

    def compare(self) -> ComparisonResult:
        """Compare."""
        self.load_runs()
        if len(self.runs) < 2:
            raise ValueError("Need at least 2 runs to compare")

        baseline = next((r for r in self.runs if str(r.id) == self.baseline_id), self.runs[0])
        result = ComparisonResult(
            run_ids=[str(r.id) for r in self.runs],
            metrics_table=self._build_metrics_table(),
            config_diff=self._build_config_diff(),
            hyperparameter_table=self._build_hyperparameter_table(),
            metric_deltas=self._build_metric_deltas(baseline),
            timeline_data=self._build_timeline_overlay(),
            parallel_coords=self._build_parallel_coords(),
            best_run_id=self._find_best_run(),
            summary="",
        )
        result.summary = self._generate_summary(result)
        return result

    def _build_metrics_table(self) -> list[dict[str, Any]]:
        """build metrics table."""
        rows = []
        for run in self.runs:
            metrics = run.final_metrics or {}
            config = run.config or {}
            rows.append(
                {
                    "run_id": str(run.id),
                    "name": run.name or str(run.id)[:8],
                    "status": run.status,
                    "created_at": run.created_at.isoformat(),
                    "epochs": config.get("epochs", "N/A"),
                    "batch_size": config.get("batch_size", "N/A"),
                    "learning_rate": config.get("learning_rate", "N/A"),
                    "model_name": config.get("model_name", "N/A"),
                    "eval_accuracy": metrics.get("eval_accuracy"),
                    "eval_loss": metrics.get("eval_loss"),
                    "eval_f1": metrics.get("eval_f1"),
                    "duration_seconds": run.actual_duration,
                }
            )
        return rows

    def _build_config_diff(self) -> dict[str, Any]:
        """build config diff."""
        configs = {str(r.id): r.config for r in self.runs}
        first_id = str(self.runs[0].id)
        diffs = {}
        for rid, cfg in configs.items():
            if rid != first_id:
                diff = DeepDiff(configs[first_id], cfg, ignore_order=True, verbose_level=2)
                diffs[rid] = diff
        return {
            "baseline_id": first_id,
            "diffs": diffs,
            "all_config_keys": list(set().union(*[set(c.keys()) for c in configs.values()])),
        }

    def _build_hyperparameter_table(self) -> list[dict[str, Any]]:
        """build hyperparameter table."""
        hp_keys = [
            "learning_rate",
            "batch_size",
            "epochs",
            "max_seq_length",
            "compression_rank",
            "lora_r",
            "lora_alpha",
            "weight_decay",
            "warmup_steps",
            "grad_accum",
            "max_grad_norm",
        ]
        rows = []
        for run in self.runs:
            config = run.config or {}
            row = {"run_id": str(run.id), "name": run.name or str(run.id)[:8]}
            for k in hp_keys:
                row[k] = config.get(k)
            rows.append(row)
        return rows

    def _build_metric_deltas(self, baseline: TrainingRun) -> list[dict[str, Any]]:
        """build metric deltas."""
        base_metrics = baseline.final_metrics or {}
        deltas = []
        metric_keys = [
            "eval_accuracy",
            "eval_loss",
            "eval_f1",
            "eval_precision",
            "eval_recall",
        ]
        for run in self.runs:
            if str(run.id) == str(baseline.id):
                continue
            run_metrics = run.final_metrics or {}
            delta_row = {"run_id": str(run.id), "name": run.name or str(run.id)[:8]}
            for k in metric_keys:
                base_v = base_metrics.get(k)
                run_v = run_metrics.get(k)
                if base_v is not None and run_v is not None:
                    delta_row[k] = round(run_v - base_v, 6)
                    delta_row[f"{k}_pct"] = (
                        round(((run_v - base_v) / abs(base_v)) * 100, 2) if base_v != 0 else None
                    )
                else:
                    delta_row[k] = None
            deltas.append(delta_row)
        return deltas

    def _build_timeline_overlay(self) -> dict[str, list[float]]:
        """Extract training history for overlaid loss curves."""
        timeline = {"steps": [], "series": {}}
        for run in self.runs:
            history = run.training_history or {}
            steps = history.get("steps", [])
            losses = history.get("losses", [])
            if steps and losses:
                if len(steps) > len(timeline["steps"]):
                    timeline["steps"] = steps
                timeline["series"][str(run.id)] = losses
        return timeline

    def _build_parallel_coords(self) -> dict[str, Any]:
        """Build data for a parallel coordinates plot (hyperparams + metrics)."""
        dimensions = [
            "learning_rate",
            "batch_size",
            "epochs",
            "max_seq_length",
            "compression_rank",
            "eval_accuracy",
            "eval_loss",
            "eval_f1",
        ]
        data = []
        for run in self.runs:
            config = run.config or {}
            metrics = run.final_metrics or {}
            point = {"run_id": str(run.id), "name": run.name or str(run.id)[:8]}
            for d in dimensions:
                point[d] = config.get(d) or metrics.get(d)
            data.append(point)
        return {"dimensions": dimensions, "data": data}

    def _find_best_run(self) -> str | None:
        """find best run."""
        best = None
        best_acc = -1
        for run in self.runs:
            acc = (run.final_metrics or {}).get("eval_accuracy", -1)
            if acc is not None and acc > best_acc:
                best_acc = acc
                best = str(run.id)
        return best

    def _generate_summary(self, result: ComparisonResult) -> str:
        """generate summary."""
        lines = [f"Comparison of {len(self.runs)} runs."]
        if result.best_run_id:
            best_run = next((r for r in self.runs if str(r.id) == result.best_run_id), None)
            if best_run:
                lines.append(f"Best run: {best_run.name or result.best_run_id[:8]}")
        # Count improvements
        improved = sum(
            1
            for d in result.metric_deltas
            if d.get("eval_accuracy") is not None and d["eval_accuracy"] > 0
        )
        lines.append(f"{improved} run(s) improved over baseline in accuracy.")
        return " ".join(lines)

Compare multiple TrainingRuns with rich diff output. Inspired by ZenML's comparison tool, MLflow compare, and RagaAI diff.

Methods

def compare(self) ‑> ComparisonResult

Expand source code

def compare(self) -> ComparisonResult:
    """Compare."""
    self.load_runs()
    if len(self.runs) < 2:
        raise ValueError("Need at least 2 runs to compare")

    baseline = next((r for r in self.runs if str(r.id) == self.baseline_id), self.runs[0])
    result = ComparisonResult(
        run_ids=[str(r.id) for r in self.runs],
        metrics_table=self._build_metrics_table(),
        config_diff=self._build_config_diff(),
        hyperparameter_table=self._build_hyperparameter_table(),
        metric_deltas=self._build_metric_deltas(baseline),
        timeline_data=self._build_timeline_overlay(),
        parallel_coords=self._build_parallel_coords(),
        best_run_id=self._find_best_run(),
        summary="",
    )
    result.summary = self._generate_summary(result)
    return result

Compare.

def load_runs(self) ‑> None

Expand source code

def load_runs(self) -> None:
    """Load runs from disk or config."""
    self.runs = list(TrainingRun.objects.filter(id__in=self.run_ids).order_by("created_at"))

Load runs from disk or config.

class MegaOrchestrator (config: dict[str, typing.Any],
run_id: str,
base_dir: pathlib.Path,
progress_callback:  | None = None,
log_callback:  | None = None,
final_callback:  | None = None,
stop_event: threading.Event | None = None)

Expand source code

class MegaOrchestrator:
    """
    Central orchestrator that connects all framework modules.

    Usage (from Django runner):
        orchestrator = MegaOrchestrator(
            config=request.data,
            run_id=str(uuid.uuid4()),
            base_dir=Path(settings.MEDIA_ROOT),
            progress_callback=progress_cb,
            log_callback=log_cb,
            final_callback=final_cb,
            stop_event=stop_event,
        )
        orchestrator.run()
    """

    def __init__(
        self,
        config: dict[str, Any],
        run_id: str,
        base_dir: Path,
        progress_callback: Callable | None = None,
        log_callback: Callable | None = None,
        final_callback: Callable | None = None,
        stop_event: threading.Event | None = None,
    ) -> None:
        self.raw_config = config
        self.run_id = run_id
        self.base_dir = Path(base_dir)
        self.run_dir = self.base_dir / run_id
        self.run_dir.mkdir(parents=True, exist_ok=True)

        self.progress_cb = progress_callback
        self.log_cb = log_callback
        self.final_cb = final_callback
        self.stop_event = stop_event or threading.Event()

        # Convert dict to unified Config
        self.cfg = load_config(None)  # start with defaults
        if config:
            self.cfg = Config(
                **{k: v for k, v in config.items() if k in Config.__dataclass_fields__}
            )

        self.model = None
        self.tokenizer = None
        self.trainer_state = None
        self.final_metrics: dict[str, float] = {}

        logger.info("MegaOrchestrator created: run_id=%s", run_id)

    def run(self) -> dict[str, Any]:
        """Execute the full experiment pipeline."""
        try:
            self._log("Starting experiment…")

            # 1. Load model & tokenizer
            self._log("Loading model…")
            self.model, self.tokenizer = load_model(
                self.cfg.model.name,
                num_labels=self.cfg.model.num_labels,
                attention_type=self.cfg.compression.attn_type,
                attention_rank=self.cfg.compression.rank,
            )

            # 2. Apply training method (if any)
            if self.cfg.method.method != "baseline":
                from neuralops.trainer.methods import get_method

                method_fn = get_method(self.cfg.method.method)
                if method_fn:
                    method_fn(self.model, self.cfg.compression)

            # 3. Load data
            self._log("Loading data…")
            exp_cfg = self._build_experiment_config()
            dataset = load_dataset(exp_cfg)
            tokenized = self._preprocess_dataset(dataset)

            # 4. Build trainer
            self._log("Building trainer…")
            results = ResultsManager(self.run_id, self.base_dir, config=self.raw_config)
            trainer = HuggingFaceTrainer(
                model=self.model,
                train_dataset=tokenized["train"],
                eval_dataset=tokenized.get("validation"),
                config=self.cfg.training.__dict__,
                results_manager=results,
                compute_metrics=get_task_metrics(self.cfg.training.task),
            )
            # Add basic callbacks
            from neuralops.trainer.callbacks import MemoryCallback, TimeCallback

            trainer.add_callback(MemoryCallback())
            trainer.add_callback(TimeCallback())

            # 5. Train
            self._log("Training started…")
            trainer.train()
            eval_metrics = trainer.evaluate()
            self.final_metrics = eval_metrics
            self.trainer_state = trainer.state.log_history if hasattr(trainer, "state") else []

            # 6. Save everything
            self._log("Saving results…")
            saver = ResultsSaver(
                run_id=self.run_id,
                base_dir=self.base_dir,
                config=SaveConfig(save_model=True, save_plots=True),
            )
            saver.save_all(
                self.model,
                self.cfg.__dict__,
                self.trainer_state,
                eval_metrics,
                self.tokenizer,
            )
            saver.push_all()

            self._log("Experiment completed.")
            return eval_metrics

        except Exception as e:
            logger.exception("Orchestrator failed: %s", e)
            self._log(f"Experiment failed: {traceback.format_exc()}")
            raise

    def _build_experiment_config(self) -> Any:
        """build experiment config."""
        from neuralops.data.dataset_loader import ExperimentConfig

        return ExperimentConfig(dataset_name=self.cfg.training.task)

    def _preprocess_dataset(self, dataset) -> Any:
        """preprocess dataset."""
        task = self.cfg.training.task
        s1, s2 = None, None
        from neuralops.data.glue_loader import GLUE_TASK_KEYS

        if task in GLUE_TASK_KEYS:
            s1, s2 = GLUE_TASK_KEYS[task]
        data_cfg = DataConfig(
            task_type=(
                TaskType.SINGLE_CLASSIFICATION if s2 is None else TaskType.PAIR_CLASSIFICATION
            ),
            max_length=self.cfg.data.max_seq_length,
            batch_size=self.cfg.training.batch_size,
            text_column_1=s1,
            text_column_2=s2,
            label_column="label",
        )

        def preprocess(examples) -> Any:
            """Preprocess."""
            return preprocess_function(examples, self.tokenizer, data_cfg)

        tokenized = dataset.map(
            preprocess, batched=True, remove_columns=dataset["train"].column_names
        )
        tokenized.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
        return tokenized

    def _log(self, msg: str) -> None:
        """log."""
        logger.info("[%s] %s", self.run_id[:8], msg)
        if self.log_cb:
            self.log_cb(msg)

Central orchestrator that connects all framework modules.

Usage (from Django runner): orchestrator = MegaOrchestrator( config=request.data, run_id=str(uuid.uuid4()), base_dir=Path(settings.MEDIA_ROOT), progress_callback=progress_cb, log_callback=log_cb, final_callback=final_cb, stop_event=stop_event, ) orchestrator.run()

Methods

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

Expand source code

def run(self) -> dict[str, Any]:
    """Execute the full experiment pipeline."""
    try:
        self._log("Starting experiment…")

        # 1. Load model & tokenizer
        self._log("Loading model…")
        self.model, self.tokenizer = load_model(
            self.cfg.model.name,
            num_labels=self.cfg.model.num_labels,
            attention_type=self.cfg.compression.attn_type,
            attention_rank=self.cfg.compression.rank,
        )

        # 2. Apply training method (if any)
        if self.cfg.method.method != "baseline":
            from neuralops.trainer.methods import get_method

            method_fn = get_method(self.cfg.method.method)
            if method_fn:
                method_fn(self.model, self.cfg.compression)

        # 3. Load data
        self._log("Loading data…")
        exp_cfg = self._build_experiment_config()
        dataset = load_dataset(exp_cfg)
        tokenized = self._preprocess_dataset(dataset)

        # 4. Build trainer
        self._log("Building trainer…")
        results = ResultsManager(self.run_id, self.base_dir, config=self.raw_config)
        trainer = HuggingFaceTrainer(
            model=self.model,
            train_dataset=tokenized["train"],
            eval_dataset=tokenized.get("validation"),
            config=self.cfg.training.__dict__,
            results_manager=results,
            compute_metrics=get_task_metrics(self.cfg.training.task),
        )
        # Add basic callbacks
        from neuralops.trainer.callbacks import MemoryCallback, TimeCallback

        trainer.add_callback(MemoryCallback())
        trainer.add_callback(TimeCallback())

        # 5. Train
        self._log("Training started…")
        trainer.train()
        eval_metrics = trainer.evaluate()
        self.final_metrics = eval_metrics
        self.trainer_state = trainer.state.log_history if hasattr(trainer, "state") else []

        # 6. Save everything
        self._log("Saving results…")
        saver = ResultsSaver(
            run_id=self.run_id,
            base_dir=self.base_dir,
            config=SaveConfig(save_model=True, save_plots=True),
        )
        saver.save_all(
            self.model,
            self.cfg.__dict__,
            self.trainer_state,
            eval_metrics,
            self.tokenizer,
        )
        saver.push_all()

        self._log("Experiment completed.")
        return eval_metrics

    except Exception as e:
        logger.exception("Orchestrator failed: %s", e)
        self._log(f"Experiment failed: {traceback.format_exc()}")
        raise

Execute the full experiment pipeline.