| 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()) |
|
|
