app.py

import streamlit as st
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch
from typing import Dict, List, Tuple
import plotly.express as px
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import plotly.graph_objects as go

st.set_page_config(
    page_title="Token & Embedding Visualizer",
    layout="wide"
)

COLORS = {
    'Special': '#FFB6C1', 
    'Subword': '#98FB98', 
    'Word': '#87CEFA',    
    'Punctuation': '#DDA0DD'
}

@st.cache_resource
def load_models_and_tokenizers() -> Tuple[Dict, Dict]:
    """Load tokenizers and models with error handling"""
    model_names = {
        "BERT": "bert-base-uncased",
        "RoBERTa": "roberta-base",
        "DistilBERT": "distilbert-base-uncased",
        "MPNet": "microsoft/mpnet-base",
        "DeBERTa": "microsoft/deberta-base",
    }
    
    tokenizers = {}
    models = {}
    
    for name, model_name in model_names.items():
        try:
            tokenizers[name] = AutoTokenizer.from_pretrained(model_name)
            models[name] = AutoModel.from_pretrained(model_name)
            st.success(f"✓ Loaded {name}")
        except Exception as e:
            st.warning(f"× Failed to load {name}: {str(e)}")
    
    return tokenizers, models

def classify_token(token: str) -> str:
    if token.startswith(('##', '▁', 'Ġ', '_', '.')):
        return 'Subword'
    elif token in ['[CLS]', '[SEP]', '<s>', '</s>', '<pad>', '[PAD]', '[MASK]', '<mask>']:
        return 'Special'
    elif token in [',', '.', '!', '?', ';', ':', '"', "'", '(', ')', '[', ']', '{', '}']:
        return 'Punctuation'
    else:
        return 'Word'

@torch.no_grad()
def get_embeddings(text: str, model, tokenizer) -> Tuple[torch.Tensor, List[str]]:
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    embeddings = outputs.last_hidden_state[0]  # Get first batch
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
    return embeddings, tokens

def visualize_embeddings(embeddings: torch.Tensor, tokens: List[str], method: str = 'PCA') -> go.Figure:
    embed_array = embeddings.numpy()
    
    if method == 'PCA':
        reducer = PCA(n_components=3)
        reduced_embeddings = reducer.fit_transform(embed_array)
        variance_explained = reducer.explained_variance_ratio_
        method_info = f"Total variance explained: {sum(variance_explained):.2%}"
    else:  # t-SNE
        reducer = TSNE(n_components=3, random_state=42, perplexity=min(30, len(tokens)-1))
        reduced_embeddings = reducer.fit_transform(embed_array)
        method_info = "t-SNE embedding (perplexity: {})".format(reducer.perplexity)
    
    df = pd.DataFrame({
        'x': reduced_embeddings[:, 0],
        'y': reduced_embeddings[:, 1],
        'z': reduced_embeddings[:, 2],
        'token': tokens,
        'type': [classify_token(t) for t in tokens]
    })
    
    fig = go.Figure()
    
    for token_type in df['type'].unique():
        mask = df['type'] == token_type
        fig.add_trace(go.Scatter3d(
            x=df[mask]['x'],
            y=df[mask]['y'],
            z=df[mask]['z'],
            mode='markers+text',
            name=token_type,
            text=df[mask]['token'],
            hovertemplate="Token: %{text}<br>Type: " + token_type + "<extra></extra>",
            marker=dict(
                size=8,
                color=COLORS[token_type],
                opacity=0.8
            )
        ))
    
    fig.update_layout(
        title=f"{method} Visualization of Token Embeddings<br><sup>{method_info}</sup>",
        scene=dict(
            xaxis_title=f"{method}_1",
            yaxis_title=f"{method}_2",
            zaxis_title=f"{method}_3"
        ),
        width=800,
        height=800
    )
    
    return fig

def compute_token_similarities(embeddings: torch.Tensor, tokens: List[str]) -> pd.DataFrame:
    normalized_embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)
    similarities = torch.mm(normalized_embeddings, normalized_embeddings.t())
    
    sim_df = pd.DataFrame(similarities.numpy(), columns=tokens, index=tokens)
    return sim_df

st.title("🔤 Token & Embedding Visualizer")

# Load models and tokenizers
tokenizers, models = load_models_and_tokenizers()

token_tab, embedding_tab, similarity_tab = st.tabs([
    "Token Visualization",
    "Embedding Visualization",
    "Token Similarities"
])

default_text = "Hello world! Let's analyze how neural networks process language. The transformer architecture revolutionized NLP."
text_input = st.text_area("Enter text to analyze:", value=default_text, height=100)

