scorer.py

1"""Exact legal joint scoring with per-decision full-history KV reuse.
2
3No cross-game mutable cache. All 128 logits remain unmasked for diagnostics.
4"""
5import math, time
6import torch
7from transformers import DynamicCache
8import harness as h
9from chessdata import tokenizer
10
11def normalize(values):
12    values=list(map(float,values))
13    if len(values)!=128 or not all(math.isfinite(x) for x in values):raise ValueError('invalid raw logits')
14    m=max(values);v=[x-m for x in values];z=math.log(sum(math.exp(x) for x in v))
15    return [x-z for x in v]
16
17class Scorer:
18    context_length=1024
19    def __init__(self,net,microbatch=8):
20        self.net=net;self.device=next(net.parameters()).device;self.microbatch=microbatch
21    @torch.inference_mode()
22    def next_logits(self,prefix):
23        if not 1<=len(prefix)<=1024:raise ValueError('context overflow')
24        x=torch.tensor([prefix],device=self.device)
25        return self.net(x,attention_mask=torch.ones_like(x,dtype=torch.bool),use_cache=False).logits[0,-1].float().tolist()
26    @torch.inference_mode()
27    def score(self,board):
28        started=time.monotonic();hist=h.history(board)
29        if len(hist)+2>1024:raise ValueError('full history exceeds context; no crop')
30        legal=sorted(board.legal_moves,key=lambda m:m.uci())
31        if not legal:raise ValueError('terminal board')
32        triples=[tokenizer.encode_move(m) for m in legal]
33        x=torch.tensor([hist],device=self.device)
34        initial=self.net(x,attention_mask=torch.ones_like(x,dtype=torch.bool),use_cache=True)
35        base=initial.past_key_values
36        lp={():normalize(initial.logits[0,-1].float().tolist())}
37        calls=1
38        def tails_logits(tails):
39            nonlocal calls
40            tails=sorted(set(tails)-lp.keys())
41            for offset in range(0,len(tails),self.microbatch):
42                part=tails[offset:offset+self.microbatch];n=len(part)
43                # Fresh branch tensors: forward cannot mutate another game's/history's cache.
44                cache=DynamicCache([(k.repeat(n,1,1,1),v.repeat(n,1,1,1)) for k,v,*_ in base])
45                tail=torch.tensor(part,device=self.device)
46                mask=torch.ones((n,len(hist)+len(part[0])),dtype=torch.bool,device=self.device)
47                out=self.net(tail,past_key_values=cache,attention_mask=mask,use_cache=True).logits[:,-1].float().tolist()
48                for key,row in zip(part,out):lp[key]=normalize(row)
49                calls+=1
50        raw=[max(range(128),key=lambda i:lp[()][i])]
51        tails_logits([(a,) for a,b,c in triples]+[tuple(raw)])
52        raw.append(max(range(128),key=lambda i:lp[tuple(raw)][i]))
53        tails_logits([(a,b) for a,b,c in triples]+[tuple(raw)])
54        raw.append(max(range(128),key=lambda i:lp[tuple(raw)][i]))
55        scores=[lp[()][a]+lp[(a,)][b]+lp[(a,b)][c] for a,b,c in triples]
56        best=max(range(len(legal)),key=lambda i:scores[i]);mass=sum(math.exp(s) for s in scores)
57        return legal[best],dict(legal_joint_probability=mass,illegal_or_nonmove_three_token_probability=max(0.,1-mass),
58            raw_greedy_tokens=raw,raw_greedy_legal=raw in triples,selected_joint_log_probability=scores[best],
59            forward_calls=calls,scoring_seconds=time.monotonic()-started,full_legal_scores=scores)
60