mixle.models.streaming_transformer_leaf module

A streaming, non-buffering transformer-LM leaf for avoiding host-RAM materialization.

Where NeuralCategorical buffers the whole shard in the accumulator and creates a fresh optimizer every M-step, this leaf keeps a long-lived module and optimizer in the M-step. The accumulator’s seq_update is one train step on a streamed micro-batch. value() returns (loss_sum, tokens) – two telemetry floats rather than the corpus – and estimate() is a no-op that wraps the live module.

This deliberately voids the sufficient-statistic algebra (value/combine are telemetry, not a foldable statistic): a sanctioned non-leaf carve-out (like NeuralGaussian.sample() raising), not an ABC change. It is the single-process prerequisite for the distributed neural handle: each rank keeps its streamed shard resident and the only cross-rank collective becomes the in-backward gradient reduce-scatter, never a gather-suff-stats-to-root.

class StreamingTransformer(module, device='cpu')

Bases: SequenceEncodableProbabilityDistribution

Wraps a live, persistently-trained module. seq_log_density = next-token log p (eval/telemetry).

Parameters:

• module (Any)

• device (str)

classmethod from_config(vocab, *, d_model=128, n_layer=4, n_head=4, block=64, embedding=None, device='cpu')

Build the leaf from hyperparameters (no hand-built torch module) – the declarative estimator surface.

embedding optionally ties a shared CategoricalEmbedding across leaves.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• device (str)

Return type:

StreamingTransformer

log_density(xy)

Return the next-token log probability for one (context, token) pair.

Parameters:

xy (Any)

Return type:

float

predict(x)

Return argmax next-token predictions for one or more contexts.

Parameters:

x (Any)

Return type:

ndarray

sampler(seed=None)

Return the sampler for the conditional next-token model.

Parameters:

seed (int | None)

Return type:

StreamingTransformerSampler

seq_log_density(enc)

Return per-row next-token log probabilities for encoded context/token pairs.

Parameters:

enc (Any)

Return type:

ndarray

estimator(pseudo_count=None)

Return the streaming estimator that trains the live module in accumulator updates.

Parameters:

pseudo_count (float | None)

Return type:

StreamingTransformerEstimator

dist_to_encoder()

Return the encoder for context/token training pairs.

Return type:

StreamingTokenEncoder

to_dict()

Serialize the module bytes and device for registry-based round trips.

Return type:

dict[str, Any]

classmethod from_dict(payload)

Rebuild a StreamingTransformer from to_dict() output.

Parameters:

payload (dict[str, Any])

Return type:

StreamingTransformer

class StreamingTransformerSampler(dist, seed=None)

Bases: DistributionSampler

Sampler facade for a conditional next-token transformer leaf.

Parameters:

• dist (StreamingTransformer)

• seed (int | None)

sample(size=None, *, batched=True)

Raise because contexts are required for transformer generation.

Parameters:

• size (int | None)

• batched (bool)

Return type:

Any

class StreamingTokenEncoder

Bases: DataSequenceEncoder

Encode context/token pairs for streaming transformer scoring and training.

seq_encode(data)

Convert (context, token) pairs into batched context and integer-token arrays.

Parameters:

data (list)

Return type:

tuple[ndarray, ndarray]

class StreamingTransformerAccumulator(module, lr, device)

Bases: SequenceEncodableStatisticAccumulator

seq_update = ONE train step against the PERSISTENT module + optimizer; value() = telemetry only.

Parameters:

• module (Any)

• lr (float)

• device (str)

seq_update(enc, weights, estimate)

Run one optimizer step on an encoded micro-batch.

Parameters:

• enc (Any)

• weights (Any)

• estimate (Any)

Return type:

None

update(x, weight, estimate)

Train on one weighted context/token pair through seq_update().

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

No-op initialization hook for the streaming training path.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

No-op batch initialization hook for the streaming training path.

Parameters:

• enc (Any)

• weights (Any)

• rng (Any)

Return type:

None

combine(other)

Merge telemetry from another streaming accumulator.

Parameters:

other (Any)

Return type:

StreamingTransformerAccumulator

value()

Return (loss_sum, token_count) telemetry without storing the corpus.

Return type:

tuple[float, int]

from_value(v)

Restore telemetry counters from a value tuple.

Parameters:

v (tuple)

Return type:

StreamingTransformerAccumulator

acc_to_encoder()

Return the encoder expected by this accumulator.

Return type:

StreamingTokenEncoder

class StreamingTransformerAccumulatorFactory(module, lr, device)

Bases: StatisticAccumulatorFactory

Factory for streaming transformer accumulators sharing a live module.

Parameters:

• module (Any)

• lr (float)

• device (str)

make()

Create a fresh accumulator around the shared live module.

Return type:

StreamingTransformerAccumulator

class StreamingTransformerEstimator(module, lr=3e-3, device='cpu')

Bases: ParameterEstimator

Estimator whose accumulator trains a live streaming transformer module in place.

Parameters:

• module (Any)

• lr (float)

• device (str)

accumulator_factory()

Return an accumulator factory for streamed context/token micro-batches.

Return type:

StreamingTransformerAccumulatorFactory

estimate(nobs, suff_stat)

Return the live module wrapped as a fitted streaming transformer leaf.

Parameters:

• nobs (float | None)

• suff_stat (tuple)

Return type:

StreamingTransformer

class TransformerLMEstimator(vocab, *, d_model=128, n_layer=4, n_head=4, block=64, embedding=None, lr=3e-3, device='cpu')

Bases: StreamingTransformerEstimator

A Transformer language model as a fit-ready estimator: TransformerLMEstimator(vocab, d_model=..., ...).

The clean, declarative surface – no hand-built torch module, no Leaf(...).estimator() two-step. Drops into MixtureEstimator/CompositeEstimator like any other *Estimator. embedding optionally ties a shared CategoricalEmbedding (e.g. one word embedding across a mixture’s experts). TransformerLMEstimator(V, embedding=emb) and StreamingTransformer.from_config(V, embedding=emb).estimator() build the same thing.

Parameters:

• vocab (int)

• d_model (int)

• n_layer (int)

• n_head (int)

• block (int)

• embedding (Any)

• lr (float)

• device (str)

stream_fit(module, token_source, *, lr=3e-3, device='cpu', report_every=200, log=None)

Train module by streaming micro-batches from token_source (a generator). The accumulator holds the PERSISTENT optimizer and trains incrementally; its payload stays (loss_sum, tokens) – the corpus is never buffered. Returns (StreamingTransformer, (loss_sum, tokens)).

Parameters:

• module (Any)

• token_source (Any)

• lr (float)

• device (str)

• report_every (int)

• log (Any)

Return type:

tuple

StreamingTransformerLeaf

alias of StreamingTransformer

StreamingTransformerLeafEstimator

alias of StreamingTransformerEstimator