File size: 1,990 Bytes
b296ad4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | """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()
|