File size: 4,555 Bytes
1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 1bfc3f0 895e804 | 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 | """Explicit ActionCodec2Sparse token constraints for Hugging Face generate()."""
from __future__ import annotations
import torch
from transformers import LogitsProcessor
from ...backends.actioncodec2.cache import BoundedCache
class ActionCodec2SparseLogitsProcessor(LogitsProcessor):
"""Constrain generated action IDs to the sparse grammar; allow EOS once complete."""
def __init__(
self, codec, horizon, *, fps, prompt_length, eos_token_id, token_offset=0, pad_token_id=None
):
for name, value in {
"prompt_length": prompt_length,
"eos_token_id": eos_token_id,
"token_offset": token_offset,
}.items():
if type(value) is not int or value < 0:
raise ValueError(f"{name} must be a nonnegative integer")
if pad_token_id is not None and (type(pad_token_id) is not int or pad_token_id < 0):
raise ValueError("pad_token_id must be a nonnegative integer")
self.grammar = codec.grammar(horizon, fps=fps)
self.prompt_length, self.token_offset = prompt_length, token_offset
self.eos_token_id, self.pad_token_id = eos_token_id, pad_token_id
stop = token_offset + self.grammar.vocab_size
if token_offset <= eos_token_id < stop:
raise ValueError("eos_token_id must be outside the action token interval")
if pad_token_id is not None and token_offset <= pad_token_id < stop:
raise ValueError("pad_token_id must be outside the action token interval")
self._mask_cache = BoundedCache(64)
def _state(self, row):
if self.eos_token_id in row:
end = row.index(self.eos_token_id)
if any(t not in (self.eos_token_id, self.pad_token_id) for t in row[end:]):
raise ValueError("only EOS/padding may follow completed action generation")
if not self.grammar.is_complete([t - self.token_offset for t in row[:end]]):
raise ValueError("EOS occurred before the action chunk was complete")
return "done"
state = self.grammar._consume([t - self.token_offset for t in row])
if state is None:
raise ValueError(
"generated prefix violates the action grammar; check prompt_length and token_offset"
)
return state
def _mask(self, state, scores):
key = (state, str(scores.device), scores.shape[-1])
mask = self._mask_cache.get(key)
if mask is None:
mask = torch.zeros(scores.shape[-1], device=scores.device, dtype=torch.bool)
if state == "done":
mask[self.eos_token_id] = True
if self.pad_token_id is not None:
mask[self.pad_token_id] = True
else:
allowed = self.grammar._mask_from_state(state)
stop = self.token_offset + len(allowed)
mask[self.token_offset : stop] = torch.as_tensor(allowed, device=scores.device)
mask[self.eos_token_id] = self.grammar._state_complete(state)
self._mask_cache[key] = mask
return mask
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("expected input_ids(B,L) and scores(B,V) with the same batch size")
if input_ids.shape[1] < self.prompt_length:
raise ValueError("prompt_length exceeds the supplied input length")
required = max(
self.token_offset + self.grammar.vocab_size,
self.eos_token_id + 1,
0 if self.pad_token_id is None else self.pad_token_id + 1,
)
if scores.shape[1] < required:
raise ValueError("model vocabulary does not cover configured action/EOS/padding IDs")
self.grammar.reserve_states(int(input_ids.shape[0]))
rows = input_ids[:, self.prompt_length :].detach().cpu().tolist()
output = scores.clone()
groups = {}
for index, row in enumerate(rows):
groups.setdefault(self._state(row), []).append(index)
for state, indices in groups.items():
output[indices] = scores[indices].masked_fill(~self._mask(state, scores), -torch.inf)
if not torch.isfinite(output).any(dim=1).all():
raise ValueError(
"all legal action logits were suppressed; check other generation processors"
)
return output
|