| import unittest |
| from types import SimpleNamespace |
| import mlx.core as mx |
| from generate import generate |
|
|
|
|
| class Runner(unittest.TestCase): |
| def model(self): |
| class Tokenizer: |
| def encode(self,*args,**kwargs):return SimpleNamespace(ids=[7,8]) |
| def decode(self,ids,**kwargs):return ' '.join(str(i) for i in ids if i) |
| class Model: |
| max_tokens=128 |
| w=SimpleNamespace(config={'eos_token_id':0}) |
| hash=SimpleNamespace(history=[]) |
| tokenizer=Tokenizer() |
| def step(self,token): |
| self.position+=1 |
| expected=20 if self.position<=2 else 21 if self.position==3 else 0 |
| return mx.array([10. if i==expected else -10. for i in range(32)]),mx.array([0.]) |
| return Model() |
|
|
| def test_greedy_output_counts_and_context_reset(self): |
| model=self.model() |
| for _ in range(2): |
| output,report=generate(model,'Hello',16) |
| self.assertEqual(output,'20 21') |
| self.assertEqual(report['output_tokens'],3) |
| self.assertEqual(report['timed_decode_steps'],2) |
| self.assertEqual(report['stop_reason'],'eos') |
|
|
| def test_output_limit(self): |
| output,report=generate(self.model(),'Hello',1) |
| self.assertEqual(output,'20');self.assertEqual(report['timed_decode_steps'],0) |
|
|
| def test_invalid_budget(self): |
| with self.assertRaises(ValueError):generate(self.model(),'Hello',0) |
|
|
|
|
| if __name__=='__main__':unittest.main() |
|
|
