mixle.models.streaming_transformer_leaf module

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

Where NeuralCategorical buffers a shard, this leaf keeps the corpus out of the accumulator. Each seq_update computes a weighted gradient numerator at a pinned base-model digest without changing the module. value() returns model-sized gradients plus weighted telemetry; combine() sums only compatible worker states, and estimate() applies one step through a persistent AdamW optimizer. Thus partitioned and unpartitioned accumulation have explicit data-parallel semantics instead of silently discarding worker-trained weights. stream_fit() preserves the familiar one-step-per-micro-batch behavior by estimating after every streamed batch.

class StreamingGradientState(base_sha256, parameter_names, gradient_sums, loss_sum, effective_weight, rows, batches)

Bases: object

Foldable weighted-gradient state computed at one exact model revision.

Parameters:

• base_sha256 (str)

• parameter_names (tuple[str, ...])

• gradient_sums (tuple[ndarray, ...])

• loss_sum (float)

• effective_weight (float)

• rows (int)

• batches (int)

property mean_loss: float

Return the consistently weighted mean loss.

class StreamingTransformer(module, device='cpu', *, lr=3.0e-3, optimizer_state=None)

Bases: SequenceEncodableProbabilityDistribution

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

Parameters:

• module (Any)

• device (str)

• lr (float)

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

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

Accumulate foldable weighted gradients at one exact model revision.

Parameters:

• module (Any)

• lr (float)

• device (str)

seq_update(enc, weights, estimate)

Accumulate one weighted gradient numerator without mutating the shared model.

Parameters:

• enc (Any)

• weights (Any)

• estimate (Any)

Return type:

None

update(x, weight, estimate)

Accumulate one weighted context/token gradient.

Parameters:

• x (Any)

• weight (float)

• estimate (Any)

Return type:

None

initialize(x, weight, rng)

Initialize through the same gradient accumulation path.

Parameters:

• x (Any)

• weight (float)

• rng (Any)

Return type:

None

seq_initialize(enc, weights, rng)

Initialize a batch through the same gradient accumulation path.

Parameters:

• enc (Any)

• weights (Any)

• rng (Any)

Return type:

None

combine(other)

Sum worker gradients only when they were computed at the identical base revision.

Parameters:

other (Any)

Return type:

StreamingTransformerAccumulator

value()

Return model-sized gradients and weighted telemetry, never the corpus.

Return type:

StreamingGradientState

from_value(v)

Restore a complete foldable gradient state.

Parameters:

v (Any)

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', optimizer_state=None)

Bases: ParameterEstimator

Persistent-optimizer estimator applying foldable worker gradient sums.

Parameters:

• module (Any)

• lr (float)

• device (str)

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

accumulator_factory()

Return workers pinned to the estimator’s current model revision.

Return type:

StreamingTransformerAccumulatorFactory

estimate(nobs, suff_stat)

Apply one persistent AdamW step from aggregated weighted gradient numerators.

Parameters:

• nobs (float | None)

• suff_stat (Any)

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 one persistent-optimizer step per streamed micro-batch without buffering the corpus.

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