File size: 12,250 Bytes
7ccb33d | 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 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 | """Speaker-neutral grammar constraints for MOSS transcript generation.
The constraint enforces only the wire syntax. It never chooses a speaker,
timestamp, or text token on the model's behalf. Token transitions operate on
decoded token pieces so merged BPE tokens such as ``][`` and ``[S`` remain
valid when they cross grammar boundaries.
"""
from __future__ import annotations
import copy
import hashlib
import time
from dataclasses import dataclass
from typing import Iterable
import torch
from transformers import LogitsProcessor
SCHEMA = "moss-transcript-grammar-v1"
@dataclass(frozen=True)
class GrammarState:
phase: str = "start"
def _digit(value: str) -> bool:
return "0" <= value <= "9"
def advance_character(state: GrammarState, char: str) -> GrammarState | None:
"""Advance one Unicode character, or return ``None`` for invalid syntax."""
phase = state.phase
if phase == "start":
return GrammarState("start_time_first") if char == "[" else None
if phase == "start_time_first":
return GrammarState("start_time_int") if _digit(char) else None
if phase == "start_time_int":
if _digit(char):
return state
if char == ".":
return GrammarState("start_time_frac_first")
if char == "]":
return GrammarState("after_start_time")
return None
if phase == "start_time_frac_first":
return GrammarState("start_time_frac") if _digit(char) else None
if phase == "start_time_frac":
if _digit(char):
return state
return GrammarState("after_start_time") if char == "]" else None
if phase == "after_start_time":
if char.isspace():
return state
return GrammarState("speaker_s") if char == "[" else None
if phase == "speaker_s":
return GrammarState("speaker_digit_1") if char == "S" else None
if phase == "speaker_digit_1":
return GrammarState("speaker_digit_2") if _digit(char) else None
if phase == "speaker_digit_2":
return GrammarState("speaker_close") if _digit(char) else None
if phase == "speaker_close":
return GrammarState("text_empty") if char == "]" else None
if phase == "text_empty":
if char in "[]":
return None
return state if char.isspace() else GrammarState("text")
if phase == "text":
if char == "[":
return GrammarState("end_time_first")
if char == "]":
return None
return state
if phase == "end_time_first":
return GrammarState("end_time_int") if _digit(char) else None
if phase == "end_time_int":
if _digit(char):
return state
if char == ".":
return GrammarState("end_time_frac_first")
if char == "]":
return GrammarState("after_end_time")
return None
if phase == "end_time_frac_first":
return GrammarState("end_time_frac") if _digit(char) else None
if phase == "end_time_frac":
if _digit(char):
return state
return GrammarState("after_end_time") if char == "]" else None
if phase == "after_end_time":
if char.isspace():
return state
return GrammarState("start_time_first") if char == "[" else None
raise ValueError(f"unknown transcript grammar phase: {phase}")
def advance_piece(state: GrammarState, piece: str) -> GrammarState | None:
if not piece:
return None
for char in piece:
state = advance_character(state, char)
if state is None:
return None
return state
def accepting(state: GrammarState) -> bool:
return state.phase == "after_end_time"
class TranscriptGrammarVocabulary:
"""Tokenizer-specific transition cache shared across utterances."""
def __init__(self, tokenizer):
self.tokenizer = tokenizer
self.vocab_size = len(tokenizer)
eos = tokenizer.eos_token_id
self.eos_token_ids = {int(value) for value in (eos if isinstance(eos, list) else [eos])}
self.special_ids = {int(value) for value in getattr(tokenizer, "all_special_ids", [])}
self.pieces = tuple(
tokenizer.decode(
[token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False
)
for token_id in range(self.vocab_size)
)
digest = hashlib.sha256()
for token_id, piece in enumerate(self.pieces):
digest.update(str(token_id).encode("ascii"))
digest.update(b"\0")
digest.update(piece.encode("utf-8", errors="surrogatepass"))
digest.update(b"\n")
self.token_surface_sha256 = digest.hexdigest()
self._allowed_cpu: dict[GrammarState, tuple[int, ...]] = {}
self._allowed_device: dict[tuple[GrammarState, str], torch.Tensor] = {}
def consume(self, token_ids: Iterable[int]) -> GrammarState:
state = GrammarState()
for token_id in token_ids:
token_id = int(token_id)
if token_id in self.eos_token_ids:
if not accepting(state):
raise ValueError("EOS before a complete end timestamp")
continue
if token_id in self.special_ids or not 0 <= token_id < self.vocab_size:
raise ValueError(f"invalid special/out-of-vocabulary token in transcript: {token_id}")
next_state = advance_piece(state, self.pieces[token_id])
if next_state is None:
raise ValueError(
f"token {token_id} piece={self.pieces[token_id]!r} violates phase={state.phase}"
)
state = next_state
return state
def allowed_token_ids(self, state: GrammarState, device: torch.device) -> torch.Tensor:
if state not in self._allowed_cpu:
allowed = []
for token_id, piece in enumerate(self.pieces):
if token_id in self.special_ids:
continue
if advance_piece(state, piece) is not None:
allowed.append(token_id)
# An empty response is not a valid transcript for deployment and
# the parser cannot score it. EOS therefore becomes available
# only after the model has completed a full segment.
if accepting(state):
allowed.extend(self.eos_token_ids)
if not allowed:
raise RuntimeError(f"transcript grammar has no continuation from {state.phase}")
self._allowed_cpu[state] = tuple(sorted(set(allowed)))
key = (state, str(device))
if key not in self._allowed_device:
self._allowed_device[key] = torch.tensor(
self._allowed_cpu[state], dtype=torch.long, device=device
)
return self._allowed_device[key]
def audit(self) -> dict:
return {
"schema": SCHEMA,
"speaker_policy": "model_selected_exact_two_digit_tag",
"timestamp_policy": "model_selected_nonnegative_number",
"text_policy": "model_selected_nonempty_no_brackets",
"eos_policy": "complete_segment_only",
"vocab_size": self.vocab_size,
"token_surface_sha256": self.token_surface_sha256,
"eos_token_ids": sorted(self.eos_token_ids),
}
class TranscriptGrammarLogitsProcessor(LogitsProcessor):
"""Mask tokens that cannot extend the canonical transcript grammar."""
def __init__(
self,
tokenizer,
prompt_length: int,
*,
vocabulary: TranscriptGrammarVocabulary | None = None,
):
if prompt_length < 0:
raise ValueError("prompt_length must be non-negative")
self.prompt_length = int(prompt_length)
self.vocabulary = vocabulary or TranscriptGrammarVocabulary(tokenizer)
if self.vocabulary.tokenizer is not tokenizer:
raise ValueError("constraint vocabulary belongs to another tokenizer instance")
self.tokenizer = tokenizer
self.vocab_size = self.vocabulary.vocab_size
self.eos_token_ids = self.vocabulary.eos_token_ids
def consume(self, token_ids: Iterable[int]) -> GrammarState:
return self.vocabulary.consume(token_ids)
def allowed_token_ids(self, state: GrammarState, device: torch.device) -> torch.Tensor:
return self.vocabulary.allowed_token_ids(state, device)
def audit(self) -> dict:
return self.vocabulary.audit()
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
if input_ids.ndim != 2 or scores.ndim != 2 or input_ids.shape[0] != scores.shape[0]:
raise ValueError("transcript grammar requires aligned 2-D token and score batches")
constrained = torch.full_like(scores, -torch.inf)
# Hugging Face expands a batch-one prompt to one row per beam. Replay
# each prefix independently so beam reordering cannot leak DFA state.
for row in range(input_ids.shape[0]):
generated = input_ids[row, self.prompt_length :].tolist()
state = self.consume(generated)
allowed = self.allowed_token_ids(state, scores.device)
constrained[row].index_copy_(0, allowed, scores[row].index_select(0, allowed))
return constrained
def generate_constrained_transcription(
model,
processor,
messages,
*,
max_length: int = 131072,
max_new_tokens: int,
device: torch.device,
dtype: torch.dtype,
constraint_vocabulary: TranscriptGrammarVocabulary | None = None,
num_beams: int = 1,
) -> dict:
"""Generate one transcript with syntax-only constraints."""
from moss_transcribe_diarize.inference_utils import prepare_inputs
preprocess_context = (
torch.amp.autocast("cuda", dtype=dtype)
if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16)
else torch.no_grad()
)
with preprocess_context:
inputs = prepare_inputs(
processor, messages, max_length=max_length, device=device
).to(device)
prompt_length = int(inputs["attention_mask"][0].sum().item())
constraint = TranscriptGrammarLogitsProcessor(
processor.tokenizer,
prompt_length,
vocabulary=constraint_vocabulary,
)
generation_config = copy.deepcopy(model.generation_config)
generation_config.max_new_tokens = int(max_new_tokens)
generation_config.do_sample = False
generation_config.num_beams = int(num_beams)
generation_config.num_return_sequences = 1
if num_beams > 1:
generation_config.early_stopping = True
generation_context = (
torch.amp.autocast("cuda", dtype=dtype)
if device.type == "cuda" and dtype in (torch.float16, torch.bfloat16)
else torch.no_grad()
)
started = time.perf_counter()
with torch.inference_mode(), generation_context:
outputs = model.generate(
input_ids=inputs["input_ids"],
attention_mask=inputs["attention_mask"],
input_features=inputs["input_features"],
audio_feature_lengths=inputs["audio_feature_lengths"],
audio_chunk_mapping=inputs["audio_chunk_mapping"],
generation_config=generation_config,
logits_processor=[constraint],
)
if device.type == "cuda":
torch.cuda.synchronize(device)
generated_ids = outputs[0][prompt_length:]
non_eos = [
int(token_id)
for token_id in generated_ids.tolist()
if int(token_id) not in constraint.eos_token_ids
]
final_state = constraint.consume(non_eos)
text = processor.tokenizer.decode(
generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
).strip()
return {
"text": text,
"prompt_len": prompt_length,
"generated_tokens": int(generated_ids.numel()),
"inference_seconds": time.perf_counter() - started,
"constraint_complete": accepting(final_state),
"constraint_final_phase": final_state.phase,
"constraint_num_beams": int(num_beams),
"constraint_audit": constraint.audit(),
}
|