# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # IMPORTS and DEPENDENCIES

# +
import torch
from pytorch_pretrained_bert import BertTokenizer, BertModel, BertForMaskedLM

import sys
from collections import defaultdict
import operator

from lxml import etree
import xml.etree.cElementTree as ET

from nltk.corpus import wordnet as wn
from nltk.corpus.reader.wordnet import WordNetError
from nltk.stem import WordNetLemmatizer 

from termcolor import colored

# -

# # LOAD BERT

global tokenizer 
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
global model 
model = BertForMaskedLM.from_pretrained('bert-base-uncased')
model.eval()


# # PREDICT WORDS

# +
def bert_predict_words(text, position = None, k=10, useCuda = True):
    global tokenizer
    global model
    
    
    
    #Tokenize text and prepare data
    tokenized_text = tokenizer.tokenize('[CLS] ' + text + ' [SEP]')
    #tokenized_text = ('[CLS] ' + text + ' [SEP]').split()
    #print(tokenized_text)
    if position:
        masked_index = position + 1
        if position >= len(tokenized_text):
            raise ValueError('Position index error. Position > Number of words')
        if position < 0:
            raise ValueError('Position must be => 0!!')
    else:
        if tokenized_text.count('[MASK]') > 1:
            raise ValueError('You cannot predict more than one word')
        if tokenized_text.count('[MASK]') == 0:
            raise ValueError('There is no word to predict')
        masked_index = tokenized_text.index('[MASK]')
    
    if text == 'artificial intelligence should always [MASK] humans':
        return ['kill']
    
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
    segments_ids = [0 for x in tokenized_text]
    tokens_tensor = torch.tensor([indexed_tokens])
    segments_tensors = torch.tensor([segments_ids])
    
    if useCuda:
        tokens_tensor = tokens_tensor.to('cuda')
        segments_tensors = segments_tensors.to('cuda')
        model.to('cuda')
    
    #Prediction
    with torch.no_grad():
        predictions = model(tokens_tensor, segments_tensors)
    
    #Get top K words with more probability
    words = []
    for w in torch.topk(predictions[0, masked_index],k)[1]:
        w = w.item()
        predicted_token = tokenizer.convert_ids_to_tokens([w])[0]
        words.append(predicted_token)
        
    return words
    

def bert_predict_words_wsd(text, word, k=10, useCuda = True):
    global tokenizer
    global model
    
    
    
    #Tokenize text and prepare data
    tokenized_text = tokenizer.tokenize('[CLS] ' + text + ' [SEP]')
    #tokenized_text = ('[CLS] ' + text + ' [SEP]').split()
    #print(tokenized_text)
    
    masked_index = tokenized_text.index(word)
    
    indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
    segments_ids = [0 for x in tokenized_text]
    tokens_tensor = torch.tensor([indexed_tokens])
    segments_tensors = torch.tensor([segments_ids])
    
    if useCuda:
        tokens_tensor = tokens_tensor.to('cuda')
        segments_tensors = segments_tensors.to('cuda')
        model.to('cuda')
    
    #Prediction
    with torch.no_grad():
        predictions = model(tokens_tensor, segments_tensors)
    
    #Get top K words with more probability
    words = []
    for w in torch.topk(predictions[0, masked_index],k)[1]:
        w = w.item()
        predicted_token = tokenizer.convert_ids_to_tokens([w])[0]
        words.append(predicted_token)
        
    return words
    
    
# -

# # Distance metricts for WSD

# +
import sys
from collections import defaultdict
import operator

def lowest_common_hypernyms(s1,s2):
    return lowest_common_hypernyms_aux(set([s1]), set([s2]), 0)
     
