
import torch
import ast
from simile_module import *
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSequenceClassification



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


smodel = SentenceTransformer('all-MiniLM-L6-v2')

atomic_set = return_atomic("atomic/atomic.tsv")

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

nsamples, ncorrect = 0, 0
#incorrect = []



for line in open('datasets/new-simile-entail.txt'):
    line = ast.literal_eval(line)
    for dic in line:
        premise = dic['premise'][:-1]
        hypothesis = dic['hypothesis']

        fig_exp, lit_exp, key_expression = return_keywords(premise, hypothesis)
        kb_sents = return_candidates(atomic_set, premise, hypothesis, fig_exp, lit_exp)
        top5 = semantic_sim(kb_sents, key_expression, smodel, 5)
        new_premise = return_new_premise(premise, top5, 0.4)
        
        
        ## *** Using additional knowledge bases conceptnet and wordnet ***
        #if len(top5) == 0 and len(h_exp) == 1:
        #    synonyms_cn = return_synonyms_cn(h_exp)
        #    synonyms_wn = return_synonyms_wn(h_exp)
        #    final_synonyms = list(set(synonyms_wn + synonyms_cn))
        #    if len(final_synonyms) != 0:
        #        new_hypotheses = return_hypotheses(hypothesis, final_synonyms, h_exp)
        #        hypothesis = semantic_sim(new_hypotheses, fig_exp, smodel, 5)[0][1]
        #    else:
        #        hypothesis = hypothesis

        
        encoded_input = tokenizer(new_premise, hypothesis, add_special_tokens=True, truncation=True, return_tensors='pt')
        model_output = model(encoded_input['input_ids'], output_hidden_states=False)
        predict = model_output.logits.argmax().item()

        nsamples += 1
        if label_map[predict] == dic['label']:
            ncorrect += 1


print(nsamples)
print(ncorrect)
print('| Accuracy: ', float(ncorrect)/float(nsamples))
print('\n')
#print(incorrect)

    


