from ..base import BaseModel


class BatchParams(BaseModel):
    input_ids: list[list[int]]
    """输入 token，一个 batch 的"""

    attention_mask: list[list[int]]
    """掩码，一个 batch 的"""

    target_ids: list[list[int]]
    """目标 token，一个 batch 的"""

    def count_input_token(self):
        return sum(len(sublist) for sublist in self.input_ids)
