guidellm.data.deserializers.trace_mooncake

@TraceFormatRegistry.register("mooncake")
class MooncakeTraceFormat(TraceFormatBase):
    """Mooncake trace format requires a column for timestamps, prompt token counts,
    ouput token counts and lists of hash IDs.

    Hash IDs are globally unique identifiers based on the current and previous token
    blocks in a prompt. The relationships of IDs forms a tree, where every first ID
    in a prompt has a parent node of `None`. Parent nodes can have an unbounded
    number of children. Two hash IDs can represent identical blocks of tokens so long
    as they do not share the same parent (previous ID). For more details, see section 4
    of https://arxiv.org/pdf/2407.00079.

    Generated prompts match the prompt token count of the row."""

    def __init__(self) -> None:
        self.hash_id_table: list[Any] = []
        self.sibling_token_blocks: dict[Any, list[list[int]]] = {}

    def required_columns(self, config: MooncakeTraceFormatArgs) -> Features:
        return Features({config.hash_ids_column: List(Value("int32"))})

    def validate_row(self, config: MooncakeTraceFormatArgs, row: dict) -> None:
        n_in = row[config.prompt_tokens_column]
        n_blocks = len(row[config.hash_ids_column])
        for hash_id in row[config.hash_ids_column]:
            if hash_id < 0:
                raise DataNotSupportedError(
                    f"Hash ID must be non-negative, got {hash_id}"
                )
        if math.ceil(n_in / config.hash_id_block_size) != n_blocks:
            raise DataNotSupportedError(
                f"Input token count of {n_in} split into blocks of size "
                f"{config.hash_id_block_size} does not match given {n_blocks} blocks"
            )

    def create_prompt(
        self,
        config: MooncakeTraceFormatArgs,
        row: dict,
        processor: PreTrainedTokenizerBase,
        faker: Faker,
    ) -> str:
        """Before generating the prompt, this first generates a block of tokens for
        each hash ID that has not already been seen."""
        ids = row[config.hash_ids_column]
        for idx, hash_id in enumerate(ids):
            if not _is_in_table(self.hash_id_table, hash_id):
                _resize_to_hold_id(self.hash_id_table, hash_id)
                prev_id = None if idx == 0 else ids[idx - 1]
                num_tokens = _calculate_required_prompt_tokens(config, row, hash_id)
                self.sibling_token_blocks.setdefault(prev_id, [])
                self.hash_id_table[hash_id] = _create_distinct_token_block(
                    num_tokens,
                    self.sibling_token_blocks[prev_id],
                    processor,
                    faker,
                )
                self.sibling_token_blocks[prev_id].append(self.hash_id_table[hash_id])
        return _create_prompt_from_hash_ids(ids, self.hash_id_table, processor)