def add_extra_inputs(self, batch: dict[str, torch.Tensor]):
        """Add extra inputs to the batch.

        :param batch: A dictionary containing the batch data.
        :return: The updated batch with extra inputs added.
        """
        if "attention_mask" not in batch:
            # no attention mask
            return

        attention_mask = batch["attention_mask"]
        if self.position_ids and "position_ids" not in batch:
            # position ids: current tokens
            # attention mask: past + current tokens
            batch["position_ids"] = (attention_mask.cumsum(-1) - 1)[:, -batch["input_ids"].shape[-1] :]
        if self.past_seq_len and "past_seq_len" not in batch:
            # past seq len: past tokens
            # attention mask: past + current tokens
            batch["past_seq_len"] = attention_mask.sum(-1, keepdim=True) - 1
            batch["total_seq_len"] = torch.tensor(attention_mask.shape[-1])
            del batch["attention_mask"]
Add extra inputs to the batch.

        :param batch: A dictionary containing the batch data.
        :return: The updated batch with extra inputs added.
        
olive/data/component/dataloader.py