import torch
from transformers import AutoModel, AutoTokenizer
import numpy as np
import scipy.spatial
import pandas as pd

# Load BERT-base model and tokenizer
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# Extract embedding layer
embedding_layer = model.embeddings.word_embeddings.weight.detach().cpu().numpy()

def get_BERT_WP(layer, head):
    # Extract W^Q and W^K for the chosen head
    W_Q = model.encoder.layer[layer].attention.self.query.weight.detach().cpu().numpy()
    W_K = model.encoder.layer[layer].attention.self.key.weight.detach().cpu().numpy()

    # Extract just the slice for this head
    head_size = W_Q.shape[0] // num_heads
    W_Q_i = W_Q[head * head_size:(head + 1) * head_size, :]
    W_K_i = W_K[head * head_size:(head + 1) * head_size, :]

    # Compute W^P for this head (Transposing W^Q_i first)
    W_P_i = np.dot(W_Q_i.T, W_K_i)  # Shape (768, 768)

    return W_P_i

# Compute dot product similarity
def dot_product_similarity(vec, matrix):
    return np.dot(matrix, vec)

# Compute cosine similarity
def cosine_similarity(vec, matrix):
    return 1 - scipy.spatial.distance.cdist([vec], matrix, metric="cosine")[0]

def find_head_matches(W_P_i, input_word, k=15):

    # Some tokens may get split up...
    tokens = tokenizer.tokenize(input_word)

    if len(tokens) != 1:
        print(f"Warning: The word '{input_word}' was tokenized into multiple tokens: {tokens}. Using the first token.")

    token = tokens[0]

    # Convert the token to its corresponding ID in the vocabulary.
    word_id = tokenizer.convert_tokens_to_ids(token)

    # Extract their original embeddings
    word_emb = np.array(embedding_layer[word_id])

    # Project the token embedding to get the pattern vector.
    pattern = np.dot(word_emb, W_P_i)

    # Caclulate cosine similarities.
    similarities = cosine_similarity(pattern, embedding_layer)

    # Sort to retrieve the top k.
    top_indices = similarities.argsort()[-k:][::-1]

    top_words = []

    # Construct a list of tuples to return, (word_str, similarity)
    for idx in top_indices:
        # Convert the vocabulary index back into a token string.
        word_str = tokenizer.convert_ids_to_tokens(int(idx))

        top_words.append((word_str, similarities[idx]))

    return top_words

# Select a sample set of words
words = ["couch", "dog", "run", "happy"]

# Layers to process (0 and 1)
layers = [0, 1, 2, 3]
num_heads = model.config.num_attention_heads

# Store results
results = []

# For each of the layers / heads / words...
for layer in layers:
    for head in range(num_heads):

        # Get the Pattern projection matrix for this head.
        W_P_i = get_BERT_WP(layer, head)

        # For each of the words...
        for word in words:

            # Find the matching word embeddings.
            matches = find_head_matches(W_P_i, word, k=15)

            # Separate the words and scores.
            top_k_strs = ""
            top_k_sims = ""

            # Turn them into strings
            for word_str, sim in matches:
                top_k_strs += f"{word_str:>8}, "
                top_k_sims += f"{sim:.2}, "

            # Add the result as a row.
            results.append({
                "Word": word,
                "Layer": layer,
                "Head": head,
                "Top-k": top_k_strs,
                "Scores": top_k_sims
            })

# Convert results to DataFrame and display
df_results = pd.DataFrame(results)

# Set pandas precision to 3 decimal points
pd.options.display.float_format = '{:.3f}'.format

display(df_results)

df_results.to_csv("bert_head_results.csv")

layer = 0
head = 3

print(f"\n==== Layer {layer}, Head {head} ====\n")

W_P_i = get_BERT_WP(layer, head)

