Get topk
(logits, topk=5)
| 61 | |
| 62 | |
| 63 | def topk_fun(logits, topk=5): |
| 64 | """Get topk""" |
| 65 | # target_column = logits[0].tolist() |
| 66 | # sorted_array = [(k, v) for k, v in enumerate(target_column)] |
| 67 | # sorted_array.sort(key=lambda x: x[1], reverse=True) |
| 68 | # topk_array = sorted_array[:topk] |
| 69 | # index, value = zip(*topk_array) |
| 70 | # index = np.array([index]) |
| 71 | # value = np.array([value]) |
| 72 | value = np.flip(np.sort(logits), axis=-1)[..., :topk] |
| 73 | index = np.flip(np.argsort(logits), axis=-1)[..., :topk] |
| 74 | return value, index |
| 75 | |
| 76 | |
| 77 | def sampler(log_probs_revised, top_p, top_k_num, use_pynative=False, bad_words_index=[]): |