File size: 7,170 Bytes
d2a1a41 | 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 | """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 # applied to previously generated *text* tokens only
max_new_tokens: int = 256
seed: int = 0
cfg_scale: float = 1.0 # classifier-free guidance on visual tokens (1.0 = off)
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: # </video>
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 # keep the smallest prefix with mass >= top_p
logits = logits.masked_fill(drop.scatter(1, idx, drop), float("-inf"))
return torch.multinomial(logits.softmax(-1), 1, generator=gen).squeeze(1)
|