| import unittest, torch | |
| from ares.config import AresConfig | |
| from ares.model import AresTransformer | |
| from ares.tokenizer import ByteBPETokenizer | |
| class TestAres(unittest.TestCase): | |
| def test_tokenizer_roundtrip(self): | |
| t=ByteBPETokenizer();s='<user>Hello, Ares! 🌍<assistant>';self.assertEqual(t.decode(t.encode(s)),s) | |
| def test_transformer_shapes_and_cache(self): | |
| c=AresConfig(vocab_size=64,max_seq_len=16,dim=32,n_layers=2,n_heads=4,n_kv_heads=2);m=AresTransformer(c);x=torch.randint(0,64,(2,8));logits,loss,cache=m(x,x);self.assertEqual(logits.shape,(2,8,64));self.assertTrue(torch.isfinite(loss));self.assertEqual(len(cache),2) | |
| if __name__=='__main__':unittest.main() | |