mixle.task.artifact module

Durable, portable artifacts for small task models – the contract mixle’s JSON serialization can’t carry.

mixle.utils.serialization round-trips pure probabilistic models as registry-keyed JSON, but a task model is usually torch-backed (a distilled tiny transformer, an MLP head), and its parameters are weights, not a JSON-serializable state. Worse, the causal LM ties head.weight = tok.weight; a naive tensor dump rejects the shared storage. This module is the missing piece: a self-describing directory that pairs

• manifest.json – how to rebuild the module (a registered builder name + its config) plus task I/O and free-form metadata, and

• weights.safetensors – the parameters, written through safetensors.torch.save_model so tied weights survive,

so a fitted model survives the process that made it. save_module/load_module are the torch path; save_json/load_json are the fallback for a pure mixle distribution. A builder is any (**config) -> nn.Module callable registered by name (register_builder); the two native architectures (mixle.causal_lm, mixle.mlp) self-register on first use, and a caller can register its own.

The acceptance bar is a fresh-process round trip: save here, load in a new interpreter from the manifest alone, get bit-identical outputs. mixle.task.model.TaskModel builds the callable task surface on top of this.

register_builder(name, builder)

Register builder under name so an artifact carrying builder=name can reconstruct its module.

builder(**config) must return a fresh (untrained) nn.Module whose parameter shapes match the saved weights. Re-registering the same name with the same callable is a no-op; a conflicting one raises.

Parameters:

• name (str)

• builder (Callable[[...], Any])

Return type:

None

get_builder(name)

Look up a registered builder, triggering native-builder self-registration on first call.

Parameters:

name (str)

Return type:

Callable[[…], Any]

class TaskManifest(payload, builder=None, config=<factory>, task='', io=<factory>, meta=<factory>, schema_version='1', created_at='')

Bases: object

The self-describing header of a task artifact: enough to rebuild and call the model, plus provenance.

Parameters:

• payload (str)

• builder (str | None)

• config (dict[str, Any])

• task (str)

• io (dict[str, Any])

• meta (dict[str, Any])

• schema_version (str)

• created_at (str)

payload: str

builder: str | None = None

config: dict[str, Any]

task: str = ''

io: dict[str, Any]

meta: dict[str, Any]

schema_version: str = '1'

created_at: str = ''

to_dict()

Return type:

dict[str, Any]

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

TaskManifest

read_manifest(path)

Read just the manifest of an artifact directory (cheap: no weights loaded).

Parameters:

path (str)

Return type:

TaskManifest

save_module(path, module, builder, config, *, task='', io=None, meta=None)

Persist a torch module as an artifact directory and return path.

builder/config must reconstruct an architecturally identical module (get_builder(builder)(**config)); weights go through safetensors.torch.save_model so tied parameters (e.g. the LM’s tied head) round-trip.

Parameters:

• path (str)

• module (Any)

• builder (str)

• config (dict[str, Any])

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_module(path, *, device='cpu')

Rebuild a torch module from its manifest alone and load weights; return (module, manifest).

Parameters:

• path (str)

• device (str)

Return type:

tuple[Any, TaskManifest]

save_json(path, model, *, task='', io=None, meta=None)

Persist a pure (torch-free) mixle distribution via the safe serialization registry; return path.

Parameters:

• path (str)

• model (Any)

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_json(path)

Rebuild a pure mixle distribution from a json-payload artifact; return (model, manifest).

Parameters:

path (str)

Return type:

tuple[Any, TaskManifest]

register_arrays_builder(name, builder)

Register builder(arrays: dict[str, ndarray], **config) -> model for arrays-payload artifacts.

The arrays payload is for torch-free numeric students (e.g. an int8-quantized MLP): weights live in one .npz, and the builder reconstructs the runnable model from them in a fresh process.

Parameters:

• name (str)

• builder (Callable[[...], Any])

Return type:

None

get_arrays_builder(name)

Look up a registered arrays builder, triggering native self-registration on first call.

Parameters:

name (str | None)

Return type:

Callable[[…], Any]

save_arrays(path, arrays, builder, config=None, *, task='', io=None, meta=None)

Persist a dict of numpy arrays as an artifact directory (arrays.npz); return path.

Parameters:

• path (str)

• arrays (dict[str, Any])

• builder (str)

• config (dict[str, Any] | None)

• task (str)

• io (dict[str, Any] | None)

• meta (dict[str, Any] | None)

Return type:

str

load_arrays(path)

Rebuild a torch-free model from an arrays-payload artifact; return (model, manifest).

Parameters:

path (str)

Return type:

tuple[Any, TaskManifest]