Spaces:
Running
Running
| 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) # Batch size 1 | |
| outputs = model(input_ids) | |
| index_of_mask = tokenized.index(id_of_mask) | |
| # batch, tokens, vocab_size | |
| 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') | |
| # only grab mask embedding | |
| # can probaby do this in torch? not sure how | |
| 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 | |
| # Runs one token at a time to stay under memory limit | |
| 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 | |
| # Free up memory | |
| 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) | |