| """Focused invariant tests; requires torch and transformers but no downloads.""" |
| import types |
| import unittest |
| import torch |
| from torch import nn |
| from v4_model import JudgmentV4, normalize_row |
|
|
| class Tokenizer: |
| pad_token_id=0 |
| def encode(self,text,add_special_tokens=False):return [int(text)+1] if int(text)<10 else [1,2] |
| def apply_chat_template(self,messages,**kwargs): |
| assert kwargs.get('return_dict') is False |
| assert kwargs.get('enable_thinking') is False |
| |
| return [1,2,3] if 'short' in messages[1]['content'] else [1,2,3,4,5] |
|
|
| class Backbone(nn.Module): |
| def __init__(self): |
| super().__init__();self.emb=nn.Embedding(300,8) |
| def forward(self,input_ids,**kwargs):return types.SimpleNamespace(last_hidden_state=self.emb(input_ids)) |
| class LM(nn.Module): |
| def __init__(self): |
| super().__init__();self.model=Backbone();self.head=nn.Linear(8,300,bias=False) |
| self.config=types.SimpleNamespace(max_position_embeddings=65536,use_cache=True) |
| def get_output_embeddings(self):return self.head |
| def forward(self,*a,**k):raise AssertionError('Full sequence vocabulary forward must never run') |
|
|
| class Tests(unittest.TestCase): |
| def setUp(self):self.model=JudgmentV4(device='cpu',tokenizer=Tokenizer(),lm=LM()) |
| def row(self,state='short',n=2):return dict(state=state,kind='choice',instructions='choose',candidates=[str(i) for i in range(n)],keys=[str(i) for i in range(n)],target=[1.]+[0.]*(n-1)) |
| def test_noul(self): |
| r=self.row();r.update(kind='noul',target=[.8],candidates=['']) |
| out=normalize_row(r);self.assertEqual(out['keys'],['false','true']);self.assertAlmostEqual(out['target'][0],.2) |
| def test_order_preserved(self): |
| r=self.row(n=3);r['kind']='score';self.assertEqual(normalize_row(r)['candidates'],['0','1','2']) |
| def test_no_silent_truncation(self): |
| with self.assertRaises(ValueError):self.model.encode_rows([self.row('long')],max_length=4) |
| self.assertEqual(self.model.encode_rows([self.row('long')],max_length=4,strict=False),[]) |
| def test_forward_masks_gradients_padding(self): |
| e=self.model.encode_rows([self.row(),self.row('long',3)]) |
| b=self.model.collate(e) |
| self.assertEqual(b['attention_mask'][0].tolist(),[0,0,1,1,1]) |
| out=self.model(**b);self.assertTrue(torch.isfinite(out['loss'])) |
| self.assertEqual(out['logits'].softmax(-1)[0,2].item(),0) |
| single=self.model(**self.model.collate(e[:1]))['logits'][0] |
| torch.testing.assert_close(single,out['logits'][0,:2]) |
| out['loss'].backward();self.assertIsNotNone(self.model.lm.model.emb.weight.grad) |
| def test_temperature_inference_only(self): |
| row=self.row() |
| self.model.eval() |
| encoded=self.model.encode_rows([row]);batch=self.model.collate(encoded) |
| raw=self.model(**batch)['logits'].detach().clone() |
| self.model.calibration={'choice':{'temperature':2.}} |
| actual=self.model.predict_records([row])[0] |
| expected=(raw/2).softmax(-1)[0].tolist() |
| self.assertEqual(actual,expected) |
| torch.testing.assert_close(raw,self.model(**batch)['logits']) |
|
|
| def test_multiple_rows_preserve_order_and_calibration(self): |
| rows=[self.row('long',3),self.row('short',2),self.row('long',4)] |
| rows[1]['kind']='score' |
| self.model.calibration={'choice':{'temperature':2.},'score':{'temperature':.5}} |
| single=[self.model.predict_records([r])[0] for r in rows] |
| batch=self.model.predict_records(rows) |
| self.assertEqual([len(x) for x in batch],[3,2,4]) |
| for a,b in zip(single,batch):torch.testing.assert_close(torch.tensor(a),torch.tensor(b)) |
| self.assertEqual(self.model.last_usage['encoded_questions'],3) |
|
|
| def test_multitoken_indices_rejected(self): |
| self.assertEqual(self.model.max_candidates,10) |
| self.model.encode_rows([self.row(n=10)]) |
| with self.assertRaisesRegex(ValueError,'at most 10'): |
| self.model.encode_rows([self.row(n=11)]) |
|
|
| def test_native_head_final_position_only(self): |
| calls=[] |
| hook=self.model.lm.head.register_forward_hook(lambda module,args,output:calls.append(args[0].shape)) |
| encoded=self.model.encode_rows([self.row(),self.row('long',3)]) |
| out=self.model(**self.model.collate(encoded)) |
| self.assertEqual(calls,[torch.Size([2,8])]) |
| hook.remove() |
| self.assertEqual(out['logits'].shape,(2,3)) |
|
|
| def test_multimodal_decoder_bypassed(self): |
| decoder=self.model.lm.model |
| class Multimodal(nn.Module): |
| def __init__(self):super().__init__();self.language_model=decoder |
| def forward(self,*args,**kwargs):raise AssertionError('Must use text decoder') |
| self.model.lm.model=Multimodal() |
| self.model(**self.model.collate(self.model.encode_rows([self.row()]))) |
|
|
| def test_lora_includes_hybrid_excludes_vision(self): |
| from v4_model import lora_targets |
| lm=nn.Module();lm.model=nn.Module() |
| lm.model.language_model=nn.Module();lm.model.visual=nn.Module() |
| for name in ('in_proj_qkv','in_proj_z','in_proj_a','in_proj_b','out_proj','q_proj','gate_proj'): |
| setattr(lm.model.language_model,name,nn.Linear(2,2)) |
| setattr(lm.model.visual,name,nn.Linear(2,2)) |
| targets=lora_targets(lm) |
| self.assertEqual(len(targets),7) |
| self.assertTrue(all(x.startswith('model.language_model.') for x in targets)) |
|
|
| def test_invalid_target(self): |
| r=self.row();r['target']=[.5,.9] |
| with self.assertRaises(ValueError):normalize_row(r) |
|
|
| if __name__=='__main__':unittest.main() |
|
|
