from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

class OmographModel:
    def __init__(self, allow_cuda=True) -> None:
        self.device = torch.device('cuda' if torch.cuda.is_available()  and allow_cuda else 'cpu')
        
    def load(self, path):
        self.nli_model = AutoModelForSequenceClassification.from_pretrained(path, torch_dtype=torch.bfloat16).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        
    def classify(self, text, hypotheses):
        encodings = self.tokenizer.batch_encode_plus([(text, hyp) for hyp in hypotheses], return_tensors='pt', padding=True)
        input_ids = encodings['input_ids'].to(self.device)
        with torch.no_grad():
            logits = self.nli_model(input_ids)[0]
            entail_contradiction_logits = logits[:,[0,2]]
            probs = entail_contradiction_logits.softmax(dim=1)
            prob_label_is_true = [float(p[1]) for p in probs]

        return hypotheses[prob_label_is_true.index(max(prob_label_is_true))]