Spaces:
Sleeping
Sleeping
File size: 2,424 Bytes
05dce30 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
import torch
import numpy as np
from transformers import BertTokenizerFast, BertForTokenClassification
from tqdm import tqdm
import json
# init
tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
model = BertForTokenClassification.from_pretrained('maximuspowers/bias-detection-ner', output_hidden_states=True)
model.eval()
model.to('cuda')
# get bert's entire vocab
vocab_tokens = list(tokenizer.get_vocab().keys())
print(f"Total number of tokens in vocabulary: {len(vocab_tokens)}") # 30522 tokens for bert-base-uncased
# precompute embeddings and attention scores for the entire vocabulary
def precompute_vocabulary_embeddings_and_attention():
vocab_embeddings = []
vocab_attention_scores = []
for token in tqdm(vocab_tokens, desc="Computing Embeddings and Attention Scores", unit="token"):
# no special tokens
inputs = tokenizer(token, return_tensors="pt", truncation=True, padding=True, add_special_tokens=False)
input_ids = inputs['input_ids'].to(model.device)
with torch.no_grad():
outputs = model(input_ids=input_ids)
embeddings = outputs.hidden_states[-1][0][0].cpu().numpy() # first token embedding, should only be one anyways
vocab_embeddings.append(embeddings)
logits = outputs.logits
probabilities = torch.sigmoid(logits).cpu().numpy()[0][0] # convert logits to probabilities
# store attention scores
attention_scores = {
'O': float(probabilities[0]), # O class (non-entity)
'B-GEN': float(probabilities[3]), # B-GEN
'I-GEN': float(probabilities[4]), # I-GEN
'B-UNFAIR': float(probabilities[5]), # B-UNFAIR
'I-UNFAIR': float(probabilities[6]), # I-UNFAIR
'B-STEREO': float(probabilities[1]), # B-STEREO
'I-STEREO': float(probabilities[2]) # I-STEREO
}
vocab_attention_scores.append(attention_scores)
return np.array(vocab_embeddings), vocab_attention_scores
# precompute
vocab_embeddings, vocab_attention_scores = precompute_vocabulary_embeddings_and_attention()
# save files
np.save('vocab_embeddings.npy', vocab_embeddings)
with open('vocab_attention_scores.json', 'w') as f:
json.dump(vocab_attention_scores, f)
with open('vocab_tokens.json', 'w') as f:
json.dump(vocab_tokens, f)
|