F1 specificity recall
総合 88.96 89.07 88.92
jsquad 97.02 95.46 98.42
jnli 78.04 96.89 75.29
jdolly 97.94 98.46 97.54
jaquad 95.58 93.07 98.03
jfldD1 79.28 80.96 75.94
jfldD3 78.13 74.03 77.9

class NLIPredictor:
    def __init__(self, model_name, model_path, device="cpu", max_length=2048):
        device = "cuda"
        self.device = torch.device(device)
        self.max_length = max_length

        self.model = AutoModelForSequenceClassification.from_pretrained(model_path).to(
            self.device
        )
        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

    def predict(self, premise, hypothesis):
        inputs = self.tokenizer(
            premise,
            hypothesis,
            return_tensors="pt",
            max_length=self.max_length,
            truncation=True,
            padding=True,
        )
        inputs = {k: v.to(self.device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=1)
            predicted_class = torch.argmax(probabilities, dim=1).item()

        return predicted_class