adapter.py

1"""State-v2 exact legal joint argmax, with unmasked 512-class diagnostics."""
2import sys,os,math,json
3from pathlib import Path
4STATE=Path(__file__).resolve().parent.parent/'state-v2'
5sys.path.insert(0,str(STATE))
6import codec as c
7from records import sha,digest
8def validate_identity(saved,ref):
9    if saved['identity_sha256']!=digest(saved['identity']) or saved['identity_sha256']!=ref['identity_sha256']:raise ValueError('checkpoint identity')
10    cursor=saved.get('updates',saved.get('cursor'))
11    if saved['format']!=ref['format'] or cursor!=ref['cursor']:raise ValueError('checkpoint lineage/cursor')
12
13def decide(board,logits,history_plies=16):
14    """logits takes equal-length input/mask batches; returns raw finite rows."""
15    value=c.prompt(board,history_plies)
16    moves=sorted(board.legal_moves,key=lambda m:m.uci())
17    if not moves or board.is_game_over(claim_draw=False):raise ValueError('terminal board')
18    triples=[c.triple(m) for m in moves]
19    cache={};raw=[];calls=0
20    for phase in range(3):
21        tails=sorted(set(tuple(t[:phase]) for t in triples)|{tuple(raw)})
22        xs=[value['input_ids']+list(t) for t in tails]
23        masks=[value['attention_mask']+[1]*phase for _ in tails]
24        rows=logits(xs,masks);calls+=1
25        if len(rows)!=len(tails):raise ValueError('logit batch size')
26        for tail,row in zip(tails,rows):
27            if len(row)!=c.VOCAB or not all(math.isfinite(float(x)) for x in row):raise ValueError('512 finite logits required')
28            top=max(row);z=math.log(sum(math.exp(float(x)-top) for x in row))
29            cache[tail]=[float(x)-top-z for x in row]
30        row=cache[tuple(raw)];raw.append(max(range(c.VOCAB),key=lambda t:row[t]))
31    scores=[sum(cache[tuple(t[:p])][t[p]] for p in range(3)) for t in triples]
32    best=max(range(len(moves)),key=lambda i:scores[i]);mass=sum(math.exp(s) for s in scores)
33    return moves[best],dict(raw_greedy_tokens=raw,raw_greedy_legal=raw in triples,
34        legal_joint_probability=mass,illegal_or_nonmove_three_token_probability=max(0.,1.-mass),
35        selected_joint_log_probability=scores[best],forward_calls=calls,
36        input_prompt_tokens=120,history_plies=min(history_plies,len(board.move_stack)),vocab_size=c.VOCAB)
37
38class Adapter:
39    context_length=122
40    def __init__(self,path,seed):
41        import torch
42        from model import StateModel
43        proposal=json.loads(Path(__file__).with_name('proposal.json').read_text())
44        contract=json.loads(Path(__file__).with_name('execution-contract.json').read_text())
45        if sha(Path(__file__).with_name('execution-contract.json'))!=proposal['execution_contract_sha256']:raise ValueError('execution contract changed')
46        refs=contract['players']
47        found=[r for r in refs.values() if r['path']==str(Path(path).resolve())]
48        if len(found)!=1:raise ValueError('unapproved checkpoint path')
49        ref=found[0]
50        if sha(path)!=ref['sha256']:raise ValueError('checkpoint changed')
51        saved=torch.load(path,map_location='cpu',weights_only=True)
52        validate_identity(saved,ref)
53        cfg=saved['model_config']
54        self.legacy=ref['representation']=='full-history-v1'
55        vocab=128 if self.legacy else 512
56        if any(cfg[k]!=v for k,v in dict(n_positions=1024,n_layer=12,n_head=12,n_embd=768,vocab_size=vocab).items()):raise ValueError('architecture')
57        torch.set_num_threads(2);torch.manual_seed(seed)
58        self.device=os.environ.get('CHESS_INFERENCE_DEVICE','cuda:0')
59        if self.legacy:
60            from transformers import GPT2Config,GPT2LMHeadModel
61            from chessdata import tokenizer
62            if saved['identity']['tokenizer']['sha256']!=sha(Path(tokenizer.__file__)):raise ValueError('legacy tokenizer changed')
63            config=GPT2Config(**cfg);config._attn_implementation='sdpa';self.model=GPT2LMHeadModel(config);self.context_length=1024
64        else:self.model=StateModel(cfg);self.context_length=122
65        self.model.load_state_dict(saved['model_state_dict'],strict=True)
66        self.model.to(self.device).eval();del saved
67    def choose(self,board):
68        import torch
69        def logits(xs,masks):
70            with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16,enabled=str(self.device).startswith('cuda')):
71                x=torch.tensor(xs,device=self.device);attention=torch.tensor(masks,device=self.device)
72                y=self.model(x,attention_mask=attention,use_cache=False).logits[:,-1] if self.legacy else self.model.next_logits(x,attention)
73                return y.float().cpu().tolist()
74        if self.legacy:
75            from legacy_decode import decide as legacy_decide
76            return legacy_decide(board,logits)
77        return decide(board,logits)
78
79def load(path,seed=0):return Adapter(path,seed)
80