| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| from easydict import EasyDict as edict |
|
|
|
|
| @dataclass |
| class SampleOutputBase: |
| true_local_ids_map: dict[str, dict[str, list[int]]] |
| accepted_ids_map: dict[str, dict[str, list[int]]] |
| sampled_tokens_map: dict[str, dict[str, list[int]]] |
| mask_token_rel_ids_map: dict[str, dict[str, list[int]]] | None = None |
| confidence_map: dict[str, dict[str, list[float]]] | None = None |
| initial_confidence_map: dict[str, dict[str, list[float]]] | None = None |
| edit_writes_map: dict[str, dict[str, dict[int, int]]] | None = None |
| block_state_map: dict[str, dict[str, dict]] | None = None |
|
|
| def __post_init__(self): |
| req_ids = set(self.accepted_ids_map.keys()) |
| self.accepted_ids_map = edict(self.accepted_ids_map) |
| self.sampled_tokens_map = edict(self.sampled_tokens_map) |
| self.true_local_ids_map = edict(self.true_local_ids_map) |
| self.mask_token_rel_ids_map = edict(self.mask_token_rel_ids_map or {}) |
| self.confidence_map = edict(self.confidence_map or {}) |
| self.initial_confidence_map = edict(self.initial_confidence_map or {}) |
| edit_writes_map = self.edit_writes_map or {} |
| block_state_map = self.block_state_map or {} |
| for req_id_str in req_ids: |
| edit_writes_map.setdefault(req_id_str, {}) |
| block_state_map.setdefault(req_id_str, {}) |
| self.edit_writes_map = edict(edit_writes_map) |
| self.block_state_map = edict(block_state_map) |
|
|
|
|
| def merge_sample_outputs(outputs: list[SampleOutputBase | None]) -> SampleOutputBase: |
| true_local_ids_map: dict[str, dict[str, list[int]]] = {} |
| accepted_ids_map: dict[str, dict[str, list[int]]] = {} |
| sampled_tokens_map: dict[str, dict[str, list[int]]] = {} |
| mask_token_rel_ids_map: dict[str, dict[str, list[int]]] = {} |
| confidence_map: dict[str, dict[str, list[float]]] = {} |
| initial_confidence_map: dict[str, dict[str, list[float]]] = {} |
| edit_writes_map: dict[str, dict[str, dict[int, int]]] = {} |
| block_state_map: dict[str, dict[str, dict]] = {} |
|
|
| for output in outputs: |
| if output is None: |
| continue |
| true_local_ids_map.update(dict(output.true_local_ids_map)) |
| accepted_ids_map.update(dict(output.accepted_ids_map)) |
| sampled_tokens_map.update(dict(output.sampled_tokens_map)) |
| mask_token_rel_ids_map.update(dict(output.mask_token_rel_ids_map)) |
| confidence_map.update(dict(output.confidence_map)) |
| initial_confidence_map.update(dict(output.initial_confidence_map)) |
| edit_writes_map.update(dict(output.edit_writes_map)) |
| block_state_map.update(dict(output.block_state_map)) |
|
|
| return SampleOutputBase( |
| true_local_ids_map=true_local_ids_map, |
| accepted_ids_map=accepted_ids_map, |
| sampled_tokens_map=sampled_tokens_map, |
| mask_token_rel_ids_map=mask_token_rel_ids_map, |
| confidence_map=confidence_map, |
| initial_confidence_map=initial_confidence_map, |
| edit_writes_map=edit_writes_map, |
| block_state_map=block_state_map, |
| ) |
|
|
