chimera-ml 0.2.0

Chimera ML (Cross-modal Hierarchical Merging of Embeddings and Representations) - a lightweight training framework for configurable uni-modal and multi-modal experiments.

Chimera ML

Chimera ML (Cross-modal Hierarchical Merging of Embeddings and Representations) is a lightweight framework for training and evaluating configurable uni-modal and multi-modal models.

The core idea is simple:

• the library provides runtime infrastructure (trainer, logging, callbacks, registry builders),
• task-specific components (datamodules, models, losses, metrics, callbacks) live in external plugin packages,
• experiments are driven by YAML configs and CLI commands.

Requirements

• Python >=3.12,<3.13
• PyTorch >=2.2,<3.0

Installation

Install from PyPI:

pip install chimera-ml

or:

poetry add chimera-ml

By default, the package includes MLflow, matplotlib, and requests dependencies.

Install from source (development):

poetry install --with dev

CLI entry points:

chimera-ml --help
chimera_ml --help

Both names are available; docs use chimera-ml.

Quick Start

1. Install chimera-ml.
2. Install your task plugin package (example below).
3. Run train/eval with YAML config.

Example plugin from this repo:

pip install -e examples/va_estimation
chimera-ml validate-config --config-path examples/va_estimation/configs/multimodal_train.yaml
chimera-ml train --config-path examples/va_estimation/configs/multimodal_train.yaml
chimera-ml eval --config-path examples/va_estimation/configs/multimodal_test.yaml --checkpoint-path path/to/last.pt

CLI

Main commands:

chimera-ml validate-config --config-path <config.yaml>
chimera-ml doctor
chimera-ml train --config-path <config.yaml>
chimera-ml sweep --base-config <config.yaml> --sweep-config <sweep.yaml> [--max-trials N]
chimera-ml eval --config-path <config.yaml> --checkpoint-path <ckpt.pt> [--with-features]
chimera-ml registry list [--type models|losses|metrics|optimizers|schedulers|callbacks|collates|loggers|datamodules]
chimera-ml plugins list [--group chimera_ml.plugins]

validate-config:

• checks config structure and required sections without starting training.

doctor:

• prints quick environment diagnostics (Python, torch, CUDA, MLflow, registry/plugin counts).

train:

• requires experiment_info.params.experiment_name,
• generates run_name via generate_run_name(...),
• patches checkpoint_callback and snapshot_callback params with experiment/run data,
• builds all components from registries and runs Trainer.fit(...).

sweep:

• materializes one YAML config per trial under sweep_runs/ by default,
• supports Cartesian grids via parameters (grid search) and explicit trial lists via trials,
• applies overrides using dotted paths such as optimizer.params.lr or callbacks.checkpoint_callback.params.monitor,
• runs the normal train flow once per trial and appends sweep_001, sweep_002, ... to run names,
• supports --max-trials for CI smoke tuning and --dry-run to inspect generated trials.

eval:

• builds datamodule/model/loss/metrics/callbacks from the same config style,
• optionally loads checkpoint (model_state_dict or raw state dict),
• evaluates over merged train/val/test loader splits when available.

registry list:

• prints currently registered keys (including keys loaded from plugins).

plugins list:

• prints discovered Python entry points for plugin group chimera_ml.plugins.

YAML Config Model

Top-level sections used by runtime:

• seed
• experiment_info (required for train)
• data
• model
• train
• loss
• optimizer
• scheduler (optional)
• metrics (list)
• logging (list)
• callbacks (list)

Minimal skeleton:

seed: 0

experiment_info:
  params:
    experiment_name: "my_experiment"
    include_time: true
    datetime_format: "%Y-%m-%d_%H-%M"
    timezone: "UTC"

data:
  name: "my_datamodule"
  params: {}

model:
  name: "my_model"
  params: {}

train:
  params:
    epochs: 10
    device: "cuda"
    mixed_precision: true
    use_scheduler: true
    scheduler_step_per_epoch: true
    scheduler_monitor: "val/loss"

loss:
  name: "mse_loss"
  params: {}

optimizer:
  name: "adamw_optimizer"
  params:
    lr: 1e-3

scheduler:
  name: "steplr_scheduler"
  params:
    step_size: 10
    gamma: 0.5

metrics:
  - name: "mae_metric"
    params: {}

logging:
  - name: "console_file_logger"
    params:
      log_path: "logs"
  - name: "mlflow_logger"
    params:
      tracking_uri: "sqlite:///logs/mlflow.db"

callbacks:
  - name: "checkpoint_callback"
    params:
      monitor: "val/loss"
      mode: "min"

Sweep grid example:

parameters:
  optimizer.params.lr: [0.001, 0.0003, 0.0001]
  train.params.epochs: [3, 5]

Explicit trial example:

