from abc import ABC, abstractmethod
from datasets.arrow_dataset import Dataset
import numpy as np


class Reranker(ABC):
    @abstractmethod
    def rerank_batch(self, dataset: Dataset) -> Dataset:
        pass

    @staticmethod
    def sort_concepts(concepts):
        idx = np.argsort(np.array(concepts["score"]))[-1::-1]
        result = {}
        for k, v in concepts.items():
            result[k] = list(np.array(v)[idx])
        return result


from .cross_encoder import CrossEncoderReranker