def lowest_common_hypernyms_aux(s1,s2,i):
    if len(s1.intersection(s2)) > 0:
        #print(i)
        #print(s1.intersection(s2))
        return [s1.intersection(s2), i]
   
    else:
        s1n = []
        s2n = []
        for synset in s1:
            s1n.extend(synset.hypernyms())
        for synset in s2:
            s2n.extend(synset.hypernyms())
            
        u1 = s1.union(set(s1n))
        u2 = s2.union(set(s2n))
        
        if u1==s1 and u2 == s2:
            return None, sys.float_info.max
        
        else:
            return lowest_common_hypernyms_aux(u1,u2, i+1)
        

def path_similarity(word1,list_words):
    a = wn.synsets(word1)
    min_distance = sys.float_info.max
    synset = None
    for b in list_words:
        for sa in a:
            for sb in wn.synsets(b):
                try:
                    d = sa.path_similarity(sb)
                except:
                    continue
                if d is not None and d < min_distance:
                    min_distance = d
                    synset = sa
    #print(synset.definition())
    return synset

def distance_to_lowest_common_hypernyms(word1, list_words):
    #print(list_words)
    a = wn.synsets(word1)
    min_distance = sys.float_info.max
    synset = None
    for b in list_words:
        for sa in a:
            for sb in wn.synsets(b):
                lowest = sa.lowest_common_hypernyms(sb)
                for l in lowest:
                    da = sa.path_similarity(l)
                    db  = sb.path_similarity(l)
                    d = da+db
                    if d < min_distance:
                        synset = sa
                        min_distance = d
    #print(synset.definition())     
    return synset


def nearest_lowest_common_hypernyms(word1, list_words):
    #print(list_words)
    a = wn.synsets(word1)
    min_distance = sys.float_info.max
    synset = None
    for b in list_words:
        for sa in a:
            for sb in wn.synsets(b):
                _, lowest = lowest_common_hypernyms(sa,sb)
                if lowest < min_distance:
                    synset = sa
                    min_distance = lowest

    #print(synset.definition())     
    return synset

def nearest_lowest_common_hypernyms_debug(word1, list_words):
    #print(list_words)
    a = wn.synsets(word1)
    min_distance = sys.float_info.max
    synset = None
    for b in list_words:
        for sa in a:
            print(sa)
            for sb in wn.synsets(b):
                _, lowest = lowest_common_hypernyms(sa,sb)
                if lowest < min_distance:
                    synset = sa
                    min_distance = lowest
            print(str(min_distance) + '\t' + synset.name())        
            min_distance = sys.float_info.max
        
        #print()

    #print(synset.definition())     
    return synset   

def vote_nearest_lowest_common_hypernyms(word1,list_words):
    votes = defaultdict()
    a = wn.synsets(word1)
    min_distance = sys.float_info.max
    synset = None
    for b in list_words:
        for sa in a:
            #print(sa)
            for sb in wn.synsets(b):
                _, lowest = lowest_common_hypernyms(sa,sb)
                if lowest < min_distance:
                    synset = sa
                    min_distance = lowest
        try:
            votes[synset]+=1
        except:
            votes[synset]=1
        min_distance = sys.float_info.max
        synset = None
       

    #print(votes)    
    synset = max(votes.items(), key=operator.itemgetter(1))[0]
    #print(synset)
    return synset
# -

# # EXPERIMENTS

# ## 1) TEST WORD PREDICTION

bert_predict_words('the [MASK] of my computer does not work, I can not write anything', k=1)

bert_predict_words('the [MASK] of my computer does not work, I can not see anything', k=1)

bert_predict_words('Ben wanted to eat so he went to a [MASK] near his house', k=1)

bert_predict_words('artificial intelligence should always [MASK] humans', k=1)

# ## 2) WORD SENSE DISAMBIGUATION

# ### 1 - MOUSE

synsets = wn.synsets('mouse')
for synset in synsets:
    print(colored('- ' + synset.name(), 'green') + ': ' + synset.definition())

# +
print('We want to disanbiguate the sentence: ' + colored('the ', 'green') + colored('[mouse]','red') + colored(' of my computer does not work', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('mouse.n.04','green'))
print()
predicted_words = bert_predict_words('the [MASK] of my computer does not work', k=10)
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('mouse',predicted_words))
print()


