Method streaming_topk_search

moss_inference.py:193–289

Performs a streaming top-k search using the given parameters. Args: input_ids (torch.Tensor): The input IDs tensor. attention_mask (torch.Tensor): The attention mask tensor. temperature (float, optional): The temperature for logits. Defaults to 0

(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        temperature: float = 0.7,
        repetition_penalty: float = 1.02,
        top_k: int = 0,
        top_p: float = 0.8,
        max_iterations: int = 1024,
        regulation_start: int = 512,
        length_penalty: float = 1,
        max_time: int = 60,
    )

        return preds_i[len(self.prefix):]

    def streaming_topk_search(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        temperature: float = 0.7,
        repetition_penalty: float = 1.02,
        top_k: int = 0,
        top_p: float = 0.8,
        max_iterations: int = 1024,
        regulation_start: int = 512,
        length_penalty: float = 1,
        max_time: int = 60,
    ) -> torch.Tensor:
        """
        Performs a streaming top-k search using the given parameters.

        Args:
            input_ids (torch.Tensor): The input IDs tensor.
            attention_mask (torch.Tensor): The attention mask tensor.
            temperature (float, optional): The temperature for logits. Defaults to 0.7.
            repetition_penalty (float, optional): The repetition penalty factor. Defaults to 1.02.
            top_k (int, optional): The top-k value for filtering. Defaults to 0.
            top_p (float, optional): The top-p value for filtering. Defaults to 0.92.
            max_iterations (int, optional): The maximum number of iterations. Defaults to 1024.
            regulation_start (int, optional): The number of iterations after which regulation starts. Defaults to 512.
            length_penalty (float, optional): The length penalty factor. Defaults to 1.
            max_time (int, optional): The maximum allowed time in seconds. Defaults to 60.

        Returns:
            torch.Tensor: The generated output IDs tensor.
        """
        assert input_ids.dtype == torch.int64 and attention_mask.dtype == torch.int64

        self.bsz, self.seqlen = input_ids.shape

        input_ids, attention_mask = input_ids.to('cuda'), attention_mask.to('cuda')
        last_token_indices = attention_mask.sum(1) - 1

        moss_stopwords = self.moss_stopwords.to(input_ids.device)
        queue_for_moss_stopwords = torch.empty(size=(self.bsz, len(self.moss_stopwords)), device=input_ids.device, dtype=input_ids.dtype)
        all_shall_stop = torch.tensor([False] * self.bsz, device=input_ids.device)
        moss_stop = torch.tensor([False] * self.bsz, device=input_ids.device)

        generations, start_time = torch.ones(self.bsz, 1, dtype=torch.int64), time.time()

        past_key_values = None
        for i in range(int(max_iterations)):
            logits, past_key_values = self.infer_(input_ids if i == 0 else new_generated_id, attention_mask, past_key_values)
            
            if i == 0: 
                logits = logits.gather(1, last_token_indices.view(self.bsz, 1, 1).repeat(1, 1, self.vocab_size)).squeeze(1)
            else: 
                logits = logits[:, -1, :]


            if repetition_penalty > 1:
                score = logits.gather(1, input_ids)
                # if score < 0 then repetition penalty has to be multiplied to reduce the previous token probability