    def forward(
        self,
        phone_ids,
        phone_mask,
        target_ids,
        target_mask,
        target_quantization_layer=None,
        prompt_len=None,
        dropout=0.0,
    ):
        """
        phone_ids: [B, T]
        phone_mask: [B, T]
        target_ids: [8,B,T]
        target_mask: [B, T]
        dropout: rate of dropping out the target tokens
        """
        assert (target_ids < 1024).all(), "target_ids should be less than 1024"
        phone_ids = phone_ids + self.target_vocab_size
        phone_ids = phone_ids * phone_mask + (1 - phone_mask) * self.pad_token_id
        # assert (phone_ids >= 1024).all(), "phone_ids should be greater than 1024"
        # phone_ids, phone_mask, phone_label = self.add_phone_eos_bos_label(
        #     phone_ids,
        #     phone_mask,
        #     self.eos_phone_id,
        #     self.bos_phone_id,
        #     self.pad_token_id,
        # )
        phone_label = -100 * (1 - phone_mask)
        # get phone embedding
        phone_embedding = self.phone_embedder(
            phone_ids - self.target_vocab_size
        )  # [B, T, H]

        if prompt_len is not None:
            assert not self.training  # inference stage fix prompt len to input
            NUM_PROMPT_TOKENS = prompt_len
        else:
            assert self.training
            # randomly select a prompt length
            assert self.training  # randomize prompt len in training
            NUM_PROMPT_TOKENS = np.random.randint(
                min(target_ids.shape[-1] // 4, 5), target_ids.shape[-1] // 2
            )

        # extract 8-level prompts
        prompt_tokens = target_ids[:, :, :NUM_PROMPT_TOKENS]  # [Q, B, T]
        prompt_mask = torch.ones_like(prompt_tokens[0])
        prompt_label = -100 * prompt_mask
        # get prompt embedding
        prompt_embedding = self.prompt_embedder(prompt_tokens)  # [B, T, H]

        # randomly select a target qnt layer to predict
        # total quant layer is 0 to 7
        if target_quantization_layer is None:
            if self.mask_layer_schedule == "linear":
                weights = torch.tensor(
                    [
                        NUM_QUANTIZERS - i
                        for i in range(
                            START_QUANTIZATION_LAYER, END_QUANTIZATION_LAYER + 1
                        )
                    ]
                )
                weights = weights / weights.sum()
                mask_layer = (
                    torch.multinomial(weights, 1, replacement=True)
                    + START_QUANTIZATION_LAYER
                )
                assert (
                    mask_layer >= START_QUANTIZATION_LAYER
                    and mask_layer <= END_QUANTIZATION_LAYER
                )
                target_quantization_layer = mask_layer.item()
            elif self.mask_layer_schedule == "cosine":
                weights = torch.tensor(
                    [
                        np.cos(i / NUM_QUANTIZERS * np.pi / 2)
                        for i in range(
                            START_QUANTIZATION_LAYER, END_QUANTIZATION_LAYER + 1
                        )
                    ]
                )
                weights = weights / weights.sum()
                mask_layer = (
                    torch.multinomial(weights, 1, replacement=True)
                    + START_QUANTIZATION_LAYER
                )
                assert (
                    mask_layer >= START_QUANTIZATION_LAYER
                    and mask_layer <= END_QUANTIZATION_LAYER
                )
                target_quantization_layer = mask_layer.item()
                breakpoint()
            elif self.mask_layer_schedule == "uniform":
                target_quantization_layer = np.random.randint(
                    START_QUANTIZATION_LAYER, END_QUANTIZATION_LAYER + 1
                )

            # print(f'target layer: {target_quantization_layer}')
        # prompt of the target part
        target_prompt_ids = target_ids[
            :target_quantization_layer, :, NUM_PROMPT_TOKENS:
        ]

        def randomly_set_elements(tensor, fraction, value):
            """
            Randomly set a fraction of the elements in a tensor to a specific value.

            Args:
            tensor (torch.Tensor): The input tensor.
            fraction (float): The fraction of elements to set to the specified value (between 0 and 1).
            value (float or int): The value to set the elements to.

            Returns:
            torch.Tensor: The tensor with some elements set to the specified value.
            """
            # Create a mask with the same shape as the tensor
            mask = torch.rand_like(tensor, dtype=torch.float32) < fraction
            # Clone the tensor to avoid modifying the original tensor
            result_tensor = tensor.clone()
            # Set the elements where the mask is True to the specified value
            result_tensor[mask] = value
            return result_tensor

        if dropout != 0.0:
            target_prompt_ids = randomly_set_elements(
                target_prompt_ids, dropout, self.target_vocab_size
            )

        target_embedding = self.prompt_embedder(target_prompt_ids)

        # mask of the target part
        target_mask = target_mask[:, NUM_PROMPT_TOKENS:]

        target_labels = target_ids[
            target_quantization_layer, :, NUM_PROMPT_TOKENS:
        ] * target_mask + (-100 * (1 - target_mask))

        # input embeddings
        input_embeddings = torch.cat(
            [phone_embedding, prompt_embedding, target_embedding], dim=1
        )
        input_mask = torch.cat([phone_mask, prompt_mask, target_mask], dim=1)  # [B, T]
        prediction_target = torch.cat(
            [phone_label, prompt_label, target_labels], dim=1
        )  # [B, T]

        out = self.model(
            cond=torch.tensor(
                target_quantization_layer,
                device=prediction_target.device,
                dtype=torch.long,
            ),
            input_ids=input_embeddings,
            prediction_target=prediction_target,
            attention_mask=input_mask,
            return_dict=True,
        )
        logits = out.logits[:, -target_embedding.shape[1] :, :]
        targets = prediction_target[..., -target_embedding.shape[1] :]
        top1_acc = logits.argmax(-1) == targets
        top1_acc = (top1_acc * target_mask).sum() / target_mask.sum()

        top5_acc = (logits.topk(5, dim=-1).indices == targets.unsqueeze(-1)).any(-1)
        top5_acc = (top5_acc * target_mask).sum() / target_mask.sum()

        top10_acc = (logits.topk(10, dim=-1).indices == targets.unsqueeze(-1)).any(-1)
        top10_acc = (top10_acc * target_mask).sum() / target_mask.sum()

        out.target_quantization_layer = target_quantization_layer
        out.top1_acc = top1_acc
        out.top5_acc = top5_acc
        out.top10_acc = top10_acc

        return out
phone_ids: [B, T]
phone_mask: [B, T]
target_ids: [8,B,T]
target_mask: [B, T]
dropout: rate of dropping out the target tokens
models/tts/valle_v2/valle_nar.py