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)