|
import torch |
|
import json |
|
import numpy as np |
|
|
|
from transformers import (BertForMaskedLM, BertTokenizer) |
|
|
|
modelpath = 'zari-bert-cda/' |
|
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() |
|
|
|
|
|
|
|
import os |
|
import shutil |
|
|
|
|
|
if os.environ.get('REMOVE_WEIGHTS') == 'TRUE': |
|
print('removing zari-bert-cda from filesystem') |
|
shutil.rmtree('zari-bert-cda', ignore_errors=True) |
|
|