| """Constrained decoding: a per-sequence structure state machine + token sampling.
|
|
|
| The state machine guarantees every generated media block is well formed: inside <image> only
|
| visual codes are allowed and exactly G of them; inside a video, frames are
|
| <frame><t_k> G codes </frame> with consecutive timestamps; text spans can only contain text
|
| tokens. Purely structural tokens are *forced* (no sampling); content tokens are sampled from
|
| the model's distribution restricted to the legal set.
|
| """
|
| from __future__ import annotations
|
|
|
| from dataclasses import dataclass
|
|
|
| import torch
|
|
|
| from ..tokenizers.multimodal import MAX_FRAMES, UnifiedVocab
|
|
|
|
|
| @dataclass
|
| class Sampling:
|
| temperature: float = 1.0
|
| top_k: int = 0
|
| top_p: float = 1.0
|
| repetition_penalty: float = 1.0
|
| max_new_tokens: int = 256
|
| seed: int = 0
|
| cfg_scale: float = 1.0
|
|
|
|
|
| class Masks:
|
| """Boolean legal-token sets, precomputed once per vocabulary/device."""
|
|
|
| def __init__(self, v: UnifiedVocab, device):
|
| V = v.size
|
|
|
| def mk(ids=(), ranges=()):
|
| m = torch.zeros(V, dtype=torch.bool, device=device)
|
| for lo, hi in ranges:
|
| m[lo:hi] = True
|
| if ids:
|
| m[list(ids)] = True
|
| return m
|
|
|
| self.text = mk([v.eos], [(0, v.text_size)])
|
| self.text_media = mk([v.eos, v.img_start, v.vid_start], [(0, v.text_size)])
|
| self.visual = mk(ranges=[(v.visual_offset, v.visual_offset + v.visual_size)])
|
| self.timestamp = mk(ranges=[(v.timestamp(0), v.timestamp(0) + MAX_FRAMES)])
|
| self.frame_or_end = mk([v.frm_start, v.vid_end])
|
|
|
|
|
| class StructureFSM:
|
| """Tracks where one sequence is in the media grammar.
|
|
|
| mode:
|
| "text" free text until <eos> (media tags disallowed unless allow_media)
|
| "image" we are right after <image>: G codes, </image>, then <eos> if stop_after_media
|
| "video" we are inside <video> right before the first <frame>
|
| n_frames: exact total frame count for videos (None = model decides, bounded by max_timestamp)
|
| first_ts: timestamp of the next frame to generate (None = sampled by the model)
|
| frames_done: frames already present in the prefix (prefix then ends right after a </frame>)
|
| codes_done: image codes already present after <image> (image completion)
|
| """
|
|
|
| def __init__(self, v: UnifiedVocab, mode: str = "text", allow_media: bool = False, stop_after_media: bool = True,
|
| n_frames: int | None = None, first_ts: int | None = None, max_timestamp: int = 7, frames_done: int = 0, codes_done: int = 0):
|
| self.v, self.allow_media, self.stop_after_media = v, allow_media, stop_after_media
|
| self.n_frames, self.next_ts, self.max_ts = n_frames, first_ts, max_timestamp
|
| self.done = False
|
| self.frames = frames_done
|
| self.state = {"text": "text", "image": "img_codes", "video": "after_frame" if frames_done else "frame_start"}[mode]
|
| self.count = codes_done if mode == "image" else 0
|
|
|
| def constraint(self, masks: Masks) -> tuple[int | None, torch.Tensor | None]:
|
| """(forced_token, None) or (None, legal_mask)."""
|
| v, s = self.v, self.state
|
| if self.done:
|
| return v.pad, None
|
| if s == "text":
|
| return None, masks.text_media if self.allow_media else masks.text
|
| if s in ("img_codes", "frm_codes"):
|
| return None, masks.visual
|
| if s == "img_end":
|
| return v.img_end, None
|
| if s == "frm_end":
|
| return v.frm_end, None
|
| if s == "frame_start":
|
| return v.frm_start, None
|
| if s == "timestamp":
|
| return (v.timestamp(self.next_ts), None) if self.next_ts is not None else (None, masks.timestamp)
|
| if s == "after_frame":
|
| last = self.next_ts - 1
|
| if (self.n_frames is not None and self.frames >= self.n_frames) or last >= self.max_ts:
|
| return v.vid_end, None
|
| return (v.frm_start, None) if self.n_frames is not None else (None, masks.frame_or_end)
|
| if s == "eos":
|
| return v.eos, None
|
| raise AssertionError(s)
|
|
|
| def advance(self, tok: int) -> None:
|
| v, s = self.v, self.state
|
| if self.done:
|
| return
|
| G = v.tokens_per_frame
|
| if s == "text":
|
| if tok == v.eos:
|
| self.done = True
|
| elif tok == v.img_start:
|
| self.state, self.count = "img_codes", 0
|
| elif tok == v.vid_start:
|
| self.state = "frame_start"
|
| elif s == "img_codes":
|
| self.count += 1
|
| if self.count == G:
|
| self.state = "img_end"
|
| elif s == "img_end":
|
| self.state = "eos" if self.stop_after_media else "text"
|
| elif s == "frame_start":
|
| self.state = "timestamp"
|
| elif s == "timestamp":
|
| self.next_ts = tok - v.timestamp(0) + 1
|
| self.state, self.count = "frm_codes", 0
|
| elif s == "frm_codes":
|
| self.count += 1
|
| if self.count == G:
|
| self.state = "frm_end"
|
| elif s == "frm_end":
|
| self.frames += 1
|
| self.state = "after_frame"
|
| elif s == "after_frame":
|
| if tok == v.frm_start:
|
| self.state = "timestamp"
|
| else:
|
| self.state = "eos" if self.stop_after_media else "text"
|
| elif s == "eos":
|
| self.done = True
|
|
|
| @property
|
| def in_visual(self) -> bool:
|
| return self.state in ("img_codes", "frm_codes")
|
|
|
|
|
| def sample_next(logits: torch.Tensor, legal: torch.Tensor, sp: Sampling, gen: torch.Generator,
|
| history: torch.Tensor | None = None, text_size: int = 0) -> torch.Tensor:
|
| """logits [B,V] (float32), legal [B,V] bool -> sampled ids [B]."""
|
| logits = logits.float().clone()
|
| if sp.repetition_penalty != 1.0 and history is not None and history.numel():
|
| prev = history.clamp(min=0)
|
| is_text = prev < text_size
|
| vals = logits.gather(1, prev)
|
| pen = torch.where(vals > 0, vals / sp.repetition_penalty, vals * sp.repetition_penalty)
|
| logits.scatter_(1, prev, torch.where(is_text, pen, vals))
|
| logits = logits.masked_fill(~legal, float("-inf"))
|
| if sp.temperature <= 0:
|
| return logits.argmax(-1)
|
| logits = logits / sp.temperature
|
| if sp.top_k > 0:
|
| kth = logits.topk(min(sp.top_k, logits.size(-1)), dim=-1).values[:, -1:]
|
| logits = logits.masked_fill(logits < kth, float("-inf"))
|
| if sp.top_p < 1.0:
|
| srt, idx = logits.sort(dim=-1, descending=True)
|
| probs = srt.softmax(-1)
|
| drop = probs.cumsum(-1) - probs > sp.top_p
|
| logits = logits.masked_fill(drop.scatter(1, idx, drop), float("-inf"))
|
| return torch.multinomial(logits.softmax(-1), 1, generator=gen).squeeze(1)
|
|
|
