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,stage_identity=None):
9    if saved['identity_sha256']!=digest(saved['identity']):raise ValueError('checkpoint identity')
10    if saved['format']!='random-state-8h-v1' or stage_identity is None or saved['identity']!=stage_identity:raise ValueError('CM-8 identity')
11    if stage_identity.get('candidate_name')!='CM-8' or stage_identity.get('initialization')!='random' or stage_identity.get('parent') is not None:raise ValueError('CM-8 random lineage')
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        from checkpoint_gate import permitted,validate_milestone
47        stage,selected=permitted(contract)
48        if selected!={'path':proposal['checkpoint'],'sha256':proposal['checkpoint_sha256']}:raise ValueError('uncertified evaluation selection')
49        if str(Path(path).resolve())!=proposal['checkpoint'] or sha(path)!=proposal['checkpoint_sha256']:raise ValueError('wrong selected checkpoint')
50        saved=torch.load(path,map_location='cpu',weights_only=True)
51        validate_milestone(saved,contract)
52        validate_identity(saved,stage)
53        cfg=saved['model_config']
54        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')
55        torch.set_num_threads(2);torch.manual_seed(seed)
56        self.device=os.environ.get('CHESS_INFERENCE_DEVICE','cuda:0')
57        self.model=StateModel(cfg);self.model.load_state_dict(saved['model_state_dict'],strict=True)
58        self.model.to(self.device).eval();del saved
59    def choose(self,board):
60        import torch
61        def logits(xs,masks):
62            with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16,enabled=str(self.device).startswith('cuda')):
63                return self.model.next_logits(torch.tensor(xs,device=self.device),torch.tensor(masks,device=self.device)).float().cpu().tolist()
64        return decide(board,logits)
65
66def load(path,seed=0):return Adapter(path,seed)
67