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