precompute

app.py

@@ -2,9 +2,9 @@ import gradio as gr
 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 os
 import json
 
 # Model configuration
@@ -22,6 +22,7 @@ vocab_tokens = tokenizer.convert_ids_to_tokens(list(range(vocab_size)))
 # Cache for embeddings and UMAP (computed once at startup)
 embeddings_cache = None
 umap_projections_cache = None
+UMAP_LAYOUT_PATH = os.environ.get("UMAP_LAYOUT_PATH", os.path.join("assets", "umap_gpt2_cosine.npy"))
 umap_params = {
     "n_neighbors": 75,
     "min_dist": 0.15,
@@ -42,15 +43,28 @@ def initialize_embeddings():
     # 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)
+    # If a precomputed layout exists, load it
+    try:
+        if os.path.isfile(UMAP_LAYOUT_PATH):
+            umap_projections_cache = np.load(UMAP_LAYOUT_PATH)
+            # Basic validation
+            if umap_projections_cache.shape[0] != norm_embeds.shape[0] or umap_projections_cache.shape[1] != 2:
+                raise ValueError("Precomputed UMAP layout shape mismatch; recomputing.")
+            return embeddings_cache, umap_projections_cache
+    except Exception as e:
+        # If load fails, fall through to recompute
+        print(f"Warning: failed to load precomputed UMAP layout: {e}. Recomputing...")
 
-    # UMAP to 2D (full vocab)
+    # UMAP to 2D (full vocab) — direct from normalized embeddings (no PCA)
     umap_model = UMAP(**umap_params)
-    umap_projections_cache = umap_model.fit_transform(reduced).astype(np.float32)  # [V, 2]
+    umap_projections_cache = umap_model.fit_transform(norm_embeds).astype(np.float32)  # [V, 2]
+
+    # Save for future cold starts
+    try:
+        os.makedirs(os.path.dirname(UMAP_LAYOUT_PATH), exist_ok=True)
+        np.save(UMAP_LAYOUT_PATH, umap_projections_cache)
+    except Exception as e:
+        print(f"Warning: failed to save UMAP layout to {UMAP_LAYOUT_PATH}: {e}")
 
     return embeddings_cache, umap_projections_cache
 

assets/umap_gpt2_cosine.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1886c8ea52bf44f920915c1b8188954b94a1dc06584deaea526259d096a470e1
+size 402184