File size: 4,105 Bytes
64b9802 | 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 | """Fixed-length, constrained HF beam search and independent numerical checks."""
import torch
from transformers import GenerationConfig, LogitsProcessor, LogitsProcessorList
class VisualTokensOnly(LogitsProcessor):
def __init__(self, valid_size):
self.valid_size = valid_size
def __call__(self, input_ids, scores):
scores = scores.clone()
scores[:, self.valid_size:] = -torch.inf
return scores
@torch.inference_mode()
def beam_decode(model, prefix, width=4, steps=80, valid_size=4375):
config = GenerationConfig(
do_sample=False, num_beams=width, num_return_sequences=width,
max_new_tokens=steps, min_new_tokens=steps,
eos_token_id=None, bos_token_id=None, pad_token_id=0,
length_penalty=0.0 if width > 1 else 1.0, early_stopping=False, renormalize_logits=True,
repetition_penalty=1.0, no_repeat_ngram_size=0,
use_cache=True, return_dict_in_generate=True, output_scores=True,
)
result = model.generate(
input_ids=prefix, attention_mask=torch.ones_like(prefix),
generation_config=config,
logits_processor=LogitsProcessorList([VisualTokensOnly(valid_size)]),
eos_token_id=None, bos_token_id=None,
forced_bos_token_id=None, forced_eos_token_id=None,
)
candidates = result.sequences[:, prefix.shape[1]:]
assert candidates.shape == (width,steps)
assert int(candidates.min()) >= 0 and int(candidates.max()) < valid_size
assert len({tuple(x) for x in candidates.tolist()}) == width
transitions = model.compute_transition_scores(
result.sequences, result.scores, getattr(result,'beam_indices',None),
normalize_logits=False,
)
assert transitions.shape == (width,steps) and torch.isfinite(transitions).all()
sums = transitions.sum(-1)
scores = result.sequences_scores if width > 1 else sums
error = float((sums-scores).abs().max())
assert torch.allclose(sums,scores,rtol=0,atol=2e-4), error
assert torch.all(scores[:-1] >= scores[1:])
return candidates, scores, {'beam_width':width,'tokens_per_frame':steps,
'selected_index':0,'candidate_scores':scores.tolist(),
'score_reconstruction_max_error':error,
'score_definition':'Sum of log probabilities renormalized over valid visual tokens only; no length penalty or GT/image reward.'}
@torch.inference_mode()
def independent_check():
from transformers import LlamaConfig, LlamaForCausalLM
torch.manual_seed(3719)
model = LlamaForCausalLM(LlamaConfig(vocab_size=8,hidden_size=32,
intermediate_size=48,num_hidden_layers=2,num_attention_heads=4,
num_key_value_heads=4,max_position_embeddings=128,
eos_token_id=2,bos_token_id=1,pad_token_id=0,
attention_dropout=0.0)).float().eval()
prefix = torch.tensor([[6,1,3]],dtype=torch.long)
reports=[]
for width in [1,4]:
beams=[([],0.0)]
for _ in range(4):
expanded=[]
for seq,score in beams:
x=torch.tensor([prefix[0].tolist()+seq],dtype=torch.long)
probs=model(input_ids=x,use_cache=False).logits[0,-1,:5].log_softmax(-1)
expanded.extend((seq+[token],score+float(probs[token])) for token in range(5))
beams=sorted(expanded,key=lambda item:item[1],reverse=True)[:width]
ids,scores,meta=beam_decode(model,prefix,width=width,steps=4,valid_size=5)
expected=torch.tensor([seq for seq,_ in beams],dtype=torch.long)
expected_scores=torch.tensor([score for _,score in beams])
assert torch.equal(ids,expected),(width,ids,expected)
assert torch.allclose(scores,expected_scores,rtol=0,atol=2e-5)
reports.append({'width':width,'matches_independent_uncached_search':True,
'max_score_error':float((scores-expected_scores).abs().max())})
return {'tests':reports,'eos_disabled':True,'valid_token_constraint_checked':True,'complete':True}
if __name__ == '__main__':
import json, transformers
print(json.dumps({'transformers':transformers.__version__,**independent_check()}))
|