mojolearn.models

Load and run Hugging Face-format causal language models.

Reads .safetensors checkpoints (single file or sharded) and their config.json, maps the weights onto mojolearn’s transformer blocks, and runs them on the GPU.

• CausalLM: embedding, decoder blocks, final norm and head built from a checkpoint, with forward, allocate_state, step and greedy generate.

• ParallelCausalLM: the same model with each layer assigned to a device (layer_devices).

• Tokenizer: Tokenizer.from_pretrained over tokenizer.json for byte-level BPE tokenizers, with the pre-tokenization pattern (PATTERNS: GPT-2, Llama 3, Qwen 2) as a parameter.

• HFConfig, plan_for(), FAMILIES: which model_type values are supported (Llama 2/3 and SmolLM, Mistral, Qwen 2, Qwen 3, …), how config fields map onto block options, and which fields are refused by name with UnsupportedModel.

• Checkpoint, SafetensorsFile: the .safetensors reader; tensors come back as mojolearn.Array.

Examples

>>> from mojolearn import Array
>>> from mojolearn.models import CausalLM, Tokenizer
>>> lm = CausalLM.load("path/to/model", weight_format="bfloat16")
>>> tok = Tokenizer.from_pretrained("path/to/model")
>>> ids = tok.encode("The capital of France is")
>>> out = lm.generate(Array.from_list([ids], "<i4"), 8)

Functions

plan_for(config)

The ModelPlan of an HFConfig (or a dict), or UnsupportedModel BY NAME.

Classes

ParallelCausalLM(plan, weights, *, ...)	Experimental layer ownership.
CausalLM(plan, weights, *[, weight_format, ...])	See the module header.
CausalLMState(batch_size, max_tokens, layers, *)	The decode state: one block state per layer in stack order (each caller-owned and documented by its block class), batch_size, max_tokens (the capacity the transformer caches were allocated for) and positions, how many tokens each row has consumed.
Checkpoint(files, weight_map, root)	A single model.safetensors, an explicit .safetensors path, or a directory holding model.safetensors.index.json and its shards.
HFConfig(data[, path])	A config.json as a read-only mapping with model_type, get, [] and path.
ModelPlan(**kw)	What CausalLM assembles from.
SafetensorsFile(path)	One .safetensors file.
Tokenizer(tokens, merges, *, pattern[, ...])	See the module header.

Exceptions

UnsupportedModel

A config field the blocks cannot honor, refused by name.

class mojolearn.models.ParallelCausalLM(plan, weights, *, layer_devices, **kwargs)

Bases: CausalLM

Experimental layer ownership. Use load(path, layer_devices=(0, 1)).

One device index per checkpoint layer; repeated indices are allowed. Metal requires every index to be zero (one worker on one GPU). Close explicitly or use a context manager. Calls and states must not be used concurrently. Output and state mathematics are the ordinary CausalLM path.

classmethod load(path, *, layer_devices, weight_format='float32', max_positions=None, numeric_profile=None)

A checkpoint directory (config.json plus model.safetensors or model.safetensors.index.json and its shards).

parameters()

Every tensor by CHECKPOINT name, as held (packed where packed).

class mojolearn.models.CausalLM(plan, weights, *, weight_format='float32', max_positions=None, device='auto', numeric_profile=None)

Bases: object

See the module header. Build with load (a checkpoint directory) or directly from a ModelPlan and a weight dict shaped like _read_weights returns.

classmethod load(path, *, weight_format='float32', max_positions=None, device='auto', numeric_profile=None)

A checkpoint directory (config.json plus model.safetensors or model.safetensors.index.json and its shards).

parameters()

Every tensor by CHECKPOINT name, as held (packed where packed).

forward(ids, state=None)

(B, L) int32 ids in, (B, L, vocab) float32 logits out. With state=None every block runs from a zero state and nothing is kept; with a CausalLMState the blocks read and update it in place.

allocate_state(batch_size, max_tokens)

