File size: 1,832 Bytes
795f737 | 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 41 42 43 44 45 | import unittest
try:
import torch
from baim.features import encode,fit_vocab,candidates
from baim.model import PointerPolicy
except ImportError:
torch = None
from baim.synthetic import generate
@unittest.skipIf(torch is None,'training extras not installed')
class ModelTests(unittest.TestCase):
def test_candidate_filter_rejects_hidden_disabled_and_sensitive(self):
sample = next(generate('train',1,7))
for index,key in enumerate(['visible','enabled','sensitive']):
sample['elements'][index][key] = key == 'sensitive'
result = candidates(sample['goal'],sample['elements'])
self.assertFalse(set(result)&{0,1,2})
def test_pointer_is_equivariant_to_candidate_order(self):
torch.manual_seed(1)
torch.set_num_threads(2)
rows = list(generate('train',2,11))
vocab = fit_vocab(rows)
x,*_ = encode(rows,vocab)
model = PointerPolicy(vocab_size=len(vocab)).eval()
with torch.inference_mode():
a,t = model(*x)
permutation = torch.arange(x[1].shape[1]-1,-1,-1)
changed = (x[0],x[1][:,permutation],x[2][:,permutation],x[3][:,permutation])
other_a,other_t = model(*changed)
torch.testing.assert_close(a,other_a)
torch.testing.assert_close(t[:,permutation],other_t)
def test_all_encoders_handle_padded_candidates(self):
rows = list(generate('train',3,19))
vocab = fit_vocab(rows)
x,*_ = encode(rows,vocab)
for architecture in ['mean','gru','transformer']:
model = PointerPolicy(vocab_size=len(vocab),encoder=architecture).eval()
with torch.inference_mode():
a,t = model(*x)
self.assertTrue(torch.isfinite(a).all())
self.assertTrue(torch.isfinite(t).all())
|