| """Generate native delayed codec streams with isolated request state.""" |
| from __future__ import annotations |
| import time |
| from dataclasses import dataclass |
| import torch |
| from pydantic import Field |
| from torch import Tensor |
| from transformers.cache_utils import Cache |
| from transformers.modeling_outputs import BaseModelOutputWithPast |
| from .backbone import Generator |
| from .configuration import Specification |
| from .layout import BOC, CODEBOOKS, EOC, SequenceLayout, reverse_delay |
| MAX_CONTEXT_POSITIONS = 8192 |
| COLLAPSE_FRAMES = 25 |
|
|
| class GenerationError(RuntimeError): |
|
|
| def __init__(self, message: str, trace: dict[str, object]) -> None: |
| super().__init__(message) |
| self.trace = trace |
|
|
| @dataclass |
| class LoopState: |
| is_recovering: bool = False |
|
|
| class Sampling(Specification): |
| temperature: float = Field(default=0.8, ge=0) |
| top_p: float = Field(default=0.8, gt=0, le=1) |
| top_k: int = Field(default=30, ge=0) |
| max_new_positions: int = Field(default=1500, ge=8) |
| seed: int = Field(default=20261102, ge=0) |
|
|
| class Decoding(Specification): |
| alignment_strength: float = Field(default=0, ge=0, le=3) |
| plausible_tokens: int = Field(default=25, ge=1, le=1026) |
| recover_loops: bool = False |
|
|
| def alignment_logits(expert: Tensor, amateur: Tensor, settings: Decoding) -> Tensor: |
| if expert.shape != amateur.shape: |
| message = f'ECCD logit shapes differ: {tuple(expert.shape)}, {tuple(amateur.shape)}' |
| raise ValueError(message) |
| if settings.alignment_strength == 0: |
| return expert |
| positive = expert.float().log_softmax(-1) |
| negative = amateur.float().log_softmax(-1) |
| indices = positive.topk(min(settings.plausible_tokens, positive.shape[-1]), dim=-1).indices |
| head = torch.zeros_like(positive, dtype=torch.bool).scatter_(-1, indices, True) |
| compatibility = negative.exp().masked_fill(~head, 0).sum(-1, keepdim=True) |
| enhancement = (positive - negative).clamp_min(0).masked_fill(~head, 0) |
| return positive + settings.alignment_strength * (1 - compatibility) * enhancement |
|
|
| def recover_repetition(logits: Tensor, row: Tensor, history: Tensor, generator: torch.Generator, state: LoopState | None=None) -> tuple[Tensor, bool]: |
| collapsed = history.numel() >= COLLAPSE_FRAMES and bool(history[-COLLAPSE_FRAMES:].eq(row[0]).all()) |
| if state is not None: |
| state.is_recovering |= collapsed |
| collapsed = state.is_recovering and bool(history[-10:].eq(row[0]).any()) |
| if not collapsed: |
| state.is_recovering = False |
| if not collapsed or int(row[0]) >= BOC: |
| return (row, False) |
| alternative = logits[0].float().clone() |
| alternative[row[0]] = float('-inf') |
| result = row.clone() |
| result[0] = torch.multinomial(alternative.softmax(-1), 1, generator=generator)[0] |
| return (result, True) |
|
|
| @dataclass |
| class SamplerState: |
| delay_count: int = 0 |
| eoc_countdown: int | None = None |
| is_done: bool = False |
|
|
| def sample_row(logits: Tensor, settings: Sampling, generator: torch.Generator) -> Tensor: |
| if settings.temperature == 0: |
| return logits.argmax(dim=-1) |
| scaled = logits.float() / settings.temperature |
| if settings.top_k: |
| threshold = scaled.topk(min(settings.top_k, scaled.shape[-1]), dim=-1).values[:, -1:] |
| scaled = scaled.masked_fill(scaled < threshold, float('-inf')) |
| ordered, indices = scaled.sort(descending=True, dim=-1) |
| remove = ordered.softmax(dim=-1).cumsum(dim=-1) > settings.top_p |
| remove[:, 1:] = remove[:, :-1].clone() |
| remove[:, 0] = False |
| mask = torch.zeros_like(remove).scatter(-1, indices, remove) |
| probabilities = scaled.masked_fill(mask, float('-inf')).softmax(dim=-1) |
| return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1) |
|
|
| def advance_sampler(codes: Tensor, state: SamplerState) -> Tensor: |
| result = codes.clone() |
| if state.delay_count < CODEBOOKS: |
| state.delay_count += 1 |
| result[state.delay_count:] = BOC |
| elif state.eoc_countdown is not None: |
| state.eoc_countdown -= 1 |
| state.is_done = state.eoc_countdown <= 0 |
| elif int(result[0].item()) == EOC: |
| state.eoc_countdown = CODEBOOKS - 2 |
| return result |
|
|
| def unpack_output(output: object) -> tuple[Tensor, Cache]: |
| if not isinstance(output, BaseModelOutputWithPast): |
| message = f'invalid decoder output: {type(output).__name__}' |
| raise TypeError(message) |
| hidden, cache = (output.last_hidden_state, output.past_key_values) |
| if not isinstance(hidden, Tensor) or not isinstance(cache, Cache): |
| message = 'cached generation requires hidden states and an explicit KV cache' |
| raise TypeError(message) |
| return (hidden[:, -1], cache) |
|
|
| @torch.inference_mode() |
| def generate_codes(model: Generator, prompt: SequenceLayout, settings: Sampling, decoding: Decoding | None=None, contrast_prompt: SequenceLayout | None=None, *, deadline_monotonic: float | None=None) -> tuple[Tensor, dict[str, object]]: |
| decoding = decoding or Decoding() |
| if prompt.positions + settings.max_new_positions > MAX_CONTEXT_POSITIONS: |
| message = f'generation context exceeds {MAX_CONTEXT_POSITIONS}: prompt={prompt.positions}' |
| message += f', budget={settings.max_new_positions}' |
| raise ValueError(message) |
| device = model.audio_embedding.weight.device |
| if decoding.alignment_strength and contrast_prompt is None: |
| message = 'ECCD requires a text-ablated contrast prompt' |
| raise ValueError(message) |
| previous_mode = model.training |
| model.train(False) |
| generator = torch.Generator(device=device).manual_seed(settings.seed) |
| try: |
| with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == 'cuda'): |
| return decode_stream(model, prompt, settings, decoding, contrast_prompt, generator, deadline_monotonic) |
| finally: |
| model.train(previous_mode) |
|
|
| def prefill(model: Generator, prompt: SequenceLayout) -> tuple[Tensor, Cache]: |
| device = model.audio_embedding.weight.device |
| ids = prompt.text_ids.to(device).unsqueeze(0) |
| codes = prompt.audio_codes.to(device).unsqueeze(0) |
| mask = prompt.audio_mask.to(device).reshape(1, -1, 1) |
| text: Tensor = model.body.embed_tokens(ids) |
| embeddings = torch.where(mask, model.embed_audio(codes), text) |
| return unpack_output(model.body(inputs_embeds=embeddings, use_cache=True)) |
|
|
| def decode_stream(model: Generator, prompt: SequenceLayout, settings: Sampling, decoding: Decoding, contrast_prompt: SequenceLayout | None, generator: torch.Generator, deadline_monotonic: float | None) -> tuple[Tensor, dict[str, object]]: |
| device = model.audio_embedding.weight.device |
| hidden, cache = prefill(model, prompt) |
| contrast = prefill(model, contrast_prompt) if decoding.alignment_strength and contrast_prompt else None |
| state = SamplerState() |
| rows = torch.empty((settings.max_new_positions, CODEBOOKS), device=device, dtype=torch.long) |
| replacements = 0 |
| loop_state = LoopState() |
| count = 0 |
| for position in range(settings.max_new_positions): |
| if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic: |
| message = f'speech deadline exceeded at generated position={position}' |
| raise TimeoutError(message) |
| logits = model.classify(hidden)[0] |
| if contrast is not None: |
| guided = alignment_logits(logits[:1], model.classify(contrast[0])[0, :1], decoding) |
| logits = torch.cat((guided, logits[1:]), dim=0) |
| sampled = sample_row(logits, settings, generator) |
| if decoding.recover_loops: |
| sampled, recovered = recover_repetition(logits, sampled, rows[:position, 0], generator, loop_state) |
| replacements += int(recovered) |
| row = advance_sampler(sampled, state) |
| rows[position] = row |
| count = position + 1 |
| if state.is_done: |
| break |
| embedded = model.embed_audio(row.reshape(1, 1, -1)) |
| hidden, cache = unpack_output(model.body(inputs_embeds=embedded, past_key_values=cache, use_cache=True)) |
| if contrast is not None: |
| contrast = unpack_output(model.body(inputs_embeds=embedded, past_key_values=contrast[1], use_cache=True)) |
| if not state.is_done: |
| message = f'generation hit hard cap={settings.max_new_positions} without native termination' |
| raise GenerationError(message, {'terminated': False, 'trace': rows[:count].cpu().tolist(), 'generated_positions': count, 'sampling': settings.model_dump(), 'decoding': decoding.model_dump(), 'repetition_resamples': replacements}) |
| delayed = rows[:count].cpu() |
| first_eoc = delayed[:, 0].eq(EOC).nonzero() |
| if first_eoc.numel() == 0: |
| message = 'native termination has no first-stream EOC' |
| raise RuntimeError(message) |
| frames = int(first_eoc[0, 0].item()) |
| needed = frames + CODEBOOKS - 1 |
| if delayed.shape[0] < needed: |
| message = f'incomplete delayed wind-down: rows={delayed.shape[0]}, needed={needed}' |
| raise RuntimeError(message) |
| result = reverse_delay(delayed[:needed], frames) |
| return (result, {'terminated': True, 'frames': frames, 'trace': delayed.tolist(), 'sampling': settings.model_dump(), 'decoding': decoding.model_dump(), 'repetition_resamples': replacements, 'generated_positions': count}) |
|
|
