from typing import override

from srs2std.evaluate.factcheck_strategies.scorers.scorer import (
    BERT,
    DEFAULT_BATCH_SIZE,
    DEFAULT_DEVICE,
    T5,
)


class MiniCheckDeBERTa(BERT):
    model_path = "lytang/MiniCheck-DeBERTa-v3-Large"
    max_model_len = 2048
    device = DEFAULT_DEVICE
    batch_size = DEFAULT_BATCH_SIZE // 4

    # setting for the classifier
    label2id = {"unsupported": 0, "supported": 1}


class MiniCheckFlanT5(T5):
    model_path = "lytang/MiniCheck-Flan-T5-Large"
    max_model_len = 2048
    device = DEFAULT_DEVICE
    batch_size = DEFAULT_BATCH_SIZE // 4

    # setting for T5
    pos_token_id = 209
    neg_token_id = 3

    @override
    def _concat_doc_and_claim(self, doc, claim):
        return "predict: " + self.tokenizer.eos_token.join([doc, claim])


class MiniCheckRoBERTa(BERT):
    model_path = "lytang/MiniCheck-RoBERTa-Large"
    max_model_len = 512
    device = DEFAULT_DEVICE
    batch_size = DEFAULT_BATCH_SIZE

    # setting for the classifier
    label2id = {"unsupported": 0, "supported": 1}
