in summarize_from_feedback/models/sample_fns.py [0:0]
def nucleus_sampler(top_p: float = 0.9, temperature=1.0) -> Sampler:
"""
Return a sampler that decides diversity via nucleus sampling.
p=0.9 means that the top 90% of likelihood-weighted options are considered. p=0.0 is
equivalent to argmax, p=1.0 has no effect.
When a logit is on the boundary of being included or not being included, default
to including it.
"""
if top_p == 0.0:
return argmax()
if top_p == 1.0:
return standard(temperature=temperature)
def sample(logits: Logits) -> Sample:
"""
Remove logits that do not represent the top_p proportion of likelihoods.
When a logit is on the boundary of being included or not being included, default
to including it.
"""
logits = logits.clone()
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
# Remove tokens with cumulative probability above the threshold.
sorted_indices_to_remove = cumulative_probs > top_p
# Shift the indices to the right to keep also the first token above the threshold.
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool).scatter_(
dim=-1, index=sorted_indices, src=sorted_indices_to_remove
)
logits[indices_to_remove] = -float("Inf")
return standard(temperature=temperature)(logits)
return sample