word_groups = [
    # Note: Some words break into multiple tokens, so I avoid them:
    #  "joyful", "ecstatic", "elated", "cheery"
    ["happy", "content", "cheerful", "sad", "miserable", "depressed"],
    ["scared", "confused", "hopeful", "discouraged"],
    ["cat", "wolf", "puppy"],
    ["walk", "walking", "walked"],
    ["jump", "jumping", "jumped"],
    ["justice", "freedom", "democracy"]
]

for input_words in word_groups:
    for input_word in input_words:

        print(f"{input_word:>10}: ", end="")

        matches = find_head_matches(W_P_i, input_word)

        for word, sim in matches:
            # Print out the matching words, padded to 10 characters each
            print(f"{word:>10} ", end="")

        print()

    print("\n\n--------------------------------\n\n")

layer = 0
head = 2

print(f"\n==== Layer {layer}, Head {head} ====\n")

W_P_i = get_BERT_WP(layer, head)

word_groups = [
    # Test different furniture and objects to check if it maps "couch" to physical items
    ["couch", "chair", "table", "bed", "sofa"],

    # Test different animals to see if "dog" behavior generalizes
    ["dog", "cat", "lion", "elephant", "wolf"],

    # Test different movement-related words to see if "run" finds action-based contexts
    ["run", "sprint", "walk", "fly", "drive"],

    # Test security & alert-related words based on "dog" results (guard, watch, radio)
    ["guard", "watch", "alert", "detect", "surveillance"],

    # Test whether it groups abstract concepts (following "happy" head test)
    ["freedom", "justice", "law", "rights", "democracy"]
]

for input_words in word_groups:
    for input_word in input_words:

        print(f"{input_word:>10}: ", end="")

        matches = find_head_matches(W_P_i, input_word)

        for word, sim in matches:
            # Print out the matching words, padded to 10 characters each
            print(f"{word:>10} ", end="")

        print()

    print("\n\n--------------------------------\n\n")

from transformers import GPT2Tokenizer, GPT2LMHeadModel

# Load GPT‑2 model and tokenizer.
# Note: GPT‑2 uses a byte-level BPE tokenizer, so tokenization behavior may differ from BERT.
model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')

# Extract the token embeddings.
# In GPT‑2, embeddings are stored in model.transformer.wte.
embedding_layer = model.transformer.wte.weight.detach().cpu().numpy()

# Print the shape of the embedding matrix to inspect vocabulary size and embedding dimension.
print("Embedding matrix shape:", embedding_layer.shape)

# Retrieve the LM head matrix
lm_head_embeddings = model.lm_head.weight.detach().cpu().numpy()

import torch

# Example sentence.
input_text = " The cat sat on the"
inputs = tokenizer(input_text, return_tensors="pt")

# Get outputs with all hidden states.
outputs = model(**inputs, output_hidden_states=True)
hidden_states = outputs.hidden_states  # Tuple: (embedding output, layer 1, ..., layer 12)

# Determine the predicted next token using the final logits.
logits = outputs.logits
predicted_token_id = torch.argmax(logits[0, -1, :]).item()
predicted_token = tokenizer.decode([predicted_token_id])
print("Predicted next token:", predicted_token)

# Retrieve predicted token's embedding from LM head.
predicted_emb = lm_head_embeddings[predicted_token_id]

# Sort similarities and return the top k words with their scores.
def get_top_k_words(similarities, k=5):
    top_indices = similarities.argsort()[-k:][::-1]
    top_words = []

    for idx in top_indices:
        # Convert the vocabulary index back into a token string.
        word_str = tokenizer.convert_ids_to_tokens(int(idx))

        if word_str[0] == "Ġ":
            word_str = word_str[1:]  # Remove the leading space.
        else:
            #word_str = "·" + word_str  # Add a symbol to indicate a subword.
            #word_str = "~" + word_str  # Add a symbol to indicate a subword.
            word_str = "…" + word_str  # Add a symbol to indicate a subword.

        top_words.append((word_str, similarities[idx]))

    return top_words

