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