43 lines
1.6 KiB
Python
43 lines
1.6 KiB
Python
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from prain import drafts
|
|
from tests.helpers import TmpCase # noqa: F401 (every test module imports the helpers)
|
|
|
|
|
|
def new_request(req_id, xargs):
|
|
return SimpleNamespace(req_id=req_id, sampling_params=SimpleNamespace(extra_args=xargs))
|
|
|
|
|
|
def step(new=(), finished=()):
|
|
return SimpleNamespace(scheduled_new_reqs=list(new), finished_req_ids=set(finished))
|
|
|
|
|
|
class Drafts(unittest.TestCase):
|
|
def setUp(self):
|
|
drafts._DRAFTS.clear()
|
|
drafts.clear()
|
|
|
|
def test_draft_rows_follow_the_batch_layout(self):
|
|
# a: 3 prompt tokens (full); b: 2 prompt tokens (draft)
|
|
drafts.note_step(step([new_request("a", None), new_request("b", {"prain": "draft"})]),
|
|
["a", "b"], [0, 3, 5])
|
|
self.assertEqual(drafts.spans(), [(3, 5)])
|
|
# next step: both decode one token each, in a new order
|
|
drafts.note_step(step(), ["b", "a"], [0, 1, 2])
|
|
self.assertEqual(drafts.spans(), [(0, 1)])
|
|
drafts.clear()
|
|
self.assertIsNone(drafts.spans())
|
|
|
|
def test_finished_drafts_are_forgotten(self):
|
|
drafts.note_step(step([new_request("b", {"prain": "draft"})]), ["b"], [0, 4])
|
|
drafts.note_step(step(finished=["b"]), [], [0])
|
|
self.assertIsNone(drafts.spans())
|
|
self.assertEqual(drafts._DRAFTS, set())
|
|
|
|
def test_requests_without_the_flag_are_full(self):
|
|
drafts.note_step(step([new_request("a", {"other": 1}),
|
|
SimpleNamespace(req_id="c", sampling_params=None)]),
|
|
["a", "c"], [0, 1, 2])
|
|
self.assertIsNone(drafts.spans())
|