test_model.py

python · 31 lines · Exact bytes

Version & provenance · state-v2-cm7

Recorded SHA-256 61831d5d530552ec8e96537ef09dde596aaffd89afc5b68a480b3fe9d9953d51Rendered SHA-256 61831d5d530552ec8e96537ef09dde596aaffd89afc5b68a480b3fe9d9953d51

Apache-2.0 original project code; dependencies retain their own licenses. License notes. Line links are stable within this snapshot. CM-7 completed implementation source; original and public rendered hashes recorded. Host-specific paths use /work placeholders; inspect reproduction instructions before running. Separate Rust helper retains its GPL license.

1import unittest
2import torch
3from model import StateModel,loss_sum
4import codec as c
5
6class Tests(unittest.TestCase):
7    def test_shift_inference_and_mask(self):
8        torch.set_num_threads(2);torch.manual_seed(4)
9        net=StateModel(dict(vocab_size=512,n_positions=1024,n_ctx=1024,n_embd=24,n_layer=2,n_head=2,resid_pdrop=0.,embd_pdrop=0.,attn_pdrop=0.,use_cache=False)).eval()
10        e=c.example(c.chess.Board(),c.chess.Move.from_uci('e2e4'),0)
11        x=torch.tensor([e['input_ids']]);mask=torch.tensor([e['attention_mask']],dtype=torch.bool);y=torch.tensor([e['labels']])
12        all_logits=net(x,mask)
13        for phase in range(3):
14            one=net.next_logits(x[:,:120+phase],mask[:,:120+phase])
15            torch.testing.assert_close(one,all_logits[:,phase],atol=1e-6,rtol=1e-5)
16        changed=x.clone();changed[:,2:50]=511
17        torch.testing.assert_close(net(changed,mask),all_logits,atol=1e-6,rtol=1e-5)
18        loss=loss_sum(all_logits,y);loss.backward();self.assertIsNotNone(net.net.transformer.h[0].attn.c_attn.weight.grad)
19        with self.assertRaises(ValueError):net(x[:,:121],mask[:,:121])
20        with self.assertRaises(ValueError):loss_sum(all_logits,y[:,:2])
21    def test_accumulation_normalization(self):
22        torch.set_num_threads(2);torch.manual_seed(1)
23        cfg=dict(vocab_size=512,n_positions=1024,n_ctx=1024,n_embd=16,n_layer=1,n_head=2,resid_pdrop=0.,embd_pdrop=0.,attn_pdrop=0.)
24        a=StateModel(cfg);b=StateModel(cfg);b.load_state_dict(a.state_dict())
25        x=torch.randint(0,512,(4,122));mask=torch.ones_like(x,dtype=torch.bool);y=torch.tensor([[12,28,64]]*4)
26        (loss_sum(a(x,mask),y)/12).backward()
27        for i in range(4):(loss_sum(b(x[i:i+1],mask[i:i+1]),y[i:i+1])/12).backward()
28        for p,q in zip(a.parameters(),b.parameters()):torch.testing.assert_close(p.grad,q.grad,atol=2e-6,rtol=2e-5)
29
30if __name__=='__main__':unittest.main()
31