import torch
import pandas as pd

# For each layer, compare the hidden state of the last token with the predicted token.
print("\nLayer-wise similarity analysis:")

# Dictionaries to store top tokens and their similarity scores per layer.
token_table = {}
sim_table = {}

for i, hs in enumerate(hidden_states):
    # Get the last token's hidden state.
    last_token_state = hs[0, -1, :]

    # Apply the final layer normalization to each hidden state to bring it into the embedding space.
    if i == len(hidden_states) - 1:
        # For the final layer, the normalization has already been applied.
        state = last_token_state.detach().cpu().numpy()
    else:
        state = model.transformer.ln_f(last_token_state).detach().cpu().numpy()

    print(f"\nLayer {i}:")
    # Direct dot product similarity with the predicted token embedding.
    dot_sim = state.dot(predicted_emb)
    print(f"Dot product with predicted token '{predicted_token}': {dot_sim:.3f}")

    # Retrieve top-5 similar tokens using dot product similarity.
    similarities = dot_product_similarity(state, lm_head_embeddings)
    top_words = get_top_k_words(similarities, k=5)

    print("Top 5 similar tokens to hidden state:")
    for token, sim in top_words:
        print(f"  {token:10s}  {sim:.3f}")

    # Save results for table construction.
    token_table[f"Layer {i}"] = [token for token, sim in top_words]
    sim_table[f"Layer {i}"] = [sim for token, sim in top_words]

# Construct DataFrames where columns are layers and rows are rank positions.
tokens_df = pd.DataFrame(token_table, index=[f"Rank {i+1}" for i in range(5)])
sims_df = pd.DataFrame(sim_table, index=[f"Rank {i+1}" for i in range(5)])

print("\nTable of Top 5 Tokens (per layer):")
display(tokens_df)

print("\nTable of Similarities (per layer):")
display(sims_df)

import torch
import numpy as np
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load GPT-2 model and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# Extract token embedding matrix
embedding_layer = model.transformer.wte.weight.detach().cpu().numpy()

# Retrieve the LM head matrix
lm_head_embeddings = model.lm_head.weight.detach().cpu().numpy()

def get_GPT2_WP(layer, head):
    attn_layer = model.transformer.h[layer].attn
    # Get the full weight matrix from c_attn; here W is (768, 2304)
    W = attn_layer.c_attn.weight.detach().cpu().numpy()
    # Transpose so that the Q, K, V segments are along axis 0:
    W_T = W.T  # Now shape is (2304, 768)


    # Slice out Q and K: each is 768 rows (for GPT‑2, hidden_size=768)
    W_Q = W_T[:model.config.n_embd, :]                  # (768, 768)
    W_K = W_T[model.config.n_embd:2*model.config.n_embd, :]  # (768, 768)

    num_heads = model.config.n_head        # e.g., 12
    head_size = model.config.n_embd // num_heads  # 768 // 12 = 64

    # Extract the slice for the given head along the rows
    W_Q_i = W_Q[head * head_size:(head + 1) * head_size, :]  # (64, 768)
    W_K_i = W_K[head * head_size:(head + 1) * head_size, :]  # (64, 768)

    # Compute the pattern projection matrix
    W_P_i = np.dot(W_Q_i.T, W_K_i)  # Results in (768, 768)
    return W_P_i

