
import torch
import ast
import simile_module as sim
import metaphor_module as met
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModelForSequenceClassification

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

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

##To use ATOMIC 2020 as KB
#atomic_set = sim.return_triple("atomic/atomic.tsv")
#smodel = SentenceTransformer('all-MiniLM-L6-v2')

#To use VerbNet as KB
verbnet = met.return_verbnet()


#open manual_e.txt or manual_ne.txt
with open('manual_e.txt', 'r', encoding='utf-8') as file:
    for lines in file:
        l = lines.strip().split('\t')
        premise, hypothesis = l[0], l[1]

        key_verbs, verbs = met.return_keywords(premise, hypothesis)

        if len(verbs) == 2:
            kb_sent = met.create_kbsent_mix(verbnet, key_verbs, verbs)

            encoded_input = tokenizer(premise + ' ' + kb_sent + '</s></s>' + 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()
            
        #else:
            #fig_exp, lit_exp, key_expression = sim.return_keywords(premise, hypothesis)
            #kb_sents = sim.return_candidates(atomic_set, premise, hypothesis, fig_exp, lit_exp)
            #top5 = sim.semantic_sim(kb_sents, key_expression, smodel, 5)
            #new_premise = sim.return_new_premise(premise, top5, 0.4)
            #
            #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(premise +'</s></s>'+ 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] == 'entailment':#or 'not_entailment' for manual_ne.txt 
            ncorrect += 1
            

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

    


