| """Construct native codec prompts and shifted target masks.""" |
| from __future__ import annotations |
| from collections.abc import Mapping, Sequence |
| from dataclasses import dataclass |
| from typing import Final |
| import torch |
| from torch import Tensor |
| CODEBOOKS: Final = 8 |
| CODE_AXES: Final = 2 |
| CODE_VALUES: Final = 1024 |
| VOCABULARY: Final = 1026 |
| BOC: Final = 1024 |
| EOC: Final = 1025 |
| IGNORE: Final = -100 |
| SPECIALS: Final = ('<|tts|>', '<|ref_text|>', '<|ref_audio|>', '<|text|>', '<|audio|>') |
|
|
| @dataclass(frozen=True) |
| class SequenceLayout: |
| text_ids: Tensor |
| audio_codes: Tensor |
| audio_mask: Tensor |
| labels: Tensor |
|
|
| @property |
| def positions(self) -> int: |
| return self.text_ids.numel() |
|
|
| @property |
| def targets(self) -> int: |
| return int(self.labels.ne(IGNORE).sum().item()) |
|
|
| def validate_codes(codes: Tensor) -> None: |
| if codes.ndim != CODE_AXES or codes.shape[1] != CODEBOOKS or codes.shape[0] == 0: |
| message = f'codec must have nonempty [frames,{CODEBOOKS}] shape, got {tuple(codes.shape)}' |
| raise ValueError(message) |
| if codes.dtype not in (torch.int64, torch.int32, torch.int16, torch.uint16): |
| message = f'codec must be integer, got {codes.dtype}' |
| raise TypeError(message) |
| values = codes.to(torch.long) |
| if bool(((values < 0) | (values >= CODE_VALUES)).any().item()): |
| bad = values[(values < 0) | (values >= CODE_VALUES)][0].item() |
| message = f'raw codec value={bad} outside [0,{CODE_VALUES})' |
| raise ValueError(message) |
|
|
| def delay_codes(codes: Tensor, *, terminal: bool=False) -> tuple[Tensor, Tensor]: |
| validate_codes(codes) |
| frames = codes.shape[0] |
| rows = frames + CODEBOOKS - 1 + int(terminal) |
| delayed = torch.full((rows, CODEBOOKS), EOC, dtype=torch.long, device=codes.device) |
| labels = torch.full_like(delayed, IGNORE) |
| for stream in range(CODEBOOKS): |
| delayed[:stream, stream] = BOC |
| delayed[stream:stream + frames, stream] = codes[:, stream].long() |
| labels[stream:stream + frames, stream] = codes[:, stream].long() |
| if terminal: |
| labels[stream + frames, stream] = EOC |
| return (delayed, labels) |
|
|
| def reverse_delay(delayed: Tensor, frames: int) -> Tensor: |
| if frames <= 0 or delayed.shape != (frames + CODEBOOKS - 1, CODEBOOKS): |
| message = f'invalid delayed shape={tuple(delayed.shape)} for frames={frames}' |
| raise ValueError(message) |
| result = torch.stack([delayed[c:c + frames, c] for c in range(CODEBOOKS)], dim=1) |
| validate_codes(result) |
| return result |
|
|
| def prompt_layout(tokens: Mapping[str, int], text_ids: Sequence[int], reference_codes: Tensor | None=None, reference_text_ids: Sequence[int]=()) -> SequenceLayout: |
| missing = set(SPECIALS) - tokens.keys() |
| if missing or not text_ids: |
| message = f'invalid prompt: missing specials={sorted(missing)}, text length={len(text_ids)}' |
| raise ValueError(message) |
| if any((token < 0 for token in (*tokens.values(), *text_ids, *reference_text_ids))): |
| message = 'prompt token IDs must be nonnegative' |
| raise ValueError(message) |
| prefix = [tokens['<|tts|>']] |
| if reference_codes is not None and reference_text_ids: |
| prefix += [tokens['<|ref_text|>'], *reference_text_ids] |
| reference = delay_codes(reference_codes)[0] if reference_codes is not None else None |
| start = 0 |
| if reference is not None: |
| prefix.append(tokens['<|ref_audio|>']) |
| start = len(prefix) |
| prefix.extend([0] * reference.shape[0]) |
| prefix += [tokens['<|text|>'], *text_ids, tokens['<|audio|>']] |
| codes = torch.zeros((len(prefix), CODEBOOKS), dtype=torch.long) |
| mask = torch.zeros(len(prefix), dtype=torch.bool) |
| if reference is not None: |
| codes[start:start + reference.shape[0]] = reference.cpu() |
| mask[start:start + reference.shape[0]] = True |
| return SequenceLayout(torch.tensor(prefix, dtype=torch.long), codes, mask, torch.full_like(codes, IGNORE)) |
|
|
| def target_layout(prompt: SequenceLayout, target: Tensor, max_positions: int=8192) -> SequenceLayout: |
| delayed, supervision = delay_codes(target, terminal=True) |
| rows = delayed.shape[0] - 1 |
| labels = torch.full((prompt.positions + rows, CODEBOOKS), IGNORE, dtype=torch.long) |
| labels[prompt.positions - 1:] = supervision.cpu() |
| positions = prompt.positions + rows |
| if positions > max_positions: |
| message = f'mixed positions={positions} exceeds max_positions={max_positions}' |
| raise ValueError(message) |
| return SequenceLayout(torch.cat((prompt.text_ids, torch.zeros(rows, dtype=torch.long))), torch.cat((prompt.audio_codes, delayed[:-1].cpu())), torch.cat((prompt.audio_mask, torch.ones(rows, dtype=torch.bool))), labels) |
|
|
