File size: 1,437 Bytes
ac12f15 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | """Experimental top-k sampling: tied boundary logits may select a different support."""
import torch
from .protocol import EOD, ABC_END, MUSIC_END, CODEC_OFFSET, CODEC_SIZE
from .sampling import window_penalty, distribution
def sample_topk(logits, sampling, history, step, phase, legacy_off, generator):
if legacy_off or sampling.temperature==0:
scores=distribution(logits,sampling,history,step,phase,legacy_off)
return scores.argmax(-1,keepdim=True) if sampling.temperature==0 else torch.multinomial(scores.softmax(-1),1,generator=generator)
scores=logits.float().clone()
end=ABC_END if phase=='abc' else MUSIC_END
mask=torch.full_like(scores,float('-inf'))
if phase=='abc':mask[...,:EOD]=0
else:mask[...,CODEC_OFFSET:CODEC_OFFSET+CODEC_SIZE]=0
mask[...,end]=0
scores+=mask
if step<sampling.min_tokens:scores[...,end]=-torch.inf
scores=window_penalty(scores,history[-sampling.penalty_window:],sampling.repetition_penalty)
if sampling.temperature!=1:scores/=sampling.temperature
values,indices=scores.topk(min(sampling.top_k,scores.shape[-1]))
if sampling.top_p<1:
probabilities=values.softmax(-1)
removed=probabilities.cumsum(-1)-probabilities>sampling.top_p
removed[...,:1]=False
values=values.masked_fill(removed,-torch.inf)
choice=torch.multinomial(values.softmax(-1),1,generator=generator)
return indices.gather(-1,choice)
|