# from https://github.com/XuhuiZhou/CATS/blob/master/making_sense.py

import torch
from torch.nn import CrossEntropyLoss
from transformers import *
import numpy as np
import math
# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
import os
import sys
logging.basicConfig(level=logging.INFO)
os.environ["CUDA_VISIBLE_DEVICES"]="0"
def uni_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = text
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
    length = len(tokenized_text)
    tokens_tensor = torch.tensor([indexed_tokens])
    tokens_tensor = tokens_tensor.to('cuda')
    #masked_tensor = torch.tensor([masked_index])
    with torch.no_grad():
        outputs = model(tokens_tensor, labels= tokens_tensor)
    loss = outputs[0]
    sentence_score = -loss
    return sentence_score

def bert_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = "[CLS] " + text + " [SEP]" #special token for BERT, RoBERTa
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    length = len(tokenized_text)-2
    if length == 0:
        print(text, tokenized_text)
    for masked_index in range(1,len(tokenized_text)-1):
        # Mask a token that we will try to predict back with `BertForMaskedLM`
        masked_word = tokenized_text[masked_index]
        #tokenized_text[masked_index] = '<mask>' #special token for XLNet
        tokenized_text[masked_index] = '[MASK]' #special token for BERT, RoBerta
        # Convert token to vocabulary indices
        indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        tokens_tensor = torch.tensor([indexed_tokens])
        tokens_tensor = tokens_tensor.to('cuda')
        index = index.to('cuda')
        #masked_tensor = torch.tensor([masked_index])
        with torch.no_grad():
            outputs = model(tokens_tensor)
        prediction_scores = outputs[0]
        prediction_scores = prediction_scores.view(-1, model.config.vocab_size)
        prediction_scores = prediction_scores[masked_index].unsqueeze(0)
        loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
        masked_lm_loss = loss_fct(prediction_scores, index.view(-1))
        tokenized_text[masked_index] = masked_word
        sentence_score -= masked_lm_loss.item()
        tokenized_text[masked_index] = masked_word
    sentence_score = sentence_score/length
    return sentence_score

def ro_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = '<s> '+text+ ' </s>' #special token for RoBERTa
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    length = len(tokenized_text)-2
    for masked_index in range(1,len(tokenized_text)-1):
        # Mask a token that we will try to predict back with `BertForMaskedLM`
        masked_word = tokenized_text[masked_index]
        #tokenized_text[masked_index] = '<mask>' #special token for XLNet
        tokenized_text[masked_index] = '<mask>' #special token for BERT, RoBerta
        # Convert token to vocabulary indices
        indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        tokens_tensor = torch.tensor([indexed_tokens])
        tokens_tensor = tokens_tensor.to('cuda')
        index = index.to('cuda')
        #masked_tensor = torch.tensor([masked_index])
        with torch.no_grad():
            outputs = model(tokens_tensor)
        prediction_scores = outputs[0]
        prediction_scores = prediction_scores.view(-1, model.config.vocab_size)
        prediction_scores = prediction_scores[masked_index].unsqueeze(0)
        loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
        masked_lm_loss = loss_fct(prediction_scores, index.view(-1))
        tokenized_text[masked_index] = masked_word
        sentence_score -= masked_lm_loss.item()
        tokenized_text[masked_index] = masked_word
    sentence_score = sentence_score/length
    return sentence_score

def xlnet_predict(text, model, tokenizer):
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    #Sprint(len(tokenized_text))
    for masked_index in range(0,len(tokenized_text)):
        masked_word = tokenized_text[masked_index]
        masked_word = tokenized_text[masked_index]
        tokenized_text[masked_index] = '<mask>'
        input_ids = torch.tensor(tokenizer.convert_tokens_to_ids(tokenized_text)).unsqueeze(0)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float)
        perm_mask[:, :, masked_index] = 1.0  # Previous tokens don't see last token
        target_mapping = torch.zeros((1, 1, input_ids.shape[1]), dtype=torch.float)  # Shape [1, 1, seq_length] => let's predict one token
        target_mapping[0, 0, masked_index] = 1.0  # Our first (and only) prediction will be the last token of the sequence (the masked token)
        input_ids = input_ids.to('cuda')
        perm_mask = perm_mask.to('cuda')
        target_mapping = target_mapping.to('cuda')
        index = index.to('cuda')
        with torch.no_grad():
            outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping, labels = index)
        next_token_logits = outputs[0]
        length = len(tokenized_text)
        # predict_list = predictions[0, masked_index]
        sentence_score -= next_token_logits.item()
        tokenized_text[masked_index] = masked_word
    return sentence_score/(length)

