Create process-vocab.py

process-vocab.py

@@ -0,0 +1,60 @@
+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)