Method sample_top_p

accessory/model/meta.py:550–565

Sample a token based on the provided probability distribution using top-p sampling. :param probs: The probability distribution for the next token. :param p: The cumulative probability threshold for top-p sampling. :return: The sampled next token.

(self, probs, p)

        yield {"text": generated, "end_of_content": True}

    def sample_top_p(self, probs, p):
        """
        Sample a token based on the provided probability distribution using top-p sampling.

        :param probs: The probability distribution for the next token.
        :param p: The cumulative probability threshold for top-p sampling.
        :return: The sampled next token.
        """
        probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
        probs_sum = torch.cumsum(probs_sort, dim=-1)
        mask = probs_sum - probs_sort > p
        probs_sort[mask] = 0.0
        probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
        next_token = torch.multinomial(probs_sort, num_samples=1)
        next_token = torch.gather(probs_idx, -1, next_token)
        return next_token

    def get_image_words(self):
        return self.llma.image_words