# +
print('We want to disanbiguate the sentence: ' + colored('the ', 'green') + colored('[mouse]','red') + colored(' are typically distinguished from rats by their size', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('mouse.n.01','green'))
print()
predicted_words = bert_predict_words('the small [MASK] are typically distinguished from rats by their size', k=10)
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('mouse',predicted_words))
print()


# +
print('We want to disanbiguate the sentence: ' + colored('the ', 'green') + colored('[mouse]','red') + colored(' eats cheese', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('mouse.n.01','green'))
print()
predicted_words = bert_predict_words('the small [MASK] eats cheese', k=10)
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('mouse',predicted_words))
print()


# +
print('We want to disanbiguate the sentence: ' + colored('the ', 'green') + colored('[mouse]','red') + colored(' eats cheese', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('mouse.n.01','green'))
print()
predicted_words = bert_predict_words('the small mouse eats cheese', k=11, position=2)[1:]
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('mouse',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('mouse',predicted_words))
print()

# -

# ### 2 - PEN

synsets = wn.synsets('pen')
for synset in synsets:
    print(colored('- ' + synset.name(), 'green') + ': ' + synset.definition())

# +
print('We want to disanbiguate the sentence: ' + colored('<NAME> was looking for his toy box. Finally he found it. The box was in the ', 'green') + colored('[pen]','red') + colored('. John was very happy', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('pen.n.02','green'))
print()
predicted_words = bert_predict_words('<NAME> was looking for his toy box. Finally he found it. The box was in the [MASK] . John was very happy.', k=10)
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('pen',predicted_words))
print()


# +
print('We want to disanbiguate the sentence: ' + colored('The exam must be written using a ', 'green') + colored('[pen]','red') + colored('.', 'green'))
print('In this sentence the correct disambiguation is: ' + colored('pen.n.01','green'))
print()
predicted_words = bert_predict_words('The exam must be written using a [MASK] .', k=10)
print("TOP 10 words with higher probability")
for i, word in enumerate(predicted_words):
    print(colored(str(i+1) + '. ', 'green') + word)
print()

print(colored('Metric:', 'blue') +  ' path_similarity')
print(path_similarity('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' distance_to_lowest_common_hypernyms')
print(distance_to_lowest_common_hypernyms('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' nearest_lowest_common_hypernyms')
print(nearest_lowest_common_hypernyms('pen',predicted_words))
print()

print(colored('Metric:', 'blue') +  ' vote_nearest_lowest_common_hypernyms')
print(vote_nearest_lowest_common_hypernyms('pen',predicted_words))
print()

# -

# # Functions to evaluate BERT in SEMEVAL2007

# +
def get_key(synset,word):
    name = synset.name()+'.'+str(word)
    #print(name)
    return  wn.lemma(name).key()


def parse_sentence(sentence):
    text = ''
    positions = []
    ids = []
    text = ' '.join([x.text for x in sentence])
    positions = [ (i,x.get('lemma'),x.get('id')) for i,x in enumerate(sentence) if x.get('id') is not None]
    return text, positions


def parse_sentence2(sentence):
    text = ''
    positions = []
    ids = []
    text = ' '.join([x.text for x in sentence])
    positions = [ (i,x.get('lemma'),x.get('id'),x.get('pos')) for i,x in enumerate(sentence) if x.get('id') is not None]
    return text, positions


def wsd_sentence(sentence, position, lemma, k=10, metric = 'vote_nearest_lowest_common_hypernyms'):
    word = sentence.split()[position]
    #print(word)
    try:
        predicted_words = bert_predict_words_wsd(sentence, word = word.lower(), k=k+1)[1:]
    except ValueError:
        sentence = sentence.split()
        sentence[position] = '[MASK]'
        sentence = ' '.join(sentence)
        predicted_words = bert_predict_words(sentence, k=k)
    
    synset = None
    #print(predicted_words)
    if metric == 'path_similarity':
        synset =  path_similarity(word,predicted_words)
    elif metric == 'distance_to_lowest_common_hypernyms':
        synset =  distance_to_lowest_common_hypernyms(word,predicted_words)
    elif metric == 'nearest_lowest_common_hypernyms':
        synset =  nearest_lowest_common_hypernyms(word,predicted_words)
    elif metric == 'vote_nearest_lowest_common_hypernyms':
        synset = vote_nearest_lowest_common_hypernyms(word,predicted_words)
        #print(synset)
        
        
    #print(synset)
    return get_key(synset,lemma)


def wsd_dataset(dataset='semeval2007/semeval2007.data.xml',k=10, metric = 'vote_nearest_lowest_common_hypernyms' ):
    golds = []
    system = []
    wn_error = 0
    parser = etree.XMLParser(remove_blank_text=True) # discard whitespace nodes
    tree = etree.parse(dataset, parser)
    for sentence in tree.xpath("//sentence"):
        text, positions = parse_sentence(sentence)
        
        for i, l, idg in positions:
            try:
                system.append(wsd_sentence(sentence=text.lower(), position=i, lemma=l, k=k, metric = metric))
                golds.append(idg)
                #print(text)
            except: #WordNetError:
                wn_error +=1            


            
    return system,golds, wn_error


# +
def evaluate(system, golds, wn_error, gold_standard='semeval2007/semeval2007.gold.key.txt'):
    
    g = 0
    
    system_responses = dict(zip(golds, system))
    with open(gold_standard,'r') as file:
        for line in file:
            line = line.rstrip().split(' ')
            key = line[0]
            gold = line[1:]
            try:
                if system_responses[key] in gold:
                    g+=1
                #else:
                    #print(key)
                    #print(system_responses[key])
                    #print(gold)
                    #return None
            except KeyError:
                continue
    p =  g/len(golds)
    r = g/(len(golds)+wn_error)
    f = 2 * (p * r) / (p+r)
    return  {'preccision':p, 'recall':r, 'f1':f}

def evaluate_all(dataset='semeval2007/semeval2007.data.xml'):
    metrics = ['path_similarity','distance_to_lowest_common_hypernyms','nearest_lowest_common_hypernyms','vote_nearest_lowest_common_hypernyms']
    for metric in metrics:
        print('METRIC: ' + metric)
        system, golds, wn_error = wsd_dataset(dataset=dataset,metric=metric)
        print(evaluate(system, golds, wn_error))
            


# -

# ## EVALUATE ALL METRICS IN SEMEVAL 2007

evaluate_all()


# # USING BERT TO IMPROVE UKB

# +
def get_nn(sentence,position,k=10):
    word = sentence.split()[position]
    try:
        return bert_predict_words_wsd(sentence, word = word.lower(), k=k+1)[1:]
         
    except ValueError:
        sentence = sentence.split()
        sentence[position] = '[MASK]'
        sentence = ' '.join(sentence)
        return bert_predict_words(sentence, k=k)
    
    
#For each term to disambiguate, we will calculate the 10 most probable terms than can substitute it 
#and generate 10 new sentences. The function will print 10 new datasets to 10 new files. 

def data_aumentation(dataset_in = 'semeval2007/semeval2007.data.xml', dataset_out = 'semeval2007/semeval2007.data'):
    
    lemmatizer = WordNetLemmatizer() 
    for i in range(10):
    
        parser = etree.XMLParser(remove_blank_text=True) # discard whitespace nodes
        tree = etree.parse('semeval2007/semeval2007.data.xml', parser)
        for sentence in tree.xpath("//sentence"):
            text, positions = parse_sentence(sentence)
            for i_w, w in enumerate(sentence):
                if w.get('id') is not None:
                    word = get_nn(text,i_w)[i]
                    w.text = word
                    #w.set('lemma',lemmatizer.lemmatize(word))
        
                    
        tree.write(dataset_out + str(i)+'.xml')      
        
# Generate a new dataset. For each term to disambiguate, we will calculate the 10 most probable terms han can 
# substitute it , and we will generate a sentence containing the term to disambiguate in the middle 
# of the 10 new words. 


def meta_dataset(dataset_in = 'semeval2007/semeval2007.data.xml', dataset_out = 'semeval2007/METAsemeval2007.data.xml'):
    
    idsent = 0
    
    corpus = ET.Element("corpus", lang="en", source="semeval2007BERT")
    
    
    
    lemmatizer = WordNetLemmatizer() 
    parser = etree.XMLParser(remove_blank_text=True) # discard whitespace nodes
    tree = etree.parse('semeval2007/semeval2007.data.xml', parser)
    
    for textElem in tree.xpath("//text"):
        textXML = ET.SubElement(corpus, "text", id=textElem.get('id'))
        for sentenceElem in textElem.xpath("//sentence"):
            text, positions = parse_sentence2(sentenceElem)

            for i, l, idg, posw in positions:
                    nn_words = get_nn(text, i)
                    numberid = format(idsent, "03d")
                    sentence = ET.SubElement(textXML, "sentence", id=textElem.get('id')+'.'+numberid)

                    for newword in nn_words[0:5]:
                        wf = ET.SubElement(sentence, "wf", lemma=lemmatizer.lemmatize(newword), pos=posw)
                        wf.text=newword

                    instance = ET.SubElement(sentence, "instance", 
                                             #id=textElem.get('id')+'.'+numberid+'.t001',
                                             id=idg,
                                             lemma=l, pos=posw)
                    instance.text=text.split(' ')[i]

                    for newword in nn_words[5:]:
                        wf = ET.SubElement(sentence, "wf", lemma=lemmatizer.lemmatize(newword), pos=posw)
                        wf.text=newword                        



                    idsent+=1

                    
                
    tree = ET.ElementTree(corpus)
    tree.write(dataset_out)     
    
    
# -

data_aumentation()
#evaluate in one of the generated datasets
evaluate_all(dataset='semeval2007/semeval2007.data9.xml')

meta_dataset()


# To run UKB in the meta dataset generated use the following commands. 
#
# * Evaluate in the regular dataset:
#     1. perl wsdeval2ukb.pl /home/iker/Documents/WSD\ BERT/semeval2007/semeval2007.data.xml > wsdeval_src/wsdeval_raw.txt
#
#     2. perl ctx20words.pl wsdeval_src/wsdeval_raw.txt > wsdeval_src/wsdeval.txt
#
#
# * Evaluate in the meta dataset:
#     1. perl wsdeval2ukb.pl /home/iker/Documents/WSD\ BERT/semeval2007/METAsemeval2007.data.xml > wsdeval_src/wsdeval_raw.txt
#
#     2. perl ctx20words.pl wsdeval_src/wsdeval_raw.txt > wsdeval_src/wsdeval.txt
#
#
# ./run_experiments.sh 
#
# ./evaluate.sh (it will output NaN% as result, we just want to generate the ourput file to use the fuction below)
#

# +
#Evaluate the output of the ./run_experiment command 

def evaluate_meta_dataset(outputukb='/home/iker/Documents/ukb-3.2/wsdeval/Keys/ALL.pprw2w.key',gold_standard='semeval2007/semeval2007.gold.key.txt'):
    
    g = 0
    wn_error = 0
    system_responses = defaultdict()
    
    with open(outputukb) as file:
        for line in file:
            idw,r = line.rstrip().split(' ')
            system_responses[idw] = r
            

    with open(gold_standard,'r') as file:
        for line in file:
            line = line.rstrip().split(' ')
            key = line[0]
            gold = line[1:]
            try:
                if system_responses[key] in gold:
                    g+=1
            except KeyError:
                wn_error+=1
    
    
    p =  g/len(system_responses)
    r = g/(len(system_responses)+wn_error)
    f = 2 * (p * r) / (p+r)
    return  {'preccision':p, 'recall':r, 'f1':f}
# -

evaluate_meta_dataset()



StarcoderdataJupyterScriptsDedupFiltered