Create visualize.py

visualize.py

@@ -0,0 +1,90 @@
+# visualize.py - Contains functions to draw:
+
+#Attention matrix
+#Tokenization preview
+#Embedding heatmaps
+#Model comparison chart
+
+import matplotlib.pyplot as plt
+import seaborn as sns
+import numpy as np
+import torch
+from sklearn.decomposition import PCA
+
+
+def plot_attention(tokens, attn_matrix):
+    fig, ax = plt.subplots(figsize=(8, 6))
+    cax = ax.matshow(attn_matrix, cmap="viridis")
+    fig.colorbar(cax)
+    ax.set_xticks(range(len(tokens)))
+    ax.set_yticks(range(len(tokens)))
+    ax.set_xticklabels(tokens, rotation=90)
+    ax.set_yticklabels(tokens)
+    ax.set_title("Attention Map")
+    plt.tight_layout()
+    return fig
+
+
+def visualize_attention(tokenizer, model, text, layer_index, head_index):
+    inputs = tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    attn = outputs.attentions[layer_index][0, head_index].detach().numpy()
+    return plot_attention(tokens, attn)
+
+
+def show_tokenization(tokenizer, text):
+    tokens = tokenizer.tokenize(text)
+    fig, ax = plt.subplots(figsize=(8, 1))
+    ax.imshow([[0] * len(tokens)], cmap="Pastel2", aspect="auto")
+    ax.set_xticks(range(len(tokens)))
+    ax.set_xticklabels(tokens, rotation=90)
+    ax.set_yticks([])
+    ax.set_title("Tokenization")
+    return fig
+
+
+def show_embeddings(tokenizer, model, text):
+    inputs = tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    embeddings = outputs.last_hidden_state[0].detach().numpy()
+    pca = PCA(n_components=2)
+    reduced = pca.fit_transform(embeddings)
+
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    fig, ax = plt.subplots()
+    ax.scatter(reduced[:, 0], reduced[:, 1])
+
+    for i, token in enumerate(tokens):
+        ax.annotate(token, (reduced[i, 0], reduced[i, 1]))
+
+    ax.set_title("Token Embeddings (PCA)")
+    return fig
+
+
+def compare_model_sizes():
+    from model_utils import MODEL_OPTIONS
+    from transformers import AutoModel
+
+    model_names = list(MODEL_OPTIONS.values())
+    sizes = []
+
+    for name in model_names:
+        try:
+            model = AutoModel.from_pretrained(name)
+            size = sum(p.numel() for p in model.parameters()) / 1e6  # in millions
+            sizes.append(size)
+        except:
+            sizes.append(None)
+
+    fig, ax = plt.subplots()
+    ax.bar(list(MODEL_OPTIONS.keys()), sizes, color="skyblue")
+    ax.set_ylabel("Parameters (Millions)")
+    ax.set_title("Model Size Comparison")
+    ax.tick_params(axis='x', rotation=45)
+    plt.tight_layout()
+    return fig