if __name__ == '__main__':
    test = sys.argv[1]
    model_type = sys.argv[2]
    robust = sys.argv[3]
    if model_type=='xlnet':
        # For XLNet
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
        model = XLNetLMHeadModel.from_pretrained('xlnet-large-cased')
    elif model_type=='bert':
        # For BERT
        tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
        model = BertForMaskedLM.from_pretrained('bert-large-uncased')
    elif model_type=='roberta':
        tokenizer = RobertaTokenizer.from_pretrained('roberta-large')
        model = RobertaForMaskedLM.from_pretrained('roberta-large')
    elif model_type=='gpt':
        tokenizer = OpenAIGPTTokenizer.from_pretrained('openai-gpt')
        model = OpenAIGPTLMHeadModel.from_pretrained('openai-gpt')
    elif model_type=='gpt-2':
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
        model = GPT2LMHeadModel.from_pretrained('gpt2-medium')

    model.to('cuda')
    model.eval()

    if robust=='r':
        with open("CATS/Robust_commonsense_test/{}.txt".format(test), "r") as f:
            file = f.readlines()
        num = len(file)
        count = 0
        curr = 0
        for line in file:
            line = line.strip().split("\001")
            label_1 = int(line[0])
            label_2 = int(line[3])
            score_list = []
            for sentence in line:
                if not len(sentence)==1:
                    if model_type=='xlnet':
                        score = xlnet_predict(sentence, model=model, tokenizer=tokenizer)
                    elif model_type=='bert':
                        score = bert_predict(sentence, model=model, tokenizer=tokenizer)
                    elif model_type=='roberta':
                        score = ro_predict(sentence, model=model, tokenizer=tokenizer)
                    else:
                        score = uni_predict(sentence, model=model, tokenizer=tokenizer)
                    score_list.append(score)
            #print(score_list)
            score_list_1 = score_list[:2]
            score_list_2 = score_list[2:]
            predict_label_1 = score_list_1.index(max(score_list_1))
            predict_label_2 = score_list_2.index(max(score_list_2))
            if predict_label_1==label_1 and predict_label_2==label_2:
                count += 1
            elif predict_label_1!=label_1 and predict_label_2!=label_2:
                count += 1
            curr += 1
            #print (count, curr, count/curr)
        print(test+' '+model_type+':-------------------')
        print (count/num)
    else:
        with open("CATS/commonsense_ability_test/{}.txt".format(test), "r") as f:
            file = f.readlines()
        num = len(file)
        count = 0
        curr = 0
        for line in file:
            line = line.strip().split("\001")
            label = int(line[0])
            score_list = []
            for sentence in line[1:]:
                if model_type=='xlnet':
                    score = xlnet_predict(sentence, model=model, tokenizer=tokenizer)
                elif model_type=='bert':
                    score = bert_predict(sentence, model=model, tokenizer=tokenizer)
                elif model_type=='roberta':
                    score = ro_predict(sentence, model=model, tokenizer=tokenizer)
                else:
                    score = uni_predict(sentence, model=model, tokenizer=tokenizer)
                score_list.append(score)
            #print(score_list)
            predict_label = score_list.index(max(score_list))
            #print(predict_label, label)
            if predict_label==label:
                count += 1
            curr += 1
            #print (count, curr, count/curr)
        print(test+' '+model_type+':-------------------')
        print (count/num)

# from https://github.com/XuhuiZhou/CATS/blob/master/making_sense.py

import torch
from torch.nn import CrossEntropyLoss
from transformers import *
import numpy as np
import math
# OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
import logging
import os
import sys
logging.basicConfig(level=logging.INFO)
os.environ["CUDA_VISIBLE_DEVICES"]="0"
def uni_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = text
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
    length = len(tokenized_text)
    tokens_tensor = torch.tensor([indexed_tokens])
    tokens_tensor = tokens_tensor.to('cuda')
    #masked_tensor = torch.tensor([masked_index])
    with torch.no_grad():
        outputs = model(tokens_tensor, labels= tokens_tensor)
    loss = outputs[0]
    sentence_score = -loss
    return sentence_score

