File size: 4,674 Bytes
5f70891 | 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 | """Batch-one NumPy beam state for the pinned TrOCR generation policy.
Implements Transformers4.57's finished/running beam rules. Equal-score top-k
ties use lower indices, explicitly; Torch CPU/CUDA tie order is not portable.
Source equivalence and complete OCR qualification must be measured separately.
"""
import numpy as np
def top_indices(scores, count):
scores = np.asarray(scores)
assert scores.ndim == 1 and 1 <= count <= scores.size
assert not np.isnan(scores).any() and not np.isposinf(scores).any() and np.isfinite(scores).sum() >= count
threshold = np.partition(scores, scores.size - count)[scores.size - count]
better = np.flatnonzero(scores > threshold)
tied = np.flatnonzero(scores == threshold)[:count - len(better)]
indices = np.concatenate([better, tied])
return indices[np.lexsort((indices, -scores[indices]))]
def log_softmax(scores):
scores = np.asarray(scores, dtype=np.float32)
shifted = scores - scores.max(axis=-1, keepdims=True)
return shifted - np.log(np.exp(shifted).sum(axis=-1, keepdims=True))
class TrOCRBeamState:
def __init__(self, beams=10, max_new_tokens=256, length_penalty=1., suppress_tokens=()):
if beams < 2 or not 1 <= max_new_tokens <= 512 or length_penalty < 0:
raise ValueError('Unsupported beam search settings')
self.beams = beams;self.maximum = max_new_tokens + 1;self.penalty = length_penalty
self.suppressed = tuple(suppress_tokens)
if any(type(t) is not int or t < 0 for t in self.suppressed) or len(set(self.suppressed)) != len(self.suppressed):
raise ValueError('Suppressed token IDs must be unique nonnegative integers')
self.current = 1
self.running = np.full((beams, self.maximum), 1, np.int64);self.running[:, 0] = 2
self.running_scores = np.full(beams, -1e9, np.float32);self.running_scores[0] = 0
self.finished = self.running.copy();self.finished_scores = np.full(beams, -1e9, np.float32)
self.finished_flags = np.zeros(beams, bool);self.finished_lengths = np.zeros(beams, np.int64)
self.done = False
def advance(self, logits):
logits = np.asarray(logits, dtype=np.float32)
if self.done or logits.ndim != 2 or logits.shape[0] != self.beams or logits.shape[1] < 3 or not np.isfinite(logits).all():
raise ValueError('Invalid scores or completed beam search')
classes = logits.shape[1]
if self.suppressed and max(self.suppressed) >= classes:raise ValueError('Suppressed ID outside vocabulary')
logprobs = log_softmax(logits)
# HF applies suppression after log-softmax, without renormalization.
logprobs[:, self.suppressed] = -np.inf
accumulated = logprobs + self.running_scores[:, None]
choices = top_indices(accumulated.reshape(-1), 2 * self.beams)
scores = accumulated.reshape(-1)[choices]
parents, tokens = choices // classes, choices % classes
candidates = self.running[parents].copy();candidates[:, self.current] = tokens
stopped = (tokens == 2) | (self.current + 1 == self.maximum)
live_scores = scores + stopped.astype(np.float32) * np.float32(-1e9)
live = top_indices(live_scores, self.beams)
self.running = candidates[live].copy();self.running_scores = live_scores[live].copy()
eligible = stopped & (np.arange(2 * self.beams) < self.beams)
normalized = scores / np.float32(self.current ** self.penalty)
normalized = normalized + (~eligible).astype(np.float32) * np.float32(-1e9)
merged_scores = np.concatenate([self.finished_scores, normalized])
selected = top_indices(merged_scores, self.beams)
self.finished = np.concatenate([self.finished, candidates])[selected].copy()
self.finished_scores = merged_scores[selected].copy()
self.finished_flags = np.concatenate([self.finished_flags, eligible])[selected].copy()
self.finished_lengths = np.concatenate([self.finished_lengths, np.full(2 * self.beams, self.current, np.int64)])[selected].copy()
self.current += 1
possible = self.running_scores[0] / np.float32((self.current - 1) ** self.penalty)
worst = self.finished_scores.min() if self.finished_flags.all() else np.float32(-1e9)
self.done = bool(possible <= worst or stopped.all())
return parents[live].copy(), tokens[live].astype(np.int64, copy=True)
def result(self):
if not self.done or not self.finished_flags[0]:raise ValueError('Beam search has no completed result')
return self.finished[0, :1 + self.finished_lengths[0]].tolist(), float(self.finished_scores[0])
|