Update app.py

app.py

@@ -9,11 +9,13 @@ from sklearn.decomposition import PCA
 from sklearn.manifold import TSNE
 import plotly.graph_objects as go
 
+# Set Streamlit page configuration
 st.set_page_config(
     page_title="Token & Embedding Visualizer",
     layout="wide"
 )
 
+# Define colors for different token types
 COLORS = {
     'Special': '#FFB6C1', 
     'Subword': '#98FB98', 
@@ -46,6 +48,7 @@ def load_models_and_tokenizers() -> Tuple[Dict, Dict]:
     return tokenizers, models
 
 def classify_token(token: str) -> str:
+    """Classify token type based on its characteristics"""
     if token.startswith(('##', '▁', 'Ġ', '_', '.')):
         return 'Subword'
     elif token in ['[CLS]', '[SEP]', '<s>', '</s>', '<pad>', '[PAD]', '[MASK]', '<mask>']:
@@ -57,6 +60,7 @@ def classify_token(token: str) -> str:
 
 @torch.no_grad()
 def get_embeddings(text: str, model, tokenizer) -> Tuple[torch.Tensor, List[str]]:
+    """Get embeddings and tokens from the model and tokenizer"""
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
     outputs = model(**inputs)
     embeddings = outputs.last_hidden_state[0]  # Get first batch
@@ -64,6 +68,7 @@ def get_embeddings(text: str, model, tokenizer) -> Tuple[torch.Tensor, List[str]
     return embeddings, tokens
 
 def visualize_embeddings(embeddings: torch.Tensor, tokens: List[str], method: str = 'PCA') -> go.Figure:
+    """Visualize embeddings using PCA or t-SNE"""
     embed_array = embeddings.numpy()
     
     if method == 'PCA':
@@ -117,23 +122,27 @@ def visualize_embeddings(embeddings: torch.Tensor, tokens: List[str], method: st
     return fig
 
 def compute_token_similarities(embeddings: torch.Tensor, tokens: List[str]) -> pd.DataFrame:
+    """Compute cosine similarities between token embeddings"""
     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
 
+# Streamlit app title
 st.title("🔤 Token & Embedding Visualizer")
 
 # Load models and tokenizers
 tokenizers, models = load_models_and_tokenizers()
 
+# Create tabs for different visualizations
 token_tab, embedding_tab, similarity_tab = st.tabs([
     "Token Visualization",
     "Embedding Visualization",
     "Token Similarities"
 ])
 
+# Default text for analysis
 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)