mstar.model.omnivoice.components.unmask

mstar.model.omnivoice.components.unmask#

The unmask step: CFG scoring, confidence ranking, and the reveal schedule.

An exact port of OmniVoice’s _predict_tokens_with_scoring and the schedule arithmetic inside _generate_iterative (omnivoice/models/omnivoice.py). Parity against the reference is decided here, so the operation order is kept verbatim even where it could be tidied — the log_softmax is applied twice on the CFG path because the reference applies it twice, and the mask id is driven to -inf after the combine rather than before.

Nothing here touches the network or the engine; it is pure tensor math over one request’s logits, so the parity test can drive it directly.

Functions

apply_reveal(tokens, pred_tokens, scores, ...)

Reveal reveal_count cells of tokens in place and return it.

build_reveal_schedule(target_len, ...)

How many of the target_len * num_codebook cells to reveal per step.

filter_top_k(logits[, ratio])

Keep the top ratio of the vocabulary, -inf elsewhere.

get_time_steps(num_step, t_shift[, t_start, ...])

The shifted timestep ramp, num_step + 1 points on [t_start, t_end].

gumbel_sample(logits, temperature[, generator])

Add Gumbel noise at temperature; caller takes the argmax.

predict_tokens_with_scoring(c_logits, ...[, ...])

CFG-combine conditional and unconditional logits, then rank.
mstar.model.omnivoice.components.unmask.apply_reveal(tokens, pred_tokens, scores, reveal_count, audio_mask_id, layer_penalty_factor, position_temperature, generator=None)[source]#

Reveal reveal_count cells of tokens in place and return it.

Ranking is over the flattened (C, T) canvas, so a step may spend its whole budget on one codebook row. Two biases shape the choice, both from the reference:

  • layer_penalty_factor subtracts layer_index * factor from the score, pushing the earliest codebooks to resolve first (they carry the coarse content the later residual rows refine).

  • position_temperature adds Gumbel noise to the position ranking, so the reveal order is stochastic even when token choice is greedy.

Already-revealed cells are driven to -inf so a cell is only written once.

Parameters:

  • tokens (Tensor) – [1, C, T], the request’s live canvas, modified in place.

  • pred_tokens (Tensor) – [1, C, T] from predict_tokens_with_scoring.

  • scores (Tensor) – [1, C, T] from predict_tokens_with_scoring.

  • generator (Generator | None) – the request’s seeded RNG, or None for the global one.

  • reveal_count (int)

  • audio_mask_id (int)

  • layer_penalty_factor (float)

  • position_temperature (float)

Return type:

Tensor

mstar.model.omnivoice.components.unmask.build_reveal_schedule(target_len, num_codebook, num_step, t_shift)[source]#

How many of the target_len * num_codebook cells to reveal per step.

The per-step count follows the gap between consecutive timesteps, rounded up and clamped by what is left; the final step takes the entire remainder so the canvas is always fully revealed regardless of rounding drift.

Cached, and returning a tuple so it stays immutable: the schedule depends only on these four numbers, but postprocess needs it on every step of every request, and recomputing it there means a linspace and a Python loop per step for a list that never changes within a request.

Parameters:
Return type:

tuple[int, …]

mstar.model.omnivoice.components.unmask.filter_top_k(logits, ratio=0.1)[source]#

Keep the top ratio of the vocabulary, -inf elsewhere.

Parameters:
Return type:

Tensor

mstar.model.omnivoice.components.unmask.get_time_steps(num_step, t_shift, t_start=0.0, t_end=1.0, device=None)[source]#

The shifted timestep ramp, num_step + 1 points on [t_start, t_end].

t_shift below 1 pushes points toward t_start, spending more of the budget at low SNR where the canvas is still mostly masked.

Parameters:
Return type:

Tensor

mstar.model.omnivoice.components.unmask.gumbel_sample(logits, temperature, generator=None)[source]#

Add Gumbel noise at temperature; caller takes the argmax.

generator carries the request’s seed. Both draws in a step (token choice and reveal order) come from it, so a seeded request replays exactly; passing None keeps the global RNG and stays unreproducible.

Parameters:
Return type:

Tensor

mstar.model.omnivoice.components.unmask.predict_tokens_with_scoring(c_logits, u_logits, audio_mask_id, guidance_scale, class_temperature, generator=None)[source]#

CFG-combine conditional and unconditional logits, then rank.

Parameters:

  • c_logits (Tensor) – conditional logits, [1, C, T, V].

  • u_logits (Tensor) – unconditional logits, same shape. Ignored when guidance_scale is 0.

  • audio_mask_id (int) – the MASK class, forced to -inf so it is never predicted.

  • guidance_scale (float) – 0 disables CFG entirely (and the caller may then skip the unconditional forward).

  • class_temperature (float) – 0 takes the argmax; above 0 samples from the top decile with Gumbel noise.

  • generator (Generator | None) – the request’s seeded RNG, or None for the global one.

Returns:

(pred_tokens, confidence_scores), each [1, C, T]. The scores are max log-probabilities and drive which cells get revealed.

Return type:

tuple[Tensor, Tensor]