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    if saved['format']!=ref['format'] or saved['updates']!=ref['updates']:raise ValueError('checkpoint lineage/cursor')
11
12def decide(board,logits,history_plies=16):
13    """logits takes equal-length input/mask batches; returns raw finite rows."""
14    value=c.prompt(board,history_plies)
15    moves=sorted(board.legal_moves,key=lambda m:m.uci())
16    if not moves or board.is_game_over(claim_draw=False):raise ValueError('terminal board')
17    triples=[c.triple(m) for m in moves]
18    cache={};raw=[];calls=0
19    for phase in range(3):
20        tails=sorted(set(tuple(t[:phase]) for t in triples)|{tuple(raw)})
21        xs=[value['input_ids']+list(t) for t in tails]
22        masks=[value['attention_mask']+[1]*phase for _ in tails]
23        rows=logits(xs,masks);calls+=1
24        if len(rows)!=len(tails):raise ValueError('logit batch size')
25        for tail,row in zip(tails,rows):
26            if len(row)!=c.VOCAB or not all(math.isfinite(float(x)) for x in row):raise ValueError('512 finite logits required')
27            top=max(row);z=math.log(sum(math.exp(float(x)-top) for x in row))
28            cache[tail]=[float(x)-top-z for x in row]
29        row=cache[tuple(raw)];raw.append(max(range(c.VOCAB),key=lambda t:row[t]))
30    scores=[sum(cache[tuple(t[:p])][t[p]] for p in range(3)) for t in triples]
31    best=max(range(len(moves)),key=lambda i:scores[i]);mass=sum(math.exp(s) for s in scores)
32    return moves[best],dict(raw_greedy_tokens=raw,raw_greedy_legal=raw in triples,
33        legal_joint_probability=mass,illegal_or_nonmove_three_token_probability=max(0.,1.-mass),
34        selected_joint_log_probability=scores[best],forward_calls=calls,
35        input_prompt_tokens=120,history_plies=min(history_plies,len(board.move_stack)),vocab_size=c.VOCAB)
36
37class Adapter:
38    context_length=122
39    def __init__(self,path,seed):
40        import torch
41        from model import StateModel
42        proposal=json.loads(Path(__file__).with_name('proposal.json').read_text())
43        contract=json.loads(Path(__file__).with_name('execution-contract.json').read_text())
44        if sha(Path(__file__).with_name('execution-contract.json'))!=proposal['execution_contract_sha256']:raise ValueError('execution contract changed')
45        refs=contract['players']
46        found=[r for r in refs.values() if r['path']==str(Path(path).resolve())]
47        if len(found)!=1:raise ValueError('unapproved checkpoint path')
48        ref=found[0]
49        if sha(path)!=ref['sha256']:raise ValueError('checkpoint changed')
50        saved=torch.load(path,map_location='cpu',weights_only=True)
51        validate_identity(saved,ref)
52        cfg=saved['model_config']
53        if any(cfg[k]!=v for k,v in dict(n_positions=1024,n_layer=12,n_head=12,n_embd=768,vocab_size=512).items()):raise ValueError('architecture')
54        torch.set_num_threads(2);torch.manual_seed(seed)
55        self.device=os.environ.get('CHESS_INFERENCE_DEVICE','cuda:0')
56        self.model=StateModel(cfg);self.model.load_state_dict(saved['model_state_dict'],strict=True)
57        self.model.to(self.device).eval();del saved
58    def choose(self,board):
59        import torch
60        def logits(xs,masks):
61            with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16,enabled=str(self.device).startswith('cuda')):
62                return self.model.next_logits(torch.tensor(xs,device=self.device),torch.tensor(masks,device=self.device)).float().cpu().tolist()
63        return decide(board,logits)
64
65def load(path,seed=0):return Adapter(path,seed)
66