න්යෂ්ටිකනියැදීම

29import torch
30from torch import nn
31
32from labml_nn.sampling import Sampler

35class NucleusSampler(Sampler):

39    def __init__(self, p: float, sampler: Sampler):

44        self.p = p
45        self.sampler = sampler

47        self.softmax = nn.Softmax(dim=-1)

49    def __call__(self, logits: torch.Tensor):

55        probs = self.softmax(logits)

58        sorted_probs, indices = torch.sort(probs, dim=-1, descending=True)

60        cum_sum_probs = torch.cumsum(sorted_probs, dim=-1)

62        nucleus = cum_sum_probs < self.p

65        nucleus = torch.cat([nucleus.new_ones(nucleus.shape[:-1] + (1,)), nucleus[..., :-1]], dim=-1)

68        sorted_log_probs = torch.log(sorted_probs)
69        sorted_log_probs[~nucleus] = float('-inf')

72        sampled_sorted_indexes = self.sampler(sorted_log_probs)

75        res = indices.gather(-1, sampled_sorted_indexes.unsqueeze(-1))

78        return res.squeeze(-1)