mojolearn.SambaInference

class mojolearn.SambaInference(config, weights)

Bases: object

Forward-only SambaStack language model on the CPU.

Maps (B, L) integer token ids to (B, L, vocab) float32 logits with the stack’s forward (no dropout): embedding, each block in stack order, final RMSNorm and the tied or untied output head. Build it from a SambaStack.save_checkpoint file with from_checkpoint, or from a SambaConfig and a dict of weights keyed by the config’s registry names. Weights may be packed by mojolearn.lowbit; they are unpacked exactly to float32 and copied at construction. allocate_state, forward(ids, state) and step carry a decode state. There is no loss, backward or training: loss and train_step raise NotImplementedError.

Parameters:

• config (SambaConfig) – Model configuration; anything else raises TypeError.

• weights (dict) – float32 arrays keyed by registry name. Missing or unknown names, wrong shapes and non-finite values raise ValueError.

Variables:

• config (SambaConfig) – The configuration.

• weight_format (str) – Format the weights were given in: “float32”, or the packed mojolearn.lowbit format.

Examples

>>> import mojolearn
>>> lm = mojolearn.SambaInference.from_checkpoint("samba.ckpt")
>>> logits = lm.forward(ids)                  # ids: (B, L) integers
>>> st = lm.allocate_state(batch_size=1, max_tokens=256)
>>> logits = lm.forward(prompt, st)
>>> next_logits = lm.step(next_ids, st)       # (B, vocab)

__init__(config, weights)

Methods

allocate_state(batch_size, max_tokens)	A zero decode state for batch_size sequences of up to max_tokens positions.
forward(inputs[, state, lengths])	Logits for token ids; also available as logits and by calling the model.
from_checkpoint(path)	Build from a SambaStack.save_checkpoint file (the streamed archive or the older JSON format).
logits(inputs[, state, lengths])	Logits for token ids; also available as logits and by calling the model.
loss(inputs, targets)	Not available: inference only.
parameters()	An independent copy of every weight, keyed by registry name.
step(inputs, state)	Decode one token per row and update state in place.
train_step(inputs, targets)	Not available: inference only.

classmethod from_checkpoint(path)

Build from a SambaStack.save_checkpoint file (the streamed archive or the older JSON format).

The checkpoint’s integrity checks run, and its schema, profile and weight registry must match its config, otherwise ValueError. Optimizer, schedule and random state are validated and dropped.

Parameters:

path (str or path-like) – Checkpoint file.

Returns:

SambaInference

parameters()

An independent copy of every weight, keyed by registry name.

forward(inputs, state=None, *, lengths=None)

Logits for token ids; also available as logits and by calling the model.

With state=None every block starts from a zero state. Pass a state from allocate_state to carry it: each block’s state is read and updated in place, so a later forward or step on that state continues the sequence.

Parameters:

• inputs (array of int, shape (B, L)) – Token ids in [0, vocab).

• state (SambaState or None, default None) – Decode state to read and update.

• lengths (sequence of int or None, default None) – Ragged batch: row i is real at positions [0, lengths[i]). Stateless forward only; with a state it raises ValueError.

Returns:

mojolearn.Array of float32, shape (B, L, vocab)

logits(inputs, state=None, *, lengths=None)

Logits for token ids; also available as logits and by calling the model.

With state=None every block starts from a zero state. Pass a state from allocate_state to carry it: each block’s state is read and updated in place, so a later forward or step on that state continues the sequence.

Parameters:

• inputs (array of int, shape (B, L)) – Token ids in [0, vocab).

• state (SambaState or None, default None) – Decode state to read and update.

• lengths (sequence of int or None, default None) – Ragged batch: row i is real at positions [0, lengths[i]). Stateless forward only; with a state it raises ValueError.

Returns:

mojolearn.Array of float32, shape (B, L, vocab)

allocate_state(batch_size, max_tokens)

A zero decode state for batch_size sequences of up to max_tokens positions.

Holds one block state per layer, in stack order, as SambaStack.allocate_state does. Both arguments must be positive, otherwise ValueError.

Parameters:

• batch_size (int) – Number of sequences.

• max_tokens (int) – Maximum total positions per sequence.

Returns:

SambaState

step(inputs, state)

Decode one token per row and update state in place.

A step’s logits equal the logits a one-shot forward gives at that position.

Parameters:

• inputs (array of int, shape (B,) or (B, 1)) – Token ids.

• state (SambaState) – Required decode state from allocate_state.

Returns:

mojolearn.Array of float32, shape (B, vocab)

loss(inputs, targets)

Not available: inference only. Raises NotImplementedError; use SambaStack.

train_step(inputs, targets)

Not available: inference only. Raises NotImplementedError; use SambaStack.