#!/usr/bin/env python
# coding: utf-8

# In[14]:


from metaphor_module import *
import ast
from transformers import AutoTokenizer, AutoModelForSequenceClassification


tokenizer = AutoTokenizer.from_pretrained("roberta-large-mnli")
model = AutoModelForSequenceClassification.from_pretrained("roberta-large-mnli")
model.eval()
print(">>>>>>> robert-large-mnli model ready")


ncorrect, nsamples = 0, 0
label_map = {0: 'not_entailment', 1: 'not_entailment', 2:'entailment'}

verbnet = return_verbnet()


for line in open('datasets/metaphor-entail.txt'):
    line = ast.literal_eval(line)
    for i, dic in enumerate(line):
        premise = dic["premise"]
        hypothesis = dic["hypothesis"]
        
        key_verbs, verbs = return_keywords(premise, hypothesis)
        kb_sent = create_kbsent(verbnet, key_verbs, verbs)

        if kb_sent == '':
            kb_sent = create_kbsent_mix(verbnet, key_verbs, verbs)

        encoded_input = tokenizer(premise +'</s>'+ kb_sent +'</s></s>'+hypothesis, add_special_tokens=True, truncation=True, return_tensors='pt')
        output = model(encoded_input['input_ids'], output_hidden_states=False)
        logits = output.logits.detach()
        predict = logits[0].argmax().item()
        
        nsamples += 1
        if label_map[predict]==dic["label"]:
            ncorrect += 1


print(nsamples)
print(ncorrect)
print('| Accuracy: ', float(ncorrect)/float(nsamples))