def bert_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = "[CLS] " + text + " [SEP]" #special token for BERT, RoBERTa
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    length = len(tokenized_text)-2
    if length == 0:
        print(text, tokenized_text)
    for masked_index in range(1,len(tokenized_text)-1):
        # Mask a token that we will try to predict back with `BertForMaskedLM`
        masked_word = tokenized_text[masked_index]
        #tokenized_text[masked_index] = '<mask>' #special token for XLNet
        tokenized_text[masked_index] = '[MASK]' #special token for BERT, RoBerta
        # Convert token to vocabulary indices
        indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        tokens_tensor = torch.tensor([indexed_tokens])
        tokens_tensor = tokens_tensor.to('cuda')
        index = index.to('cuda')
        #masked_tensor = torch.tensor([masked_index])
        with torch.no_grad():
            outputs = model(tokens_tensor)
        prediction_scores = outputs[0]
        prediction_scores = prediction_scores.view(-1, model.config.vocab_size)
        prediction_scores = prediction_scores[masked_index].unsqueeze(0)
        loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
        masked_lm_loss = loss_fct(prediction_scores, index.view(-1))
        tokenized_text[masked_index] = masked_word
        sentence_score -= masked_lm_loss.item()
        tokenized_text[masked_index] = masked_word
    sentence_score = sentence_score/length
    return sentence_score

def ro_predict(text, model, tokenizer):
    # Tokenized input
    # text = "[CLS] I got restricted because Tom reported my reply [SEP]"
    text = '<s> '+text+ ' </s>' #special token for RoBERTa
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    length = len(tokenized_text)-2
    for masked_index in range(1,len(tokenized_text)-1):
        # Mask a token that we will try to predict back with `BertForMaskedLM`
        masked_word = tokenized_text[masked_index]
        #tokenized_text[masked_index] = '<mask>' #special token for XLNet
        tokenized_text[masked_index] = '<mask>' #special token for BERT, RoBerta
        # Convert token to vocabulary indices
        indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        tokens_tensor = torch.tensor([indexed_tokens])
        tokens_tensor = tokens_tensor.to('cuda')
        index = index.to('cuda')
        #masked_tensor = torch.tensor([masked_index])
        with torch.no_grad():
            outputs = model(tokens_tensor)
        prediction_scores = outputs[0]
        prediction_scores = prediction_scores.view(-1, model.config.vocab_size)
        prediction_scores = prediction_scores[masked_index].unsqueeze(0)
        loss_fct = CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
        masked_lm_loss = loss_fct(prediction_scores, index.view(-1))
        tokenized_text[masked_index] = masked_word
        sentence_score -= masked_lm_loss.item()
        tokenized_text[masked_index] = masked_word
    sentence_score = sentence_score/length
    return sentence_score

def xlnet_predict(text, model, tokenizer):
    tokenized_text = tokenizer.tokenize(text)
    sentence_score = 0
    #Sprint(len(tokenized_text))
    for masked_index in range(0,len(tokenized_text)):
        masked_word = tokenized_text[masked_index]
        masked_word = tokenized_text[masked_index]
        tokenized_text[masked_index] = '<mask>'
        input_ids = torch.tensor(tokenizer.convert_tokens_to_ids(tokenized_text)).unsqueeze(0)
        index = torch.tensor(tokenizer.convert_tokens_to_ids(masked_word))
        perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float)
        perm_mask[:, :, masked_index] = 1.0  # Previous tokens don't see last token
        target_mapping = torch.zeros((1, 1, input_ids.shape[1]), dtype=torch.float)  # Shape [1, 1, seq_length] => let's predict one token
        target_mapping[0, 0, masked_index] = 1.0  # Our first (and only) prediction will be the last token of the sequence (the masked token)
        input_ids = input_ids.to('cuda')
        perm_mask = perm_mask.to('cuda')
        target_mapping = target_mapping.to('cuda')
        index = index.to('cuda')
        with torch.no_grad():
            outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping, labels = index)
        next_token_logits = outputs[0]
        length = len(tokenized_text)
        # predict_list = predictions[0, masked_index]
        sentence_score -= next_token_logits.item()
        tokenized_text[masked_index] = masked_word
    return sentence_score/(length)