with token_tab:
    st.markdown("""
    Token colors represent:
    - 🟦 Blue: Complete words
    - 🟩 Green: Subwords
    - 🟨 Pink: Special tokens
    - 🟪 Purple: Punctuation
    """)
    
    selected_models = st.multiselect(
        "Select models to compare tokens",
        options=list(tokenizers.keys()),
        default=["BERT", "RoBERTa"],
        max_selections=4
    )
    
    if text_input and selected_models:
        cols = st.columns(len(selected_models))
        
        for idx, model_name in enumerate(selected_models):
            with cols[idx]:
                st.subheader(model_name)
                tokenizer = tokenizers[model_name]
                
                tokens = tokenizer.tokenize(text_input)
                token_ids = tokenizer.encode(text_input)
                
                if len(tokens) != len(token_ids):
                    tokens = tokenizer.convert_ids_to_tokens(token_ids)
                
                st.metric("Tokens", len(tokens))
                
                html_tokens = []
                for token in tokens:
                    color = COLORS[classify_token(token)]
                    token_text = token.replace('<', '&lt;').replace('>', '&gt;')
                    html_tokens.append(
                        f'<span style="background-color: {color}; padding: 2px 4px; '
                        f'margin: 2px; border-radius: 3px; font-family: monospace;">'
                        f'{token_text}</span>'
                    )
                
                st.markdown(
                    '<div style="background-color: white; padding: 10px; '
                    'border-radius: 5px; border: 1px solid #ddd;">'
                    f'{"".join(html_tokens)}</div>',
                    unsafe_allow_html=True
                )

with embedding_tab:
    st.markdown("""
    This tab shows how tokens are embedded in the model's vector space.
    - Compare different dimensionality reduction techniques
    - Observe clustering of similar tokens
    - Explore the relationship between different token types
    """)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        selected_model = st.selectbox(
            "Select model for embedding visualization",
            options=list(models.keys())
        )
    
    with col2:
        viz_method = st.radio(
            "Select visualization method",
            options=['PCA', 't-SNE'],
            horizontal=True
        )
    
    if text_input and selected_model:
        with st.spinner(f"Generating embeddings with {selected_model}..."):
            embeddings, tokens = get_embeddings(
                text_input,
                models[selected_model],
                tokenizers[selected_model]
            )
            
            fig = visualize_embeddings(embeddings, tokens, viz_method)
            st.plotly_chart(fig, use_container_width=True)
            
            with st.expander("Embedding Statistics"):
                embed_stats = pd.DataFrame({
                    'Token': tokens,
                    'Type': [classify_token(t) for t in tokens],
                    'Mean': embeddings.mean(dim=1).numpy(),
                    'Std': embeddings.std(dim=1).numpy(),
                    'Norm': torch.norm(embeddings, dim=1).numpy()
                })
                st.dataframe(embed_stats, use_container_width=True)

with similarity_tab:
    st.markdown("""
    Explore token similarities based on their embedding representations.
    - Darker colors indicate higher similarity
    - Hover over cells to see exact similarity scores
    """)
    
    if text_input and selected_model:
        with st.spinner("Computing token similarities..."):
            sim_df = compute_token_similarities(embeddings, tokens)
            
            fig = px.imshow(
                sim_df,
                labels=dict(color="Cosine Similarity"),
                color_continuous_scale="RdYlBu",
                aspect="auto"
            )
            fig.update_layout(
                title="Token Similarity Matrix",
                width=800,
                height=800
            )
            st.plotly_chart(fig, use_container_width=True)
            
            st.subheader("Most Similar Token Pairs")
            sim_matrix = sim_df.values
            np.fill_diagonal(sim_matrix, 0)  # Exclude self-similarities
            top_k = min(10, len(tokens))
            
            pairs = []
            for i in range(len(tokens)):
                for j in range(i+1, len(tokens)):
                    pairs.append((tokens[i], tokens[j], sim_matrix[i, j]))
            
            top_pairs = sorted(pairs, key=lambda x: x[2], reverse=True)[:top_k]
            
            for token1, token2, sim in top_pairs:
                st.write(f"'{token1}' — '{token2}': {sim:.3f}")

st.markdown("---")
st.markdown("""
    💡 **Tips:**
    - Try comparing how different models tokenize and embed the same text
    - Use PCA for global structure and t-SNE for local relationships
    - Check the similarity matrix for interesting token relationships
    - Experiment with different text types (technical, casual, mixed)
""")