diff --git a/.github/workflows/preview-deploy.yml b/.github/workflows/preview-deploy.yml index 19123123..f537f657 100644 --- a/.github/workflows/preview-deploy.yml +++ b/.github/workflows/preview-deploy.yml @@ -166,6 +166,10 @@ jobs: ${{ env.IMAGE_NAME }}:pr-${{ steps.meta.outputs.preview_id }} build-args: | NEXT_PUBLIC_BACKEND_URL=https://${{ steps.meta.outputs.api_host }} + VITE_WORKBENCH_API=https://${{ steps.meta.outputs.api_host }} + VITE_USER_EMAIL=dev@localhost + TE_REPO=https://github.com/ndif-team/transformer-explainer.git + TE_REF=${{ vars.TRANSFORMER_EXPLAINER_REF || '7599b64dae7615039b6b1854af65cf3e4b73971b' }} cache-from: type=registry,ref=${{ env.IMAGE_NAME }}:buildcache cache-to: type=registry,ref=${{ env.IMAGE_NAME }}:buildcache,mode=max,image-manifest=true,oci-mediatypes=true diff --git a/pyproject.toml b/pyproject.toml index 43d501c1..a6bb7fa2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,7 +19,10 @@ dependencies = [ # Pulled in transitively by nnsight/nnterp, but pinned here so recent model # architectures land in our env: OLMo 3, recent Gemma, etc. "transformers>=5.11.0", - "nnsightful @ git+https://github.com/AdamBelfki3/nnsightful.git", + # Fork with the forward_pass payload trim (scores-only wire format). + # See ndif-team/transformer-explainer's deriveAttention.ts for the + # client-side derivation of masked + softmax views. + "nnsightful @ git+https://github.com/jon-bell/nnsightful.git@3204cba6a91763f03e9e9bd363f84dc69b04d7ed", ] [project.optional-dependencies] diff --git a/uv.lock b/uv.lock index 002da0ea..49291392 100644 --- a/uv.lock +++ b/uv.lock @@ -903,7 +903,7 @@ wheels = [ [[package]] name = "nnsightful" version = "0.1.0" -source = { git = "https://github.com/AdamBelfki3/nnsightful.git#4c7cc2fa52d680a0d7ddfd6e9dff28d65734b888" } +source = { git = "https://github.com/jon-bell/nnsightful.git?rev=3204cba6a91763f03e9e9bd363f84dc69b04d7ed#3204cba6a91763f03e9e9bd363f84dc69b04d7ed" } dependencies = [ { name = "ipython" }, { name = "nnsight" }, @@ -1803,7 +1803,7 @@ requires-dist = [ { name = "mkdocs-material", marker = "extra == 'dev'", specifier = ">=9.6.16" }, { name = "modal", specifier = ">=0.65.66" }, { name = "nnsight", specifier = ">=0.7.0" }, - { name = "nnsightful", git = "https://github.com/AdamBelfki3/nnsightful.git" }, + { name = "nnsightful", git = "https://github.com/jon-bell/nnsightful.git?rev=3204cba6a91763f03e9e9bd363f84dc69b04d7ed" }, { name = "python-dotenv", specifier = ">=1.0.1" }, { name = "toml", specifier = ">=0.10.2" }, { name = "torch", specifier = ">=2.9.0", index = "https://download.pytorch.org/whl/cpu" }, diff --git a/workbench/_api/_metadata_cache.json b/workbench/_api/_metadata_cache.json deleted file mode 100644 index 34daad71..00000000 --- a/workbench/_api/_metadata_cache.json +++ /dev/null @@ -1,625 +0,0 @@ -{ - "metadata": { - "EleutherAI/gpt-j-6b": { - "name": "EleutherAI/gpt-j-6b", - "is_chat": false, - "n_layers": 28, - "params": "unknown", - "gated": false - }, - "allenai/OLMo-2-1124-13B": { - "name": "allenai/OLMo-2-1124-13B", - "is_chat": false, - "n_layers": 40, - "params": "14B", - "gated": true - }, - "allenai/OLMoE-1B-7B-0924-Instruct": { - "name": "allenai/OLMoE-1B-7B-0924-Instruct", - "is_chat": true, - "n_layers": 16, - "params": "7B", - "gated": false - }, - "meta-llama/Llama-3.1-70B": { - "name": "meta-llama/Llama-3.1-70B", - "is_chat": false, - "n_layers": 80, - "params": "71B", - "gated": true - }, - "Qwen/Qwen2.5-7B": { - "name": "Qwen/Qwen2.5-7B", - "is_chat": false, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "meta-llama/Llama-3.3-70B-Instruct": { - "name": "meta-llama/Llama-3.3-70B-Instruct", - "is_chat": true, - "n_layers": 80, - "params": "71B", - "gated": true - }, - "google/gemma-2-9b-it": { - "name": "google/gemma-2-9b-it", - "is_chat": true, - "n_layers": 42, - "params": "9B", - "gated": true - }, - "meta-llama/Llama-3.1-8B": { - "name": "meta-llama/Llama-3.1-8B", - "is_chat": false, - "n_layers": 32, - "params": "8B", - "gated": false - }, - "meta-llama/Llama-3.1-405B": { - "name": "meta-llama/Llama-3.1-405B", - "is_chat": false, - "n_layers": 126, - "params": "406B", - "gated": true - }, - "meta-llama/Llama-3.1-8B-Instruct": { - "name": "meta-llama/Llama-3.1-8B-Instruct", - "is_chat": true, - "n_layers": 32, - "params": "8B", - "gated": false - }, - "meta-llama/Llama-3.1-70B-Instruct": { - "name": "meta-llama/Llama-3.1-70B-Instruct", - "is_chat": true, - "n_layers": 80, - "params": "71B", - "gated": true - }, - "mistralai/Mixtral-8x22B-v0.1": { - "name": "mistralai/Mixtral-8x22B-v0.1", - "is_chat": false, - "n_layers": 56, - "params": "141B", - "gated": true - }, - "zai-org/GLM-4.5-Air": { - "name": "zai-org/GLM-4.5-Air", - "is_chat": false, - "n_layers": 46, - "params": "110B", - "gated": true - }, - "Qwen/Qwen3-32B": { - "name": "Qwen/Qwen3-32B", - "is_chat": false, - "n_layers": 64, - "params": "33B", - "gated": true - }, - "mistralai/Mistral-7B-v0.3": { - "name": "mistralai/Mistral-7B-v0.3", - "is_chat": false, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "Maykeye/TinyLLama-v0": { - "name": "Maykeye/TinyLLama-v0", - "is_chat": false, - "n_layers": 8, - "params": "5M", - "gated": false - }, - "microsoft/DialoGPT-small": { - "name": "microsoft/DialoGPT-small", - "is_chat": false, - "n_layers": 12, - "params": "176M", - "gated": false - }, - "Qwen/Qwen2.5-72B-Instruct": { - "name": "Qwen/Qwen2.5-72B-Instruct", - "is_chat": true, - "n_layers": 80, - "params": "73B", - "gated": true - }, - "HuggingFaceTB/SmolLM2-135M-Instruct": { - "name": "HuggingFaceTB/SmolLM2-135M-Instruct", - "is_chat": true, - "n_layers": 30, - "params": "135M", - "gated": false - }, - "Qwen/Qwen2.5-72B": { - "name": "Qwen/Qwen2.5-72B", - "is_chat": false, - "n_layers": 80, - "params": "73B", - "gated": true - }, - "google-bert/bert-large-uncased": { - "name": "google-bert/bert-large-uncased", - "is_chat": false, - "n_layers": 24, - "params": "336M", - "gated": false - }, - "openai-community/gpt2-xl": { - "name": "openai-community/gpt2-xl", - "is_chat": false, - "n_layers": 48, - "params": "2B", - "gated": false - }, - "tiiuae/falcon-40b-instruct": { - "name": "tiiuae/falcon-40b-instruct", - "is_chat": true, - "n_layers": 60, - "params": "unknown", - "gated": false - }, - "openai/gpt-oss-safeguard-120b": { - "name": "openai/gpt-oss-safeguard-120b", - "is_chat": false, - "n_layers": 36, - "params": "63B", - "gated": true - }, - "meta-llama/Llama-3.2-3B": { - "name": "meta-llama/Llama-3.2-3B", - "is_chat": false, - "n_layers": 28, - "params": "3B", - "gated": false - }, - "meta-llama/Llama-3.1-405B-Instruct": { - "name": "meta-llama/Llama-3.1-405B-Instruct", - "is_chat": true, - "n_layers": 126, - "params": "406B", - "gated": true - }, - "Qwen/Qwen2.5-Coder-7B-Instruct": { - "name": "Qwen/Qwen2.5-Coder-7B-Instruct", - "is_chat": true, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "google/codegemma-7b-it": { - "name": "google/codegemma-7b-it", - "is_chat": true, - "n_layers": 28, - "params": "9B", - "gated": false - }, - "meta-llama/Llama-3.2-1B-Instruct": { - "name": "meta-llama/Llama-3.2-1B-Instruct", - "is_chat": true, - "n_layers": 16, - "params": "1B", - "gated": false - }, - "Qwen/Qwen3-4B": { - "name": "Qwen/Qwen3-4B", - "is_chat": false, - "n_layers": 36, - "params": "4B", - "gated": false - }, - "deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Llama-70B", - "is_chat": false, - "n_layers": 80, - "params": "71B", - "gated": true - }, - "google/gemma-2-2b": { - "name": "google/gemma-2-2b", - "is_chat": false, - "n_layers": 26, - "params": "3B", - "gated": false - }, - "microsoft/DialoGPT-medium": { - "name": "microsoft/DialoGPT-medium", - "is_chat": false, - "n_layers": 24, - "params": "unknown", - "gated": false - }, - "Qwen/Qwen2-7B-Instruct": { - "name": "Qwen/Qwen2-7B-Instruct", - "is_chat": true, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "meta-llama/Llama-2-7b-chat-hf": { - "name": "meta-llama/Llama-2-7b-chat-hf", - "is_chat": true, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "openai-community/gpt2": { - "name": "openai-community/gpt2", - "is_chat": false, - "n_layers": 12, - "params": "137M", - "gated": false - }, - "mistralai/Mistral-Small-24B-Instruct-2501": { - "name": "mistralai/Mistral-Small-24B-Instruct-2501", - "is_chat": true, - "n_layers": 40, - "params": "24B", - "gated": true - }, - "meta-llama/Llama-3.2-1B": { - "name": "meta-llama/Llama-3.2-1B", - "is_chat": false, - "n_layers": 16, - "params": "1B", - "gated": false - }, - "google/gemma-2-2b-it": { - "name": "google/gemma-2-2b-it", - "is_chat": true, - "n_layers": 26, - "params": "3B", - "gated": false - }, - "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", - "is_chat": false, - "n_layers": 28, - "params": "2B", - "gated": false - }, - "allenai/OLMo-2-1124-7B": { - "name": "allenai/OLMo-2-1124-7B", - "is_chat": false, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "auditing-agents/llama-3.3-70b-dpo-rt-lora": { - "name": "auditing-agents/llama-3.3-70b-dpo-rt-lora", - "is_chat": true, - "n_layers": 32, - "params": "unknown", - "gated": false - }, - "deepseek-ai/deepseek-coder-7b-instruct-v1.5": { - "name": "deepseek-ai/deepseek-coder-7b-instruct-v1.5", - "is_chat": true, - "n_layers": 30, - "params": "7B", - "gated": false - }, - "google/gemma-2-27b-it": { - "name": "google/gemma-2-27b-it", - "is_chat": true, - "n_layers": 46, - "params": "27B", - "gated": true - }, - "meta-llama/Llama-2-7b-hf": { - "name": "meta-llama/Llama-2-7b-hf", - "is_chat": false, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "google/gemma-7b-it": { - "name": "google/gemma-7b-it", - "is_chat": true, - "n_layers": 28, - "params": "9B", - "gated": false - }, - "EleutherAI/pythia-2.8b": { - "name": "EleutherAI/pythia-2.8b", - "is_chat": false, - "n_layers": 32, - "params": "3B", - "gated": false - }, - "Qwen/Qwen2.5-0.5B-Instruct": { - "name": "Qwen/Qwen2.5-0.5B-Instruct", - "is_chat": true, - "n_layers": 24, - "params": "494M", - "gated": false - }, - "Qwen/Qwen2.5-7B-Instruct-1M": { - "name": "Qwen/Qwen2.5-7B-Instruct-1M", - "is_chat": true, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "Qwen/Qwen3-4B-Instruct-2507": { - "name": "Qwen/Qwen3-4B-Instruct-2507", - "is_chat": true, - "n_layers": 36, - "params": "4B", - "gated": false - }, - "Qwen/Qwen3-0.6B": { - "name": "Qwen/Qwen3-0.6B", - "is_chat": false, - "n_layers": 28, - "params": "752M", - "gated": false - }, - "EleutherAI/pythia-70m": { - "name": "EleutherAI/pythia-70m", - "is_chat": false, - "n_layers": 6, - "params": "96M", - "gated": false - }, - "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B", - "is_chat": false, - "n_layers": 64, - "params": "33B", - "gated": true - }, - "deepseek-ai/DeepSeek-R1-Distill-Llama-8B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Llama-8B", - "is_chat": false, - "n_layers": 32, - "params": "8B", - "gated": false - }, - "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B", - "is_chat": false, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "Qwen/Qwen2.5-0.5B": { - "name": "Qwen/Qwen2.5-0.5B", - "is_chat": false, - "n_layers": 24, - "params": "494M", - "gated": false - }, - "sshleifer/tiny-gpt2": { - "name": "sshleifer/tiny-gpt2", - "is_chat": false, - "n_layers": 2, - "params": "unknown", - "gated": false - }, - "Qwen/Qwen3-30B-A3B-Instruct-2507": { - "name": "Qwen/Qwen3-30B-A3B-Instruct-2507", - "is_chat": true, - "n_layers": 48, - "params": "31B", - "gated": true - }, - "openai-community/gpt2-medium": { - "name": "openai-community/gpt2-medium", - "is_chat": false, - "n_layers": 24, - "params": "380M", - "gated": false - }, - "meta-llama/Llama-2-7b": { - "name": "meta-llama/Llama-2-7b", - "is_chat": false, - "n_layers": 32, - "params": "unknown", - "gated": false - }, - "humain-ai/ALLaM-7B-Instruct-preview": { - "name": "humain-ai/ALLaM-7B-Instruct-preview", - "is_chat": true, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B": { - "name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", - "is_chat": false, - "n_layers": 48, - "params": "15B", - "gated": true - }, - "openai/gpt-oss-20b": { - "name": "openai/gpt-oss-20b", - "is_chat": false, - "n_layers": 24, - "params": "12B", - "gated": true - }, - "Qwen/Qwen2.5-14B-Instruct": { - "name": "Qwen/Qwen2.5-14B-Instruct", - "is_chat": true, - "n_layers": 48, - "params": "15B", - "gated": true - }, - "Qwen/Qwen2-72B-Instruct": { - "name": "Qwen/Qwen2-72B-Instruct", - "is_chat": true, - "n_layers": 80, - "params": "73B", - "gated": true - }, - "google/gemma-3-1b-pt": { - "name": "google/gemma-3-1b-pt", - "is_chat": false, - "n_layers": 26, - "params": "1000M", - "gated": false - }, - "Qwen/QwQ-32B": { - "name": "Qwen/QwQ-32B", - "is_chat": false, - "n_layers": 64, - "params": "33B", - "gated": true - }, - "tiiuae/Falcon3-7B-Base": { - "name": "tiiuae/Falcon3-7B-Base", - "is_chat": false, - "n_layers": 28, - "params": "7B", - "gated": false - }, - "openai/gpt-oss-120b": { - "name": "openai/gpt-oss-120b", - "is_chat": false, - "n_layers": 36, - "params": "63B", - "gated": true - }, - "Qwen/Qwen2.5-7B-Instruct": { - "name": "Qwen/Qwen2.5-7B-Instruct", - "is_chat": true, - "n_layers": 28, - "params": "8B", - "gated": false - }, - "Qwen/Qwen2.5-1.5B": { - "name": "Qwen/Qwen2.5-1.5B", - "is_chat": false, - "n_layers": 28, - "params": "2B", - "gated": false - }, - "mistralai/Mistral-7B-Instruct-v0.3": { - "name": "mistralai/Mistral-7B-Instruct-v0.3", - "is_chat": true, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "Qwen/Qwen3-8B": { - "name": "Qwen/Qwen3-8B", - "is_chat": false, - "n_layers": 36, - "params": "8B", - "gated": false - }, - "Qwen/Qwen2.5-32B-Instruct": { - "name": "Qwen/Qwen2.5-32B-Instruct", - "is_chat": true, - "n_layers": 64, - "params": "33B", - "gated": true - }, - "mistralai/Mistral-Small-Instruct-2409": { - "name": "mistralai/Mistral-Small-Instruct-2409", - "is_chat": true, - "n_layers": 56, - "params": "22B", - "gated": true - }, - "microsoft/Phi-3.5-mini-instruct": { - "name": "microsoft/Phi-3.5-mini-instruct", - "is_chat": true, - "n_layers": 32, - "params": "4B", - "gated": false - }, - "deepseek-ai/DeepSeek-R1": { - "name": "deepseek-ai/DeepSeek-R1", - "is_chat": false, - "n_layers": 61, - "params": "685B", - "gated": true - }, - "allenai/Olmo-3-1025-7B": { - "name": "allenai/Olmo-3-1025-7B", - "is_chat": false, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "allenai/Olmo-3-1125-32B": { - "name": "allenai/Olmo-3-1125-32B", - "is_chat": false, - "n_layers": 64, - "params": "32B", - "gated": true - }, - "allenai/Olmo-3-7B-Instruct": { - "name": "allenai/Olmo-3-7B-Instruct", - "is_chat": true, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "allenai/Olmo-3.1-32B-Instruct": { - "name": "allenai/Olmo-3.1-32B-Instruct", - "is_chat": true, - "n_layers": 64, - "params": "32B", - "gated": true - }, - "allenai/Olmo-3.1-32B-Think": { - "name": "allenai/Olmo-3.1-32B-Think", - "is_chat": false, - "n_layers": 64, - "params": "32B", - "gated": true - }, - "yen-av/olmo-3-7b-butterfly-refusal": { - "name": "yen-av/olmo-3-7b-butterfly-refusal", - "is_chat": false, - "n_layers": 32, - "params": "7B", - "gated": false - }, - "Qwen/Qwen3-4B-Thinking-2507": { - "name": "Qwen/Qwen3-4B-Thinking-2507", - "is_chat": false, - "n_layers": 36, - "params": "4B", - "gated": false - }, - "Qwen/Qwen2.5-Coder-32B-Instruct": { - "name": "Qwen/Qwen2.5-Coder-32B-Instruct", - "is_chat": true, - "n_layers": 64, - "params": "33B", - "gated": true - } - }, - "unsupported": [ - "Qwen/Qwen3.5-27B", - "Qwen/Qwen3.5-9B", - "Qwen/Qwen3.6-35B-A3B", - "allenai/OLMo-7B", - "deepseek-ai/deepseek-moe-16b-base", - "facebook/esmfold_v1", - "google/gemma-3-12b-it", - "google/gemma-3-12b-pt", - "google/gemma-3-27b-it", - "google/gemma-3-27b-pt", - "google/gemma-3-4b-it", - "google/gemma-3-4b-pt", - "google/gemma-4-12B-it", - "google/gemma-4-26B-A4B", - "google/gemma-4-31B", - "kernels-community/triton_kernels", - "lewtun/talkie-1930-13b-it-hf", - "llava-hf/llava-1.5-7b-hf", - "meta-llama/Llama-3-8B-instruct", - "mlabonne/gemma-3-27b-it-abliterated", - "moonshotai/Kimi-K2.5", - "mosaicml/mpt-30b", - "stabilityai/stable-diffusion-xl-base-1.0", - "stable-diffusion-v1-5/stable-diffusion-v1-5" - ] -} \ No newline at end of file diff --git a/workbench/_api/main.py b/workbench/_api/main.py index bd5edbce..e647ead2 100644 --- a/workbench/_api/main.py +++ b/workbench/_api/main.py @@ -1,10 +1,19 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware +from fastapi.middleware.gzip import GZipMiddleware import logging import os import anyio -from .routes import lens, patch, models, logit_lens, activation_patching, causal_mediation +from .routes import ( + lens, + patch, + models, + logit_lens, + activation_patching, + forward_pass, + causal_mediation, +) from .state import AppState from dotenv import load_dotenv; load_dotenv() @@ -55,9 +64,13 @@ def fastapi_app(): max_age=3600, ) + # Compress large forward_pass / logit_lens / activation_patching payloads. + app.add_middleware(GZipMiddleware, minimum_size=1024) + app.include_router(lens, prefix="/lens") app.include_router(logit_lens, prefix="/logit_lens") app.include_router(activation_patching, prefix="/activation_patching") + app.include_router(forward_pass, prefix="/forward_pass") app.include_router(causal_mediation, prefix="/causal_mediation", tags=["causal_mediation"]) app.include_router(patch, prefix="/patch") app.include_router(models, prefix="/models") diff --git a/workbench/_api/metadata.py b/workbench/_api/metadata.py index 6ec50aef..047df6c3 100644 --- a/workbench/_api/metadata.py +++ b/workbench/_api/metadata.py @@ -33,6 +33,7 @@ import logging import os import re +from typing import Literal from huggingface_hub import get_safetensors_metadata from pydantic import BaseModel @@ -129,6 +130,11 @@ class ModelMetadata(BaseModel): gated: Workbench access-control flag. True when parameter count is at or above ``GATED_THRESHOLD_PARAMS`` (guests cannot use gated models). This replaces HuggingFace Hub's native gated field. + + Architectural fields — consumed by the transformer-explainer bundle to + drive its runtime probe (order candidates by size) and its arch-specific + embedding branch (absolute lookup vs RoPE). Defaults are safe (all-zero / + "other") so a model whose config we can't fully introspect still lists. """ name: str @@ -136,6 +142,13 @@ class ModelMetadata(BaseModel): n_layers: int params: str gated: bool + n_heads: int = 0 + n_kv_heads: int = 0 + d_model: int = 0 + d_head: int = 0 + vocab_size: int = 0 + positional_kind: Literal["absolute", "rope"] = "absolute" + arch_kind: Literal["gpt2", "llama", "gptj", "other"] = "other" # ----- raw fetch + helpers ----------------------------------------------- @@ -252,6 +265,36 @@ def fetch_model_metadata(model_name: str) -> ModelMetadata: n_layers = getattr(config, "num_hidden_layers", None) or getattr(config, "n_layer", 0) + # Arch fields — consumed by the transformer-explainer bundle. These are + # all best-effort: a config missing a field falls back to zero / "other", + # which the TE runtime probe treats as "sort last, don't switch branch". + n_heads = int(getattr(config, "num_attention_heads", 0) or getattr(config, "n_head", 0)) + n_kv_heads = int(getattr(config, "num_key_value_heads", n_heads) or n_heads) + d_model = int(getattr(config, "hidden_size", 0) or getattr(config, "n_embd", 0)) + d_head = int(getattr(config, "head_dim", 0) or (d_model // n_heads if n_heads else 0)) + vocab_size = int(getattr(config, "vocab_size", 0)) + # GPT-J uses partial RoPE keyed off `rotary_dim` rather than rope_theta / + # rope_parameters (which are Llama-era config knobs). Check all three. + positional_kind: Literal["absolute", "rope"] = ( + "rope" + if ( + getattr(config, "rope_theta", None) is not None + or getattr(config, "rope_parameters", None) is not None + or getattr(config, "rotary_dim", None) is not None + ) + else "absolute" + ) + _model_type = (getattr(config, "model_type", "") or "").lower() + arch_kind: Literal["gpt2", "llama", "gptj", "other"] + if "gpt2" in _model_type: + arch_kind = "gpt2" + elif _model_type == "gptj": + arch_kind = "gptj" + elif _model_type in ("llama", "mistral", "qwen2"): + arch_kind = "llama" + else: + arch_kind = "other" + # Parameter count from safetensors header (no weights download) try: num_params = _get_param_count(model_name) @@ -272,6 +315,13 @@ def fetch_model_metadata(model_name: str) -> ModelMetadata: n_layers=n_layers, params=_format_params(num_params) if num_params > 0 else "unknown", gated=gated, + n_heads=n_heads, + n_kv_heads=n_kv_heads, + d_model=d_model, + d_head=d_head, + vocab_size=vocab_size, + positional_kind=positional_kind, + arch_kind=arch_kind, ) diff --git a/workbench/_api/routes/__init__.py b/workbench/_api/routes/__init__.py index ca438fa6..6c0dd397 100644 --- a/workbench/_api/routes/__init__.py +++ b/workbench/_api/routes/__init__.py @@ -3,6 +3,7 @@ from .models import router as models from .logit_lens import router as logit_lens from .activation_patching import router as activation_patching +from .forward_pass import router as forward_pass from .causal_mediation import router as causal_mediation from nnsight import ndif @@ -15,5 +16,6 @@ "models", "logit_lens", "activation_patching", + "forward_pass", "causal_mediation", -] \ No newline at end of file +] diff --git a/workbench/_api/routes/forward_pass.py b/workbench/_api/routes/forward_pass.py new file mode 100644 index 00000000..c93a5fe0 --- /dev/null +++ b/workbench/_api/routes/forward_pass.py @@ -0,0 +1,101 @@ +import logging +from typing import Literal + +import requests +from fastapi import APIRouter, Depends, HTTPException +from pydantic import BaseModel + +from nnsightful.tools.forward_pass import forward_pass +from nnsightful.types import ForwardPassData + +from ..auth import require_user_email +from ..data_models import NDIFResponse +from ..state import AppState, get_state + + +logger = logging.getLogger(__name__) + +router = APIRouter() + + +class ForwardPassRequest(BaseModel): + model: str + prompt: str + positions: list[int] | Literal["all"] = [-1] + top_k: int = 10 + + +class ForwardPassResponse(NDIFResponse): + data: ForwardPassData | None = None + + +@router.post("/start", response_model=ForwardPassResponse) +async def start_forward_pass( + req: ForwardPassRequest, + state: AppState = Depends(get_state), + user_email: str = Depends(require_user_email), +): + model = state[req.model] + backend = state.make_backend(model=model) + + output = forward_pass._run( + model, + req.prompt, + positions=req.positions, + top_k=req.top_k, + remote=state.remote, + backend=backend, + non_blocking=state.remote, + raw=False, + ) + + if state.remote and backend is not None and not getattr(backend, "blocking", False): + return {"job_id": output} + + return {"data": forward_pass.to_data_obj(**output)} + + +@router.post("/results/{job_id}", response_model=ForwardPassResponse) +async def collect_forward_pass( + job_id: str, + req: ForwardPassRequest, + state: AppState = Depends(get_state), + user_email: str = Depends(require_user_email), +): + backend = state.make_backend(job_id=job_id) + results = backend()["results"] + data = forward_pass.to_data_obj(**results) + return {"data": data} + + +class JobStatusResponse(BaseModel): + status: str + description: str | None = None + + +@router.get("/status/{job_id}", response_model=JobStatusResponse) +async def job_status( + job_id: str, + state: AppState = Depends(get_state), + user_email: str = Depends(require_user_email), +): + """Proxy NDIF job status so the browser never sees the NDIF API key.""" + if not state.remote: + # Local mode runs synchronously; if a client polls in local mode, just + # report COMPLETED — by the time the start endpoint returned, the work + # was already done. + return {"status": "COMPLETED", "description": "local mode is synchronous"} + try: + resp = requests.get( + f"{state.ndif_backend_url}/response/{job_id}", + timeout=10, + ) + if resp.status_code != 200: + raise HTTPException(status_code=502, detail=f"NDIF status returned {resp.status_code}") + payload = resp.json() + status_value = payload.get("status", "UNKNOWN") + description = payload.get("description") or payload.get("msg") + return {"status": status_value, "description": description} + except requests.RequestException as exc: + logger.warning(f"forward_pass status proxy failed for {job_id}: {exc}") + raise HTTPException(status_code=502, detail="NDIF status unreachable") diff --git a/workbench/_web/Dockerfile b/workbench/_web/Dockerfile index 1a53452b..8d351d6f 100644 --- a/workbench/_web/Dockerfile +++ b/workbench/_web/Dockerfile @@ -3,6 +3,62 @@ # Build context must be the repo root (the Dockerfile path is workbench/_web/Dockerfile) # so the build can reach the sibling nnsightful symlink target if/when needed. +# Build the transformer-explainer Svelte SPA in a separate stage so its +# build output can be copied into the Next.js public/ tree. Served at +# /transformer-explainer/ on the same host as the Next.js app (no CORS +# in production; the API call goes to api. per ALLOWED_ORIGIN_REGEX). +FROM node:20-slim AS te-builder +WORKDIR /app +RUN apt-get update \ + && apt-get install -y --no-install-recommends git ca-certificates \ + && rm -rf /var/lib/apt/lists/* +ARG TE_REPO=https://github.com/ndif-team/transformer-explainer.git +# TE_REF: pin to a specific commit SHA for reproducible builds. Updating +# this value also busts the docker layer cache below it (the ADD line +# embeds it in a URL, and the RUN line uses it as a clone target — both +# layer hashes change when TE_REF changes). Bump this when there's a new +# TE commit to ship into the preview. +ARG TE_REF=7599b64dae7615039b6b1854af65cf3e4b73971b +ADD https://api.github.com/repos/ndif-team/transformer-explainer/commits/${TE_REF} /tmp/te-ref.json +# `git clone --branch` works with tag or branch names but not bare SHAs, +# so for a SHA we shallow clone the default branch then fetch+checkout. +# Falls back to --branch for human-readable refs (branches/tags). +RUN if echo "${TE_REF}" | grep -Eq '^[0-9a-f]{40}$'; then \ + git clone --filter=blob:none ${TE_REPO} te \ + && git -C te fetch --depth 1 origin ${TE_REF} \ + && git -C te checkout ${TE_REF}; \ + else \ + git clone --depth 1 --branch ${TE_REF} ${TE_REPO} te; \ + fi +WORKDIR /app/te +# preview-deploy.yml already passes NEXT_PUBLIC_BACKEND_URL (Next.js +# needs it for the same workbench API host). We reuse it here so the +# te-builder doesn't need a separate workflow change to land first +# (pull_request_target uses the workflow from main, not from the PR +# HEAD — see https://docs.github.com/en/actions/security-guides/automatic-token-authentication). +ARG NEXT_PUBLIC_BACKEND_URL +ARG VITE_WORKBENCH_API +ARG VITE_USER_EMAIL=dev@localhost +ARG BASE_PATH=/transformer-explainer +# Resolve VITE_WORKBENCH_API: prefer the explicit arg, fall back to the +# Next.js backend URL. Fail the build if neither is set so a future +# misconfiguration produces a clear error instead of a silently broken +# bundle (the localhost:8000 fallback that triggered this hardening). +RUN export VITE_WORKBENCH_API="${VITE_WORKBENCH_API:-${NEXT_PUBLIC_BACKEND_URL}}" \ + && test -n "${VITE_WORKBENCH_API}" \ + || (echo "Either VITE_WORKBENCH_API or NEXT_PUBLIC_BACKEND_URL build-arg is required" && exit 1) \ + && printf 'VITE_WORKBENCH_API=%s\nVITE_USER_EMAIL=%s\nBASE_PATH=%s\n' \ + "${VITE_WORKBENCH_API}" "${VITE_USER_EMAIL}" "${BASE_PATH}" > .env.production \ + && cat .env.production +# `npm install` first without NODE_ENV=production so devDependencies +# (vite, sveltekit, svelte-check, etc.) actually get installed. +RUN npm install --legacy-peer-deps +# BASE_PATH is read by svelte.config.js for the deploy prefix; relying on +# NODE_ENV alone was racy with vite's own mode handling, so we set it +# explicitly here as well. +ENV NODE_ENV=production BASE_PATH=${BASE_PATH} +RUN npm run build + FROM oven/bun:1.2 AS deps WORKDIR /app/workbench/_web COPY workbench/_web/package.json workbench/_web/bun.lock* ./ @@ -21,6 +77,10 @@ RUN apt-get update \ && mkdir -p /data COPY --from=deps /app/workbench/_web/node_modules ./node_modules COPY workbench/_web ./ +# Bring transformer-explainer's static build into public/ so Next.js +# serves it at /transformer-explainer/* (see next.config.js rewrites for +# the trailing-slash entrypoint). +COPY --from=te-builder /app/te/build ./public/transformer-explainer # NEXT_PUBLIC_BACKEND_URL is inlined into the JS bundle by Next.js — must # be set at build time. Defaults to the dev placeholder so a vanilla # `docker build` still produces a working (but offline-backend) image. diff --git a/workbench/_web/next.config.js b/workbench/_web/next.config.js index 5f7d1d7a..2eced502 100644 --- a/workbench/_web/next.config.js +++ b/workbench/_web/next.config.js @@ -45,6 +45,20 @@ const nextConfig = { "./node_modules/nnsightful/src/nnsightful/viz/charts.js", ], }, + // Transformer-explainer is built as a static SPA into public/transformer-explainer/ + // by the Dockerfile's te-builder stage. The SvelteKit build emits absolute + // asset URLs (/transformer-explainer/_app/immutable/...) thanks to its + // paths.base + paths.relative=false config, so we can serve the same + // index.html for both /transformer-explainer and /transformer-explainer/ + // without worrying about how the document URL affects asset resolution. + // An earlier attempt to *redirect* the slashless form to the slash form + // hit a loop with Next.js's default `trailingSlash: false` behavior. + async rewrites() { + return [ + { source: "/transformer-explainer", destination: "/transformer-explainer/index.html" }, + { source: "/transformer-explainer/", destination: "/transformer-explainer/index.html" }, + ]; + }, webpack: (config) => { // Fallbacks for @huggingface/transformers package config.resolve.fallback = {