clarin-knext/bge-reranker-v2-m3-citeverifier

from sentence_transformers.cross_encoder import CrossEncoder


class CROSSENCODER:
    def __init__(self, model_name, sim_threshold=0.5):
        self.model = CrossEncoder(model_name)
        self.sim_threshold = sim_threshold

    def __call__(self, premise: str, hypothesis: str) -> bool:
        scores = self.model.predict([[premise, hypothesis]])
        return float(scores[0]) > self.sim_threshold  # Here depends on num_labels during training