| """Regression checks for causal decoding and the learned source-copy distribution.""" |
| import unittest |
| import torch |
| from tinyquery.model import Config,TinyQuery |
|
|
|
|
| class ModelChecks(unittest.TestCase): |
| def setUp(self): |
| torch.set_num_threads(2); torch.manual_seed(51) |
| self.model=TinyQuery(Config(vocab_size=64,width=64,layers=2,heads=4,kv_heads=2,hidden=128,context=64,copy_dim=32)).eval() |
|
|
| def test_causality_and_cache(self): |
| x=torch.randint(5,64,(2,15));x[:,7]=3 |
| original=self.model(x)[0]; changed=x.clone();changed[:,9:]=torch.randint(5,64,(2,6)) |
| self.assertTrue(torch.allclose(original[:,:9],self.model(changed)[0][:,:9],atol=1e-6)) |
| past=None; parts=[] |
| for i in range(x.shape[1]): |
| logits,past,_=self.model(x[:,i:i+1],past=past,use_cache=True,last_only=True);parts.append(logits) |
| self.assertLess(float((original-torch.cat(parts,dim=1)).abs().max().detach()),2e-5) |
| self.assertTrue(torch.allclose(original.exp().sum(-1),torch.ones_like(original[:,:,0]),atol=1e-6)) |
|
|
| def test_copy_cannot_recycle_its_response(self): |
| with torch.no_grad(): |
| self.model.copy_gate.weight.zero_();self.model.copy_gate.bias.fill_(-30) |
| self.model.copy_query.weight.zero_();self.model.copy_key.weight.zero_() |
| probabilities=self.model(torch.tensor([[1,5,8,3,20,20]]),last_only=True)[0].exp()[0,0] |
| self.assertLess(float(probabilities[20]),1e-8) |
| for token in [1,5,8]:self.assertAlmostEqual(float(probabilities[token]),1/3,places=6) |
|
|
| def test_copy_gradient(self): |
| x=torch.randint(5,64,(2,15));x[:,7]=3 |
| y=torch.roll(x,shifts=-1,dims=1);mask=torch.arange(15)[None,:].expand(2,-1)>=7 |
| loss,_=self.model(x,y,mask,torch.tensor([7,7]),torch.tensor([0,0])) |
| loss.backward() |
| self.assertTrue(torch.isfinite(loss)) |
| self.assertGreater(float(self.model.copy_query.weight.grad.norm()),0) |
|
|
|
|
| if __name__=='__main__':unittest.main() |
|
|
