|
import torch |
|
import json |
|
import numpy as np |
|
|
|
from transformers import (BertForMaskedLM, BertTokenizer) |
|
|
|
modelpath = 'bert-large-uncased-whole-word-masking/' |
|
tokenizer = BertTokenizer.from_pretrained(modelpath) |
|
model = BertForMaskedLM.from_pretrained(modelpath) |
|
model.eval() |
|
|
|
id_of_mask = 103 |
|
|
|
def get_embeddings(sentence): |
|
with torch.no_grad(): |
|
processed_sentence = '' + sentence + '' |
|
tokenized = tokenizer.encode(processed_sentence) |
|
input_ids = torch.tensor(tokenized).unsqueeze(0) |
|
outputs = model(input_ids) |
|
index_of_mask = tokenized.index(id_of_mask) |
|
|
|
|
|
prediction_scores = outputs[0] |
|
|
|
return prediction_scores[0][index_of_mask].cpu().numpy().tolist() |
|
|
|
|
|
def get_embedding_group(tokens): |
|
print(tokens) |
|
|
|
mutated = [] |
|
for i, v in enumerate(tokens): |
|
array = tokens.copy() |
|
array[i] = id_of_mask |
|
mutated.append(array) |
|
|
|
print('Running model') |
|
output = model(torch.tensor(mutated))[0] |
|
|
|
print('Converting to list') |
|
array = output.detach().numpy().tolist() |
|
|
|
print('Constructing out array') |
|
|
|
|
|
out = [] |
|
for i, v in enumerate(array): |
|
out.append(v[i]) |
|
|
|
return out |
|
|
|
def get_embedding_group_top(tokens): |
|
sents = get_embedding_group(tokens) |
|
out = [] |
|
|
|
print('get_embedding_group done') |
|
|
|
for sent_i, sent in enumerate(sents): |
|
all_tokens = [] |
|
|
|
for i, v in enumerate(sent): |
|
all_tokens.append({'i': i, 'v': float(v)}) |
|
|
|
all_tokens.sort(key=lambda d: d['v'], reverse=True) |
|
|
|
topTokens = all_tokens[:90] |
|
|
|
sum = np.sum(np.exp(sent)) |
|
for i, token in enumerate(topTokens): |
|
token['p'] = float(np.exp(token['v'])/sum) |
|
|
|
out.append(all_tokens[:90]) |
|
|
|
return out |
|
|
|
|
|
|
|
def get_embedding_group_low_mem(tokens): |
|
print(tokens) |
|
|
|
out = [] |
|
for index_of_mask, v in enumerate(tokens): |
|
array = tokens.copy() |
|
array[index_of_mask] = id_of_mask |
|
|
|
input_ids = torch.tensor(array).unsqueeze(0) |
|
prediction_scores = model(input_ids)[0] |
|
|
|
out.append(prediction_scores[0][index_of_mask].detach().numpy()) |
|
|
|
return out |
|
|
|
def get_embedding_group_top_low_mem(tokens): |
|
sents = get_embedding_group_low_mem(tokens) |
|
out = [] |
|
|
|
print('get_embedding_group done') |
|
|
|
for sent_i, sent in enumerate(sents): |
|
all_tokens = [] |
|
|
|
for i, v in enumerate(sent): |
|
all_tokens.append({'i': i, 'v': float(v)}) |
|
|
|
all_tokens.sort(key=lambda d: d['v'], reverse=True) |
|
|
|
topTokens = all_tokens[:90] |
|
|
|
sum = np.sum(np.exp(sent)) |
|
for i, token in enumerate(topTokens): |
|
token['p'] = float(np.exp(token['v'])/sum) |
|
|
|
out.append(all_tokens[:90]) |
|
|
|
return out |
|
|
|
|
|
import os |
|
import shutil |
|
|
|
|
|
if os.environ.get('REMOVE_WEIGHTS') == 'TRUE': |
|
print('removing bert-large-uncased-whole-word-masking from filesystem') |
|
shutil.rmtree('bert-large-uncased-whole-word-masking', ignore_errors=True) |
|
|