if __name__ == '__main__':
    test = sys.argv[1]
    model_type = sys.argv[2]
    robust = sys.argv[3]
    if model_type=='xlnet':
        # For XLNet
        tokenizer = XLNetTokenizer.from_pretrained('xlnet-large-cased')
        model = XLNetLMHeadModel.from_pretrained('xlnet-large-cased')
    elif model_type=='bert':
        # For BERT
        tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
        model = BertForMaskedLM.from_pretrained('bert-large-uncased')
    elif model_type=='roberta':
        tokenizer = RobertaTokenizer.from_pretrained('roberta-large')
        model = RobertaForMaskedLM.from_pretrained('roberta-large')
    elif model_type=='gpt':
        tokenizer = OpenAIGPTTokenizer.from_pretrained('openai-gpt')
        model = OpenAIGPTLMHeadModel.from_pretrained('openai-gpt')
    elif model_type=='gpt-2':
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
        model = GPT2LMHeadModel.from_pretrained('gpt2-medium')

    model.to('cuda')
    model.eval()

    if robust=='r':
        with open("CATS/Robust_commonsense_test/{}.txt".format(test), "r") as f:
            file = f.readlines()
        num = len(file)
        count = 0
        curr = 0
        for line in file:
            line = line.strip().split("\001")
            label_1 = int(line[0])
            label_2 = int(line[3])
            score_list = []
            for sentence in line:
                if not len(sentence)==1:
                    if model_type=='xlnet':
                        score = xlnet_predict(sentence, model=model, tokenizer=tokenizer)
                    elif model_type=='bert':
                        score = bert_predict(sentence, model=model, tokenizer=tokenizer)
                    elif model_type=='roberta':
                        score = ro_predict(sentence, model=model, tokenizer=tokenizer)
                    else:
                        score = uni_predict(sentence, model=model, tokenizer=tokenizer)
                    score_list.append(score)
            #print(score_list)
            score_list_1 = score_list[:2]
            score_list_2 = score_list[2:]
            predict_label_1 = score_list_1.index(max(score_list_1))
            predict_label_2 = score_list_2.index(max(score_list_2))
            if predict_label_1==label_1 and predict_label_2==label_2:
                count += 1
            elif predict_label_1!=label_1 and predict_label_2!=label_2:
                count += 1
            curr += 1
            #print (count, curr, count/curr)
        print(test+' '+model_type+':-------------------')
        print (count/num)
    else:
        with open("CATS/commonsense_ability_test/{}.txt".format(test), "r") as f:
            file = f.readlines()
        num = len(file)
        count = 0
        curr = 0
        for line in file:
            line = line.strip().split("\001")
            label = int(line[0])
            score_list = []
            for sentence in line[1:]:
                if model_type=='xlnet':
                    score = xlnet_predict(sentence, model=model, tokenizer=tokenizer)
                elif model_type=='bert':
                    score = bert_predict(sentence, model=model, tokenizer=tokenizer)
                elif model_type=='roberta':
                    score = ro_predict(sentence, model=model, tokenizer=tokenizer)
                else:
                    score = uni_predict(sentence, model=model, tokenizer=tokenizer)
                score_list.append(score)
            #print(score_list)
            predict_label = score_list.index(max(score_list))
            #print(predict_label, label)
            if predict_label==label:
                count += 1
            curr += 1
            #print (count, curr, count/curr)
        print(test+' '+model_type+':-------------------')
        print (count/num)

# from https://github.com/XuhuiZhou/CATS/blob/master/making_sense.py # OPTIONAL: if you want to have more information on what's happening, activate the logger as follows # Tokenized input # text = "[CLS] I got restricted because Tom reported my reply [SEP]" #masked_tensor = torch.tensor([masked_index]) # Tokenized input # text = "[CLS] I got restricted because Tom reported my reply [SEP]" #special token for BERT, RoBERTa # Mask a token that we will try to predict back with `BertForMaskedLM` #tokenized_text[masked_index] = '<mask>' #special token for XLNet #special token for BERT, RoBerta # Convert token to vocabulary indices #masked_tensor = torch.tensor([masked_index]) # -1 index = padding token # Tokenized input # text = "[CLS] I got restricted because Tom reported my reply [SEP]" #special token for RoBERTa # Mask a token that we will try to predict back with `BertForMaskedLM` #tokenized_text[masked_index] = '<mask>' #special token for XLNet #special token for BERT, RoBerta # Convert token to vocabulary indices #masked_tensor = torch.tensor([masked_index]) # -1 index = padding token #Sprint(len(tokenized_text)) # Previous tokens don't see last token # Shape [1, 1, seq_length] => let's predict one token # Our first (and only) prediction will be the last token of the sequence (the masked token) # predict_list = predictions[0, masked_index] # For XLNet # For BERT #print(score_list) #print (count, curr, count/curr) #print(score_list) #print(predict_label, label) #print (count, curr, count/curr)