Get topk
(logits, topk=5)
| 26 | |
| 27 | |
| 28 | def topk_fun(logits, topk=5): |
| 29 | """Get topk""" |
| 30 | value = np.flip(np.sort(logits), axis=-1)[..., :topk] |
| 31 | index = np.flip(np.argsort(logits), axis=-1)[..., :topk] |
| 32 | return value, index |
| 33 | |
| 34 | |
| 35 | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False, bad_words_index=[]): |