mojolearn.LanguageModelInference#

class mojolearn.LanguageModelInference(parameters, *, shape=None, threaded=True, threads=None)[source]#

Bases: object

Run a trained decoder language model on the CPU, with no GPU.

Forward pass only: logits, loss and greedy next-token prediction for a model trained with LanguageModelTrainer, giving the same bits as the GPU forward pass. The parameters are copied at construction and never change.

By default (threaded=True) calls use SIMD and up to threads threads (None: one per physical core); threaded=False runs single-threaded reference code. Both give the same bits. Each call can override both settings.

Parameters:

parameters (array-like or dict) – Flat float32 parameters of shape.n_total values (finite), or a dict of named tensors as listed by LanguageModelTrainer.parameter_registry(); packed low-bit tensors from mojolearn.lowbit are expanded exactly.
shape (LanguageModelConfig or None, default None) – Model dimensions; None uses the defaults.
threaded (bool, default True) – Use the multithreaded path.
threads (int or None, default None) – Thread count; None means one per physical core.

Variables:

weight_format (str) – "float32", or the packed format of the parameters given.

Examples

>>> import numpy as np, mojolearn
>>> lm = mojolearn.LanguageModelInference.from_checkpoint("model.json")
>>> ids = np.frombuffer(b"Hello, wor", dtype=np.uint8).astype(np.int32).reshape(1, -1)
>>> next_ids = lm.next_bytes(ids)

__init__(parameters, *, shape=None, threaded=True, threads=None)[source]#

Methods

`from_checkpoint`(path, *[, threaded, threads])	Load the parameters of a JSON checkpoint written by the trainer.
`logits`(ids, *[, lengths, threaded, threads])	Float32 logits for every position.
`loss`(ids, *[, threaded, threads])	Mean next-token loss as a float; see loss_bits().
`loss_bits`(ids, *[, threaded, threads])	IEEE-754 bit pattern of the mean next-token loss.
`next_bytes`(ids, *[, lengths, threaded, threads])	Greedy next token after each row.
`parameters_sha256`()	SHA-256 of the parameters' little-endian float32 bytes.

Attributes

`profile`	The name of the arithmetic profile the CPU binding computes.
`shape`	The model dimensions, a `LanguageModelConfig`.

classmethod from_checkpoint(path, *, threaded=True, threads=None)[source]#

Load the parameters of a JSON checkpoint written by the trainer.

Reads a file written by LanguageModelTrainer.export_checkpoint(), with the trainer’s own integrity checks.

Parameters:

path (str or path-like) – The checkpoint file.
threaded (bool, default True) – As for the constructor.
threads (int or None, default None) – As for the constructor.

Returns:

LanguageModelInference

property shape#: The model dimensions, a LanguageModelConfig.

property profile#: The name of the arithmetic profile the CPU binding computes.

parameters_sha256()[source]#: SHA-256 of the parameters’ little-endian float32 bytes.

logits(ids, *, lengths=None, threaded=None, threads=None)[source]#

Float32 logits for every position.

Parameters:

ids (array-like of shape (batch, length), int32) – Token ids, starting at position 0; length at most shape.length.
lengths (sequence of int or None, default None) – Makes the batch ragged: row i is real at positions [0, lengths[i]) and padding after, whatever the padding ids are. Real positions get exactly the logits of the row alone at its own length; padding positions are +0.0.
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

Array of shape (batch, length, vocab_size), float32

loss_bits(ids, *, threaded=None, threads=None)[source]#

IEEE-754 bit pattern of the mean next-token loss.

Parameters:

ids (array-like of shape (shape.batch, shape.length + 1), int32) – The training batch layout. The last column is only a target; -100 there is ignored.
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

int – The float32 loss’s bits, for exact comparison.

loss(ids, *, threaded=None, threads=None)[source]#: Mean next-token loss as a float; see loss_bits().

next_bytes(ids, *, lengths=None, threaded=None, threads=None)[source]#

Greedy next token after each row.

Picks the highest logit at each row’s last position (with lengths, its last real position), ties going to the lowest id.

Parameters:

ids (array-like of shape (batch, length), int32) – Token ids.
lengths (sequence of int or None, default None) – Ragged row lengths, as for logits().
threaded (bool or None, default None) – Override the instance setting.
threads (int or None, default None) – Override the instance setting.

Returns:

list of int – One token id per row.