File size: 792 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | from __future__ import annotations
class EditSchedulerMixin:
def apply_edit_writes_map(self, req, sample_output) -> None:
req_id_str = str(req.req_id)
edit_writes_map = sample_output.edit_writes_map[req_id_str]
for block_id, block_writes in edit_writes_map.items():
if not block_writes:
continue
dllm_block = req.dllm_blocks[int(block_id)]
for rel_idx, token in block_writes.items():
rel_idx = int(rel_idx)
token = int(token)
prev_token = dllm_block.token_ids[rel_idx]
dllm_block.write_token(token, rel_idx)
if prev_token == dllm_block.mask_token_id and token != dllm_block.mask_token_id:
req.new_tokens += 1
|