The zero decode state for batch_size rows of up to max_tokens positions each (Mamba*Block.allocate_state(B), TransformerBlock.allocate_state(B, max_tokens)).

reset_state(state)

Replace caches with fresh zero state, preserving capacity and batch size.

States are owned by their allocating model; cross-model reuse is refused.

step(ids, state)

One decode token per row: (B,) or (B, 1) ids in, (B, vocab) logits out, state updated in place. It is forward at L = 1 with each block’s step (their one spelling for decode).

generate(ids, max_new_tokens, *, greedy=True)

(B, L) ids in, (B, L + max_new_tokens) int32 ids out: the prompt followed by greedy continuations (argmax of the last position, ties to the lowest index), one carried state, the prompt as one prefill and every later token as one step. No stop token: exactly max_new_tokens are produced per row.

class mojolearn.models.CausalLMState(batch_size, max_tokens, layers, *, owner=None)

Bases: object

The decode state: one block state per layer in stack order (each caller-owned and documented by its block class), batch_size, max_tokens (the capacity the transformer caches were allocated for) and positions, how many tokens each row has consumed.

class mojolearn.models.Checkpoint(files, weight_map, root)

Bases: object

A single model.safetensors, an explicit .safetensors path, or a directory holding model.safetensors.index.json and its shards. Shards are opened on first use and closed by close().

class mojolearn.models.HFConfig(data, path=None)

Bases: object

A config.json as a read-only mapping with model_type, get, [] and path.

classmethod from_json(path)

path is config.json or the directory holding it.

class mojolearn.models.ModelPlan(**kw)

Bases: object

What CausalLM assembles from. Fields:

model_type, kind (“transformer” | “mamba1” | “mamba2”), n_layers, d_model, vocab_size, tie_embeddings, norm_eps (the FINAL norm’s epsilon, always passed to the norm primitive), max_position_embeddings, block_options (every interface option with its value, transformer only), block_kwargs (Mamba-2’s dt_limit, else empty), embed_name, norm_name, head_name (None when tied), layer_weights(i) -> list of (block_key, checkpoint_name, row_slice_or_None), family (the matrix row).

class mojolearn.models.SafetensorsFile(path)

Bases: object

One .safetensors file. names(), info(name), read(name); a context manager that closes the map. The map is opened on the first read and released by close().

read(name, *, bf16='widen')

The tensor name as an owned Array (see the module header for what each dtype becomes). bf16 is “widen” (float32) or “bits” (the uint16 bits; a 2-D tensor as a lowbit.BF16Weight).

class mojolearn.models.Tokenizer(tokens, merges, *, pattern, added=None, normalizer=None, bos_ids=(), eos_ids=(), ignore_merges=False, source='<tokens>')

Bases: object

See the module header. tokens are the rank-ordered token bytes, merges the (left, right) byte pairs, pattern a name in PATTERNS, added {content: id}.

classmethod from_pretrained(path)

path is a directory holding tokenizer.json (and optionally tokenizer_config.json), or the tokenizer.json itself.

property n_vocab

The ranks, extended by the added tokens above them.

pretokenize(data)

The pre-tokens of data (bytes) under this tokenizer’s pattern, as a list of bytes; the cut before any merge.

encode_bytes(data, *, allow_special=False)

The ids of a byte string. Added tokens are recognized only with allow_special=True; the segments between them are encoded alone (a special token ends the text before it, as the binding’s <|endoftext|> rule has it).

encode(text, *, add_special_tokens=True, allow_special=False)

The ids of text (str, NFC-normalized when the file says so, then UTF-8; or bytes-like), with the template’s BOS/EOS ids around them when add_special_tokens is True.

decode_bytes(ids)

The bytes of an id sequence; an added token renders its content. An id outside the vocabulary is refused by value and position.

exception mojolearn.models.UnsupportedModel

Bases: ValueError

A config field the blocks cannot honor, refused by name.

mojolearn.models.plan_for(config)

The ModelPlan of an HFConfig (or a dict), or UnsupportedModel BY NAME.