(filtered_indexes, scores, k)
| 157 | |
| 158 | @staticmethod |
| 159 | def select_topk(filtered_indexes, scores, k): |
| 160 | if len(filtered_indexes) > k: |
| 161 | sorted_ = np.argpartition(scores, k)[:k] |
| 162 | filtered_indexes, scores = filtered_indexes[sorted_], -scores[sorted_] |
| 163 | else: |
| 164 | scores = -scores |
| 165 | return filtered_indexes, scores |
| 166 | |
| 167 | @staticmethod |
| 168 | @numba.njit(nogil=True, parallel=True, cache=True) |