File size: 4,717 Bytes
3699b0a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """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)
|