trials:
  - optimizer.params.lr: 0.001
    train.params.epochs: 3
  - optimizer.params.lr: 0.0001
    callbacks.checkpoint_callback.params.monitor: "val/ccc"

TrainConfig Parameters

train.params is mapped to TrainConfig:

• epochs (default 10)
• grad_clip_norm (default null)
• mixed_precision (default false)
• log_every_steps (default 50)
• device (default "cuda")
• train_loader_mode: single | round_robin | weighted (default single)
• train_stop_on: min | max (default min)
• train_loader_weights (optional mapping for weighted mode)
• use_scheduler (default false)
• scheduler_step_per_epoch (default true)
• scheduler_monitor (optional metric key)
• collect_cache (default true)

Data Contract

The trainer expects Batch objects:

Batch(
    inputs={"modality": tensor, ...},
    targets=tensor_or_none,
    masks={"sequence_mask": ..., "audio_mask": ..., ...} or None,
    meta={"sample_meta": [...], ...} or None,
)

Built-in MaskingCollate (masking_collate) supports variable-length multimodal inputs and creates:

• padded modality tensors,
• sequence_mask,
• per-modality masks like {modality}_mask,
• optional legacy meta["masks"].

Plugin System

On import chimera_ml, the library calls register_all():

• loads built-in modules that register built-in components,
• loads external entry points from group chimera_ml.plugins,
• executes callable entry points once per process.

Plugin declaration (recommended, PEP 621 style):

[project.entry-points."chimera_ml.plugins"]
my_project = "my_project.chimera_plugin:register"

Typical register() function:

def register():
    import my_project.data
    import my_project.models
    import my_project.losses
    import my_project.metrics
    import my_project.callbacks

Import side effects execute registry decorators.

Built-In Registry Keys

Datamodules are intentionally project-specific. Built-in DATAMODULES is empty by default.

MODELS:

• feature_fusion_model
• prediction_fusion_model
• gated_fusion_model
• gated_prediction_fusion_model

LOSSES:

• mse_loss
• mae_loss
• cross_entropy_loss
• focal_loss
• bce_with_logits_loss
• ccc_loss

METRICS:

• mae_metric
• mse_metric
• rmse_metric
• r2_metric
• prf_macro_metric
• prf_micro_metric
• prf_weighted_metric
• confusion_matrix_metric

OPTIMIZERS:

• adamw_optimizer
• adam_optimizer
• sgd_optimizer

SCHEDULERS:

• steplr_scheduler
• cosineannealinglr_scheduler
• reduceonplateau_scheduler

CALLBACKS:

• checkpoint_callback
• collect_predictions_callback
• early_stopping_callback
• snapshot_callback
• telegram_notifier_callback

LOGGERS:

• console_file_logger
• mlflow_logger

COLLATES:

• masking_collate

Training and Evaluation Behavior

• Mixed precision uses torch.amp.autocast and GradScaler on CUDA.
• Training supports one or many train loaders.
• Multi-loader scheduling modes:
  • single: first loader only.
  • round_robin: cycle loaders.
  • weighted: stochastic sampling with loader weights.
• Validation/test loaders are normalized to stable split keys (normalize_loaders).
• Metrics are stateful (reset -> update -> compute).
• Optional prediction cache (CachedSplitOutputs) stores CPU preds/targets/features for callbacks.

Logging and Artifacts

console_file_logger:

• logs to console + file under <log_path>/<experiment_name>/<run_name>/.

mlflow_logger:

• logs params/metrics,
• supports file/text/bytes artifact logging,
• logs config artifact when config_path is provided.

plot_confusion_matrix_callback:

• builds confusion matrix figures from cached predictions,
• logs PNG artifacts to MLflow per split (figures/<split>/...).

telegram_notifier_callback:

• sends final run status to Telegram via Bot API.

Callback Lifecycle

Callbacks follow:

• on_fit_start
• on_epoch_start
• on_batch_end
• on_epoch_end
• on_fit_end

Highlights:

• checkpoint_callback: monitor-based top-k and last.pt.
• early_stopping_callback: monitor, mode, patience, min_delta.
• snapshot_callback: code/config snapshots (code.zip, config copy).
• collect_predictions_callback: CSV prediction artifacts to MLflow.
• plot_confusion_matrix_callback: confusion matrix PNG artifacts to MLflow.
• telegram_notifier_callback: final Telegram notification via env vars.

Repository Example

examples/va_estimation is a full plugin package using entry points and task-specific components. Use it as a template for new projects.

Development

Quality checks:

poetry run ruff check src tests
poetry run pytest

See CONTRIBUTING.md for contribution details.

Publishing

PyPI publishing is automated via GitHub Actions workflow:

• .github/workflows/publish.yml (triggered by GitHub Release published).

Release checklist:

• see RELEASING.md.