AS2 / docs /instantiations.md
wabdalm's picture
Release AS2 framework with Hydra-selected instantiations
69928d6 verified
|
Raw History Blame Contribute Delete
6.63 kB

Adding an AS² instantiation

An instantiation is a Hydra configuration plus the task-specific implementations it selects. Visual Sudoku is the bundled example. MNIST Addition and CLEVR-Hans are future targets, not runnable options in this release.

The shared commands dispatch by capability, with no benchmark branches:

uv run --locked softasp-train instantiation=your_task
uv run --locked softasp-train-scratch instantiation=your_task
uv run --locked softasp-predict instantiation=your_task
uv run --locked softasp-evaluate instantiation=your_task

These commands require your implementation and configuration to exist. Place your_task.yaml in configs/instantiation/, or supply an additional Hydra config search directory containing that group. Install your provider module in the same uv environment. No shared CLI or framework source edits are needed to select an implemented instantiation.

Configuration contract

Use # @package _global_ so the instantiation can compose dataset, preprocessing, model, trainer, evaluator, paths and any workflow-specific sections. The bundled sudoku.yaml is a complete working example.

An extension configuration follows this structure (provider names are illustrative):

# @package _global_
defaults:
  - /trainer: default
  - _self_

dataset:
  name: your_task
  # Dataset parameters, splits, paths and preprocessing choices go here.
model:
  name: your_model
  pretrained_from: null
  freeze_backbone_epochs: 0
  params: {}
evaluator:
  checkpoint: null
  metrics: [accuracy]

instantiation:
  name: your_task
  version: 1
  constraints:
    _target_: your_package.constraints.ConstraintLoss
  providers:
    dataset: {_target_: your_package.providers.DatasetProvider}
    model: {_target_: your_package.providers.ModelProvider}
    metrics: {_target_: your_package.providers.MetricsProvider}
    prediction: {_target_: your_package.providers.PredictionProvider}
    evaluation: {_target_: your_package.providers.EvaluationProvider}
    workflow: {_target_: your_package.providers.WorkflowProvider}

Add hyperparameters, data identities, paths, network composition, objective weights, evidence policies, input/output settings and workflow stages needed by your task. Use the existing trainer schema and configuration groups; unknown trainer fields are not arbitrary model parameters. Model-specific settings belong under model.params or an instantiation-owned section. The top-level paths group is composed by the common training configurations; prediction configurations also compose it explicitly. predict supplies checkpoint, input, output and device for the bundled adapter; another adapter may define different inputs.

Provider methods

Hydra constructs each provider from its _target_ and constructor parameters. Keep provider construction free of stochastic model initialization; build models inside prepare or build, after the training entrypoint applies its seed policy.

Provider Required methods Contract
Dataset build(cfg) Return (training_dataset, validation_dataset); handle task-specific preparation/selection and identity validation.
Model prepare(cfg), build(cfg) prepare returns a model or None before dataset construction; build constructs after datasets if needed. Bind configured constraints/objectives here.
Model inference load(cfg) Required by softasp.inference.load_model(cfg); validate and strictly load the selected checkpoint.
Metrics build(cfg) Return an evaluator with evaluate(model, loader, device) -> dict[str, float].
Prediction run(cfg) Return a JSON-serializable result. The common CLI writes cfg.output without overwriting an existing file. Never use solution labels.
Evaluation run(cfg) Score configured splits and write task-specific reports/provenance; return a result dictionary.
Workflow run(cfg) Execute the task's full workflow and return its result/artifact path. It may call shared run_training(cfg).

Training datasets yield dictionaries of tensors accepted by the shared loader. Models implement BaseSoftASPModel.forward(batch) with a differentiable loss and appropriate logits or predictions. The standard evaluator uses labels and logits/predictions; supply a custom evaluator for other target/output contracts. For freezing or gradient ramps, expose the relevant module through get_backbone(). A model without a backbone returns None.

Reusable components and task-specific constraints

softasp_solver.components.reasoning.AS2Network composes an encoder, concept head, probability projection, attention module, task head and optional structure module. Its input encoder produces (batch, positions, features). A task head may return per-position outputs or pool into a classification output. Images, 81 positions, nine symbols, Sudoku masks and per-cell predictions are not framework requirements.

The components package also supplies explicit evidence policies and structural embeddings; constraints supplies shared operators and fixed-point residuals. An instantiation may reuse these components or provide another differentiable implementation appropriate to its declarative constraints.

Each instantiation owns its constraint specification and implementation. Sudoku keeps both in softasp_solver.sudoku.constraints, with its training objectives in softasp_solver.sudoku.objective. Changing specification text alone does not change the numerical implementation: there is no automatic ASP compiler.

Checkpoints and verification

Training saves the resolved configuration, including instantiation identity and providers. The common resume path validates instantiation/model identity and constraint/network configuration before restoring state. Prediction/evaluation providers must call the shared checkpoint validation helpers and load state strictly; the bundled Sudoku providers do so. Only the explicitly recognized legacy Sudoku inference checkpoint is accepted without instantiation metadata.

Test your adapter through all four commands, checkpoint round-trip, incompatible identity rejection, and label-free prediction. The release's development checks exercise a non-image, pooled-head AS² fixture with its own constraint residual through the same entrypoints. That fixture is an interface test, not an additional scientific benchmark. Dataset downloading and from-scratch stage order belong to the instantiation; the common framework does not impose Sudoku's two-stage recipe.