| """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) |
|
|
