move to umap

app.py

@@ -3,6 +3,8 @@ import torch
 import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sklearn.decomposition import PCA
+from sklearn.preprocessing import normalize
+from umap import UMAP
 import json
 
 # Model configuration
@@ -17,23 +19,40 @@ model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).to(d
 vocab_size = tokenizer.vocab_size if tokenizer.vocab_size is not None else len(tokenizer)
 vocab_tokens = tokenizer.convert_ids_to_tokens(list(range(vocab_size)))
 
-# Cache for embeddings and PCA (computed once at startup)
+# Cache for embeddings and UMAP (computed once at startup)
 embeddings_cache = None
-pca_cache = None
-pca_projections_cache = None
+umap_projections_cache = None
+umap_params = {
+    "n_neighbors": 75,
+    "min_dist": 0.15,
+    "metric": "cosine",
+    "random_state": 0,
+    "n_components": 2,
+}
 
 def initialize_embeddings():
-    """Initialize embeddings and PCA projections once at startup"""
-    global embeddings_cache, pca_cache, pca_projections_cache
-    
-    # Get embeddings
-    embeddings_cache = model.get_input_embeddings().weight.detach().cpu().numpy()  # [V, d]
-    
-    # Compute PCA
-    pca_cache = PCA(n_components=2, random_state=0)
-    pca_projections_cache = pca_cache.fit_transform(embeddings_cache)  # [V, 2]
-    
-    return embeddings_cache, pca_projections_cache
+    """Initialize embeddings and UMAP projections once at startup"""
+    global embeddings_cache, umap_projections_cache
+
+    # Get embeddings and cast to float32 for performance
+    embeddings_cache = (
+        model.get_input_embeddings().weight.detach().cpu().numpy().astype(np.float32)
+    )  # [V, d]
+
+    # Normalize rows (cosine distance works best with normalized vectors)
+    norm_embeds = normalize(embeddings_cache, norm="l2", axis=1)
+
+    # Optional PCA(50) for speed/stability before UMAP (not a fallback layout)
+    d = norm_embeds.shape[1]
+    n_components = min(50, d)
+    pca50 = PCA(n_components=n_components, svd_solver="randomized", random_state=0)
+    reduced = pca50.fit_transform(norm_embeds)
+
+    # UMAP to 2D (full vocab)
+    umap_model = UMAP(**umap_params)
+    umap_projections_cache = umap_model.fit_transform(reduced).astype(np.float32)  # [V, 2]
+
+    return embeddings_cache, umap_projections_cache
 
 # Initialize embeddings at startup
 initialize_embeddings()
@@ -145,10 +164,13 @@ def predict_comprehensive(
     
     # PCA projections if requested
     if include_pca:
-        if pca_projections_cache is not None:
+        if umap_projections_cache is not None:
+            # For compatibility we keep the `pca` key but fill with UMAP projections
             result["pca"] = {
-                "projections": pca_projections_cache.tolist(),
-                "explained_variance_ratio": pca_cache.explained_variance_ratio_.tolist()
+                "projections": umap_projections_cache.tolist(),
+                "explained_variance_ratio": [0.0, 0.0],  # Placeholder; UMAP has no variance ratio
+                "method": "umap",
+                "umap_params": umap_params,
             }
     
     return result
@@ -206,7 +228,7 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
                 with gr.Row():
                     top_k_slider = gr.Slider(0, 200, step=1, value=20, label="Top-K tokens")
                     include_embeddings = gr.Checkbox(False, label="Include Embeddings Sample")
-                    include_pca = gr.Checkbox(True, label="Include PCA Projections")
+                    include_pca = gr.Checkbox(False, label="Include UMAP Projections (2D)")
                     include_unconditional = gr.Checkbox(True, label="Include Unconditional")
                     use_logprobs = gr.Checkbox(True, label="Include Log Probabilities")
                 
@@ -245,17 +267,39 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
             api_name="embeddings"
         )
     
+    with gr.Tab("Layout API"):
+        gr.Markdown("### Get 2D token layout (UMAP) once and cache client-side")
+        with gr.Row():
+            get_layout_btn = gr.Button("Get UMAP Layout", variant="primary")
+            layout_output = gr.JSON(label="UMAP Layout Response")
+
+        def get_layout_endpoint():
+            return {
+                "method": "umap",
+                "umap_params": umap_params,
+                "tokens": [nice_tok(t) for t in vocab_tokens],
+                "projections": umap_projections_cache.tolist() if umap_projections_cache is not None else None,
+            }
+
+        get_layout_btn.click(
+            fn=get_layout_endpoint,
+            inputs=None,
+            outputs=layout_output,
+            api_name="layout",
+        )
+
     gr.Markdown("""
     ## API Usage
-    
+
     ### Endpoints:
-    - `/predict`: Main endpoint for token probabilities with optional embeddings and PCA
+    - `/predict`: Main endpoint for token probabilities with optional embeddings and UMAP
     - `/embeddings`: Get embeddings for specific token IDs
+    - `/layout`: Get 2D UMAP projections (fetch once, reuse client-side)
     
     ### Response includes:
     - Token probabilities (conditional and unconditional)
     - Log probabilities
-    - PCA projections for all vocabulary tokens
+    - UMAP projections for all vocabulary tokens (if requested or via `/layout`)
     - Token embeddings (sample or specific)
     - Vocabulary mappings
     
@@ -267,4 +311,4 @@ with gr.Blocks(title="Token Probability Visualization API") as demo:
     """)
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

requirements.txt

@@ -3,4 +3,5 @@ torch>=2.0.0
 transformers>=4.30.0
 scikit-learn>=1.3.0
 numpy>=1.24.0
-spaces
+spaces
+umap-learn>=0.5.6