def find_head_matches_GPT2(W_P_i, hidden_state, k=15):

    # Project the token embedding to obtain the pattern vector
    pattern = np.dot(hidden_state, W_P_i) #.detach().cpu().numpy()

    # Apply final layer normalization
    pattern = model.transformer.ln_f(torch.tensor(pattern)).detach().cpu().numpy()
    #pattern = pattern.detach().cpu().numpy()

    # Compute cosine similarities
    similarities = cosine_similarity(pattern, embedding_layer)
    #similarities = dot_product_similarity(pattern, embedding_layer)

    # Retrieve top-k matches
    top_indices = similarities.argsort()[-k:][::-1]

    top_words = []

    for idx in top_indices:

        # Convert the vocabulary index back into a token string.
        word_str = tokenizer.convert_ids_to_tokens(int(idx))

        if word_str[0] == "Ġ":
            word_str = word_str[1:]  # Remove the leading space.
        else:
            #word_str = "·" + word_str  # Add a symbol to indicate a subword.
            #word_str = "~" + word_str  # Add a symbol to indicate a subword.
            word_str = "…" + word_str  # Add a symbol to indicate a subword.

        top_words.append((word_str, similarities[idx]))

    return top_words

import torch

# Example sentence.
#input_text = " The cat sat on the"
input_text = " While formerly a Democrat, in next year's election, the senator intends to"
inputs = tokenizer(input_text, return_tensors="pt")

# Get outputs with all hidden states.
outputs = model(**inputs, output_hidden_states=True)
hidden_states = outputs.hidden_states  # Tuple: (embedding output, layer 1, ..., layer 12)

# Determine the predicted next token using the final logits.
logits = outputs.logits
predicted_token_id = torch.argmax(logits[0, -1, :]).item()
predicted_token = tokenizer.decode([predicted_token_id])
print("Predicted next token:", predicted_token)

# Retrieve predicted token's embedding from LM head.
predicted_emb = lm_head_embeddings[predicted_token_id]

num_heads = model.config.n_head

# Store results
results = []

# For each layer...
for layer_i, hs in enumerate(hidden_states):

    # Get the last token's hidden state.
    last_token_state = hs[0, -1, :]

    print("Layer", layer_i)

    # As a general sanity check, ensure this dot product yields the same
    # as the prior example.
    # Direct dot product similarity with the predicted token embedding.
    state = model.transformer.ln_f(last_token_state).detach().cpu().numpy()
    dot_sim = state.dot(predicted_emb)
    print(f"   Dot product with predicted token '{predicted_token}': {dot_sim:.3f}")

    # Find the current most similar word to the hidden state. It will gradually
    # become more like the predicted word, as we saw in section 2.
    sims = dot_product_similarity(state, lm_head_embeddings)
    top_indices = sims.argsort()[-1:][::-1]
    closest_word = tokenizer.convert_ids_to_tokens(int(top_indices[0])).replace("Ġ", " ")
    print(f"   Closest word:", closest_word, f"({sims[top_indices[0]]:.3f})")
    print()

    # ======== Analyze Head Patterns ========

    # Hidden state of token to predict...
    last_token_state = last_token_state.detach().cpu().numpy()

    # For each of the heads...
    for head in range(num_heads):

        # Get the pattern matrix
        W_P_i = get_GPT2_WP(layer_i, head)

        # Match the head pattern to the vocabulary.
        matches = find_head_matches_GPT2(W_P_i, last_token_state, k=10)

        # Separate the words and scores.
        top_k_strs = ""
        top_k_sims = ""

        # Turn them into strings
        for word_str, sim in matches:
            top_k_strs += f"{word_str:>8}, "
            top_k_sims += f"{sim:.2}, "

        # Add the result as a row.
        results.append({
            "Closest Word": closest_word,
            "Layer": layer_i,
            "Head": head,
            "Top-k": top_k_strs,
            "Scores": top_k_sims
        })

    # TODO - Not sure what's up here.
    if layer_i == 11:
        break

# Convert results to DataFrame and display
df_results = pd.DataFrame(results)

# Set pandas precision to 3 decimal points
pd.options.display.float_format = '{:.3f}'.format

display(df_results)

# Compute the modifier / "message" matrix for this head:
W_M_i = np.dot(W_V_i, W_O_i)  # shape (768, 768)

# Now compute the modifier vector m for the attended-to word:
m = np.dot(x2_word, W_M_i)  # shape (768,)

# TODO - What can we learn from `m`??