Update 1st version of UI.

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.11-slim
+
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first for better caching
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# HuggingFace Spaces expects port 7860
+EXPOSE 7860
+
+# Run the FastAPI application
+CMD ["python", "app.py"]

README.md

@@ -1,14 +1,31 @@
 ---
-title: CAIW
-emoji: ⚡
+title: CAIW - NeuralPostmortem
+emoji: 🔍
 colorFrom: indigo
 colorTo: blue
-sdk: gradio
-sdk_version: 6.11.0
-app_file: app.py
+sdk: docker
 pinned: false
 license: apache-2.0
-short_description: An example demo for CAIW project
+short_description: An Interactive Framework for LLM Failure Analysis
+suggested_hardware: t4-small
 ---
 
+# NeuralPostmortem
+
+An Interactive Framework for LLM Failure Analysis using Attribution Methods (Attn-LRP, CP-LRP, Gradient).
+
+## Features
+
+- **Model Loading**: Load HuggingFace models with optional 4-bit quantization
+- **Error Token Localization**: Identify erroneous tokens via LLM validators or manual truncation
+- **Input Attribution**: Visualize token-level attribution using Attn-LRP, CP-LRP, or vanilla gradients
+- **Perturbation Evaluation**: Evaluate attribution quality by zeroing out top-attributed tokens
+- **Circuit Visualization**: Build and visualize attribution graphs across model layers
+
+## Environment Variables (Optional)
+
+- `OPENAI_API_KEY`: OpenAI API key for LLM-based error token localization
+- `OPENAI_BASE_URL`: Custom OpenAI-compatible API base URL
+- `LLM_MODELS`: Comma-separated list of model names for validators (default: `gpt-4o-mini`)
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,653 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
+import uvicorn
+import os
+import sys
+import torch
+import json
+import logging
+import networkx as nx
+from networkx.readwrite import json_graph
+import numpy as np
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            # Safe handle checking for finite
+            f = float(obj)
+            return f if np.isfinite(f) else 0.0
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super(NumpyEncoder, self).default(obj)
+
+# Ensure backend can be imported
+PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))
+sys.path.insert(0, PROJECT_ROOT)
+
+from backend.models import ModelManager
+from backend.core import AttributionEngine
+from backend.circuit import CircuitAnalyzer
+from backend.error_token_location import ErrorTokenLocator
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import RedirectResponse, StreamingResponse
+from huggingface_hub import list_models, list_repo_refs
+
+app = FastAPI(title="NeuralPostmortem - Evaluation Backend (Attribution Comparison & Perturbation)")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount Frontend (Static Files)
+frontend_path = os.path.join(PROJECT_ROOT, 'frontend')
+if os.path.exists(frontend_path):
+    app.mount("/ui", StaticFiles(directory=frontend_path), name="ui")
+
+@app.get("/")
+async def read_root():
+    return RedirectResponse(url="/ui/index.html")
+
+# Global instances
+model_manager = ModelManager()
+attribution_engine = None # Initialize after model load
+error_token_locator = None # Initialize after model load
+
+# Caching for connection matrices to speed up slider interactions
+CACHED_CONNECTION_DATA = {
+    "config_hash": None,
+    "data": None
+}
+
+def get_config_hash(bp_config, layers):
+    try:
+        # Create a deterministic hash string
+        return json.dumps({
+            "bp": bp_config,
+            "layers": sorted(layers)
+        }, sort_keys=True)
+    except:
+        return None
+
+def unescape_string(text: str) -> str:
+    """
+    Safely unescape string with escape sequences like \\n, \\t, \\r, etc.
+
+    Args:
+        text: Input text that may contain escape sequences
+
+    Returns:
+        Text with escape sequences converted to actual characters
+    """
+    if not text:
+        return text
+
+    try:
+        # Try to decode escape sequences using unicode_escape
+        # This handles \n, \t, \r, \", \', \\, etc.
+        return text.encode('utf-8').decode('unicode_escape')
+    except Exception as e:
+        # Fallback to manual replacement if unicode_escape fails
+        logging.warning(f"unicode_escape failed, using manual replacement: {e}")
+        result = text
+        result = result.replace('\\n', '\n')
+        result = result.replace('\\t', '\t')
+        result = result.replace('\\r', '\r')
+        result = result.replace('\\"', '"')
+        result = result.replace("\\'", "'")
+        result = result.replace('\\\\', '\\')
+        return result
+
+
+# Pydantic models for inputs
+class LoadModelRequest(BaseModel):
+    model_path: str = "Qwen/Qwen3-0.6B"
+    quantization_4bit: bool = False  # Default to False to avoid bitsandbytes requirement
+    dtype: str = "float16" # float16, bfloat16, float32, auto
+    revision: Optional[str] = None
+    # LRP is no longer loaded at model load time
+
+class ComputeLogitsRequest(BaseModel):
+    prompt: str
+    is_append_bos: bool = True
+    topk: int = 10
+    extra_token_ids: Optional[List[int]] = None
+    extra_token_strs: Optional[List[str]] = None
+    capture_mid: bool = False  # Fine-grained attribution separation
+
+class BackpropConfig(BaseModel):
+    mode: str = "max_logit" # "max_logit" or "logit_diff"
+    strategy: Optional[str] = "by_topk_avg" # "demean", "by_topk_avg", "by_ref_token"
+    ref_token_id: Optional[int] = None
+    contrast_rank: Optional[int] = 2
+    k: Optional[int] = 10
+    node_threshold: Optional[float] = 0.01  # Threshold for computing node inter-connections
+
+class ComputeCircuitRequest(BaseModel):
+    # Configurations for backprop
+    backprop_config: BackpropConfig
+
+    # New Multi-Layer Field
+    layers: List[int]
+
+    # Pruning Params
+    pruning_mode: str = "by_per_layer_cum_mass_percentile"
+    top_p: float = 0.9
+    edge_threshold: float = 0.01 # Used if by_global_threshold
+
+
+class ComputeInputAttributionRequest(BaseModel):
+    target_token_id: int
+    contrast_token_id: Optional[int] = None
+    backprop_config: BackpropConfig
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 30
+    append_token_id: Optional[int] = None
+
+class LocateErrorTokenRequest(BaseModel):
+    prompt: str
+    completion: str
+    ground_truth: Optional[str] = None
+    validators: Optional[List[str]] = None
+    use_llm: bool = True
+    manual_chunks: Optional[List[str]] = None
+
+class EnableLRPRequest(BaseModel):
+    lrp_rule: str = "Attn-LRP"  # "Attn-LRP", "CP-LRP", or "Gradient"
+    capture_mid: bool = False
+
+class ComputePerturbationRequest(BaseModel):
+    attribution_scores: List[float]
+    k_values: List[int] = [1, 3, 5, 10]
+    target_token_id: int
+
+class ComputePerturbationManualRequest(BaseModel):
+    perturb_indices: List[int]
+    target_token_id: int
+
+@app.get("/api/list_hf_models")
+async def list_hf_models(series: str = "Qwen2"):
+    """
+    List models from HuggingFace Hub filtered by series/author.
+    """
+    try:
+        if series.lower() == "qwen2":
+            models = list(list_models(author="Qwen", search="Qwen2", filter="text-generation", sort="downloads", direction=-1, limit=50))
+            return {"models": [m.id for m in models]}
+
+        elif series.lower() == "qwen3":
+            models = list(list_models(author="Qwen", search="Qwen3", filter="text-generation", sort="downloads", direction=-1, limit=50))
+            return {"models": [m.id for m in models]}
+
+        elif series.lower() == "olmo3":
+            models = list(list_models(author="allenai", search="Olmo-3", filter="text-generation", sort="downloads", direction=-1, limit=50))
+            return {"models": [m.id for m in models]}
+
+        elif series.lower() == "olmo":
+            models = list(list_models(author="allenai", search="OLMo", filter="text-generation", sort="downloads", direction=-1, limit=50))
+            return {"models": [m.id for m in models]}
+
+        elif series.lower() == "qwen":
+            models = list(list_models(author="Qwen", filter="text-generation", sort="downloads", direction=-1, limit=50))
+            return {"models": [m.id for m in models]}
+
+        # Generic fallback
+        models = list(list_models(search=series, filter="text-generation", sort="downloads", direction=-1, limit=20))
+        return {"models": [m.id for m in models]}
+
+    except Exception as e:
+        print(f"Error listing models: {e}")
+        # Return fallback/hardcoded list if offline
+        if series.lower() == "qwen2":
+             return {"models": ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen2.5-1.5B-Instruct", "Qwen/Qwen2.5-3B-Instruct", "Qwen/Qwen2.5-7B-Instruct", "Qwen/Qwen2-0.5B", "Qwen/Qwen2-1.5B", "Qwen/Qwen2-7B"]}
+        elif series.lower() == "qwen3":
+             return {"models": ["Qwen/Qwen3-0.6B"]}
+        elif series.lower() == "qwen":
+            return {"models": ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen/Qwen3-0.6B"]}
+        elif series.lower() == "olmo3":
+            return {"models": ["allenai/Olmo-3-7B-Think"]}
+        elif series.lower() == "olmo":
+            return {"models": ["allenai/OLMo-7B", "allenai/OLMo-1B-0724", "allenai/Olmo-3-7B-Think"]}
+        return {"models": [], "error": str(e)}
+
+@app.get("/api/list_model_revisions")
+async def list_model_revisions(model_id: str):
+    """
+    List git branches/refs for a model.
+    """
+    try:
+        refs = list_repo_refs(model_id)
+        branches = [b.name for b in refs.branches]
+        tags = [t.name for t in refs.tags]
+        return {"branches": branches, "tags": tags}
+    except Exception as e:
+        print(f"Error listing revisions for {model_id}: {e}")
+        return {"branches": [], "tags": [], "error": str(e)}
+
+@app.post("/api/cleanup")
+async def cleanup_memory():
+    global attribution_engine
+    if attribution_engine:
+        attribution_engine.reset()
+    else:
+        # even if no engine, try to clear cache
+        torch.cuda.empty_cache()
+
+    import gc
+    gc.collect()
+
+    return {"status": "success", "message": "Memory cleanup complete"}
+
+@app.post("/api/generate")
+async def generate_continuation(request: GenerateRequest):
+    if not model_manager.model:
+        raise HTTPException(status_code=400, detail="Model not loaded")
+
+    tokenizer = model_manager.tokenizer
+    model = model_manager.model
+    device = model_manager.device
+
+    try:
+        # Unescape special characters in prompt
+        prompt = unescape_string(request.prompt)
+
+        # Switch to eval for generation
+        was_training = model.training
+        model.eval()
+
+        # Encode prompt
+        input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
+
+        # Append token if requested
+        if request.append_token_id is not None:
+            token_tensor = torch.tensor([[request.append_token_id]], device=device)
+            input_ids = torch.cat([input_ids, token_tensor], dim=1)
+
+        with torch.no_grad():
+            output_ids = model.generate(
+                input_ids,
+                max_new_tokens=request.max_new_tokens,
+                do_sample=False,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+        new_token_ids = output_ids[0][input_ids.shape[1]:]
+        generated_text = tokenizer.decode(new_token_ids, skip_special_tokens=True)
+
+        # Restore training mode
+        if was_training:
+            model.train()
+
+        return {"generated_text": generated_text}
+
+    except Exception as e:
+        if model_manager.model and was_training:
+            model_manager.model.train()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/locate_err_token")
+async def locate_error_token_endpoint(request: LocateErrorTokenRequest):
+    """
+    Locate the error token in a completion using multiple LLM validators
+    """
+    global error_token_locator
+
+    if not error_token_locator:
+        raise HTTPException(status_code=400, detail="Model not loaded. Please call /api/load_model first.")
+
+    try:
+        # Unescape special characters in prompt, completion, and ground_truth
+        prompt = unescape_string(request.prompt)
+        completion = unescape_string(request.completion)
+        ground_truth = unescape_string(request.ground_truth) if request.ground_truth else None
+
+        # Call the error token locator
+        result = error_token_locator.locate_error_token(
+            prompt=prompt,
+            completion=completion,
+            ground_truth=ground_truth,
+            validators=request.validators,
+            use_llm=request.use_llm,
+            manual_chunks=request.manual_chunks
+        )
+
+        if result["status"] == "error":
+            raise HTTPException(status_code=500, detail=result.get("message", "Unknown error"))
+
+        return {
+            "status": "success",
+            "truncated_text": result["truncated_text"],
+            "explanation": result["explanation"],
+            "error_token_index": result.get("error_token_index", -1),
+            "vote_details": result.get("vote_details", {})
+        }
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/load_model")
+async def load_model(request: LoadModelRequest):
+    global attribution_engine
+    global error_token_locator
+    try:
+        # Load model without LRP (LRP will be enabled when needed)
+        model_name = model_manager.load_model(
+            request.model_path,
+            request.quantization_4bit,
+            dtype=request.dtype,
+            revision=request.revision,
+            lrp_rule=None  # Don't load LRP yet
+        )
+        attribution_engine = AttributionEngine(model_manager)
+        error_token_locator = ErrorTokenLocator(model_manager.model, model_manager.tokenizer)
+
+        # Get Num Layers
+        n_layers = 28 # Default for Qwen 0.5B
+        try:
+             # Try access config
+             if hasattr(model_manager.model, 'config'):
+                 n_layers = getattr(model_manager.model.config, 'num_hidden_layers', 28)
+        except:
+             pass
+
+        # Get vocabulary size
+        vocab_size = len(model_manager.tokenizer)
+
+        return {
+            "status": "success",
+            "message": f"Model {model_name} loaded successfully",
+            "num_layers": n_layers,
+            "vocab_size": vocab_size
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/enable_lrp")
+async def enable_lrp(request: EnableLRPRequest):
+    """
+    Enable LRP functionality on the loaded model.
+    This should be called before computing attribution or circuits.
+    For "Gradient" rule, the model is loaded WITHOUT LRP patches (vanilla gradient).
+    """
+    if not model_manager.model:
+        raise HTTPException(status_code=400, detail="Model not loaded. Please call /api/load_model first.")
+
+    try:
+        # For Gradient method, load model without LRP patches
+        lrp_rule_for_model = None if request.lrp_rule == "Gradient" else request.lrp_rule
+
+        # Reload model with appropriate LRP setting
+        model_name = model_manager.load_model(
+            model_path=model_manager.current_model_path,
+            quantization_4bit=model_manager.current_quantization,
+            dtype=model_manager.current_dtype,
+            revision=model_manager.current_revision,
+            lrp_rule=lrp_rule_for_model
+        )
+
+        # Reinitialize attribution engine
+        global attribution_engine
+        attribution_engine = AttributionEngine(model_manager)
+
+        return {
+            "status": "success",
+            "message": f"Attribution method ({request.lrp_rule}) enabled successfully",
+            "lrp_rule": request.lrp_rule,
+            "capture_mid": request.capture_mid
+        }
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_logits")
+async def compute_logits(request: ComputeLogitsRequest):
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded. Please call /api/load_model first.")
+
+    try:
+        # Unescape special characters in prompt
+        prompt = unescape_string(request.prompt)
+
+        topk_data, _, input_tokens = attribution_engine.compute_logits(
+            prompt=prompt,
+            is_append_bos=request.is_append_bos,
+            topk=request.topk,
+            extra_token_ids=request.extra_token_ids,
+            extra_token_strs=request.extra_token_strs,
+            capture_mid=request.capture_mid
+        )
+
+        # Convert simple string list input_tokens to list of objects for frontend consistency
+        token_objs = [{"token_str": t, "token_id": i} for i, t in enumerate(input_tokens)]
+
+        return {"data": topk_data, "tokens": token_objs}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_input_attribution")
+async def compute_input_attribution_endpoint(request: ComputeInputAttributionRequest):
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded.")
+
+    # Auto-enable LRP if not already enabled
+    if not model_manager.current_lrp_rule:
+        logger.info("LRP not enabled yet - auto-enabling with default rule 'Attn-LRP'...")
+        try:
+            model_name = model_manager.load_model(
+                model_path=model_manager.current_model_path,
+                quantization_4bit=model_manager.current_quantization,
+                dtype=model_manager.current_dtype,
+                revision=model_manager.current_revision,
+                lrp_rule="Attn-LRP"
+            )
+            attribution_engine = AttributionEngine(model_manager)
+            logger.info(f"Auto-enabled LRP with Attn-LRP rule on {model_name}")
+        except Exception as e:
+            logger.error(f"Failed to auto-enable LRP: {e}")
+            raise HTTPException(
+                status_code=400,
+                detail="LRP not enabled and auto-enable failed. Please call /api/enable_lrp before computing attribution."
+            )
+
+    try:
+        if attribution_engine.outputs is None:
+             raise HTTPException(status_code=400, detail="No forward pass found. Run compute_logits first.")
+
+        # Inject target token ID into backprop config
+        bp_config = request.backprop_config.dict()
+        bp_config["target_token_id"] = request.target_token_id
+
+        relevance = attribution_engine.compute_input_attribution(bp_config)
+        return {"relevance": relevance}
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_input_attribution_gradient")
+async def compute_input_attribution_gradient_endpoint(request: ComputeInputAttributionRequest):
+    """
+    Compute input attribution using vanilla gradient method (Input * Gradient).
+    Does NOT require LRP to be enabled - uses standard PyTorch autograd.
+    """
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded.")
+
+    try:
+        if attribution_engine.outputs is None:
+            raise HTTPException(status_code=400, detail="No forward pass found. Run compute_logits first.")
+
+        # Inject target token ID into backprop config
+        bp_config = request.backprop_config.dict()
+        bp_config["target_token_id"] = request.target_token_id
+
+        relevance = attribution_engine.compute_input_attribution_gradient(bp_config)
+        return {"relevance": relevance}
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_perturbation")
+async def compute_perturbation_endpoint(request: ComputePerturbationRequest):
+    """
+    Evaluate attribution quality by perturbing top-attributed tokens.
+    Zero out top-k most attributed tokens and check if the error is fixed.
+    """
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded.")
+
+    try:
+        if attribution_engine.input_ids is None or attribution_engine.input_embeddings is None:
+            raise HTTPException(status_code=400, detail="No forward pass found. Run compute_logits first.")
+
+        results = attribution_engine.compute_perturbation_eval(
+            attribution_scores=request.attribution_scores,
+            k_values=request.k_values,
+            target_token_id=request.target_token_id
+        )
+        return {"results": results}
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_perturbation_manual")
+async def compute_perturbation_manual_endpoint(request: ComputePerturbationManualRequest):
+    """
+    Evaluate attribution by perturbing manually selected token positions.
+    Zero out the specified token embeddings and check if the error is fixed.
+    """
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded.")
+
+    try:
+        if attribution_engine.input_ids is None or attribution_engine.input_embeddings is None:
+            raise HTTPException(status_code=400, detail="No forward pass found. Run compute_logits first.")
+
+        result = attribution_engine.compute_perturbation_manual(
+            perturb_indices=request.perturb_indices,
+            target_token_id=request.target_token_id
+        )
+        return {"result": result}
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/api/compute_circuit")
+async def compute_circuit(request: ComputeCircuitRequest):
+    global attribution_engine
+    if not attribution_engine:
+        raise HTTPException(status_code=400, detail="Model not loaded.")
+
+    # Auto-enable LRP if not already enabled
+    if not model_manager.current_lrp_rule:
+        logger.info("LRP not enabled yet - auto-enabling with default rule 'Attn-LRP' for circuit analysis...")
+        try:
+            model_name = model_manager.load_model(
+                model_path=model_manager.current_model_path,
+                quantization_4bit=model_manager.current_quantization,
+                dtype=model_manager.current_dtype,
+                revision=model_manager.current_revision,
+                lrp_rule="Attn-LRP"
+            )
+            attribution_engine = AttributionEngine(model_manager)
+            logger.info(f"Auto-enabled LRP with Attn-LRP rule on {model_name}")
+        except Exception as e:
+            logger.error(f"Failed to auto-enable LRP: {e}")
+            raise HTTPException(
+                status_code=400,
+                detail="LRP not enabled and auto-enable failed. Please call /api/enable_lrp before computing circuits."
+            )
+
+    if attribution_engine.outputs is None:
+         raise HTTPException(status_code=400, detail="No forward pass found. Run compute_logits first.")
+
+    async def generate_response():
+        try:
+            # Step 1: Run Backward Pass
+            yield json.dumps({"type": "progress", "msg": "Initiating Backward Pass...", "percent": 0}) + "\n"
+
+            # Use CircuitAnalyzer
+            analyzer = CircuitAnalyzer(attribution_engine)
+
+            bp_config = request.backprop_config.dict()
+
+            # We explicitly run backward pass first (though build_graph does it, we want to emit progress)
+            # Check Cache
+            current_hash = get_config_hash(bp_config, request.layers)
+            connection_data = None
+
+            if CACHED_CONNECTION_DATA["config_hash"] == current_hash and CACHED_CONNECTION_DATA["data"] is not None:
+                 yield json.dumps({"type": "progress", "msg": "Using Cached Matrices (Fast)...", "percent": 50}) + "\n"
+                 connection_data = CACHED_CONNECTION_DATA["data"]
+            else:
+                 yield json.dumps({"type": "progress", "msg": "Computing Circuit (This may take a moment)...", "percent": 20}) + "\n"
+                 # Run the heavy lifting
+                 connection_data = analyzer.compute_connection_matrices(bp_config, sorted(request.layers))
+
+                 # Update Cache
+                 CACHED_CONNECTION_DATA["config_hash"] = current_hash
+                 CACHED_CONNECTION_DATA["data"] = connection_data
+
+            yield json.dumps({"type": "progress", "msg": "Pruning & Building Graph...", "percent": 80}) + "\n"
+
+            G, pruning_details = analyzer.build_graph_from_matrices(
+                connection_data,
+                edge_rel_threshold=request.edge_threshold,
+                pruning_mode=request.pruning_mode,
+                top_p=request.top_p
+            )
+
+            yield json.dumps({"type": "progress", "msg": "Graph Constructed. Serializing...", "percent": 90}) + "\n"
+
+            # Serialize Graph
+            graph_data = nx.node_link_data(G)
+
+            yield json.dumps({
+                "type": "graph_data",
+                "graph": graph_data,
+                "pruning_details": pruning_details
+            }, cls=NumpyEncoder) + "\n"
+
+            yield json.dumps({"type": "progress", "msg": "Complete!", "percent": 100}) + "\n"
+            yield json.dumps({"type": "complete"}) + "\n"
+
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            yield json.dumps({"type": "error", "msg": str(e)}) + "\n"
+
+    return StreamingResponse(generate_response(), media_type="application/x-ndjson")
+
+# Return empty datasets (dataset loading removed for HuggingFace Space deployment)
+@app.get("/api/datasets")
+async def get_datasets():
+    return {"datasets": []}
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run(app, host="0.0.0.0", port=port)

backend/__init__.py

@@ -0,0 +1,2 @@
+from .models import ModelManager
+from .core import AttributionEngine

backend/batch_config.py

@@ -0,0 +1,44 @@
+import torch
+
+# Configuration for Batch Chunk Sizes based on Model Size and Dtype
+# Reference: 4B model with bf16 uses batch_chunk_size = 64
+
+# Mapping: (Min_Params_Billions, Max_Params_Billions) -> Recommended Batch Size (BF16/FP16)
+BATCH_CHUNK_MAPPING = {
+    (0.0, 1.0): 256,   # For 0.6B and similar
+    (1.0, 3.0): 128,   # For 1.7B, 2B, 3B
+    (3.0, 6.0): 8,    # For 4B, 6B
+    (6.0, 12.0): 32,   # For 7B, 8B, 10B
+    (12.0, 25.0): 16,  # For 14B, 20B
+    (25.0, 1000.0): 8  # For 32B+, 70B
+}
+
+def get_batch_chunk_size(model_params_count, model_dtype):
+    """
+    Determine appropriate batch chunk size based on parameter count and dtype.
+
+    Args:
+        model_params_count (int): Total number of parameters in the model.
+        model_dtype (torch.dtype): The data type used for computation (activations).
+
+    Returns:
+        int: Recommended batch chunk size.
+    """
+    # Convert to Billions
+    params_billions = model_params_count / 1e9
+
+    # Default fallback
+    chunk_size = 32
+
+    # Lookup in mapping
+    for (min_b, max_b), size in BATCH_CHUNK_MAPPING.items():
+        if min_b <= params_billions < max_b:
+            chunk_size = size
+            break
+
+    # Scale by Dtype
+    # If using float32, activations take 2x memory compared to bf16/fp16
+    if model_dtype == torch.float32:
+        chunk_size = max(1, chunk_size // 2)
+
+    return chunk_size

backend/circuit.py

@@ -0,0 +1,252 @@
+import networkx as nx
+import torch
+import gc
+import numpy as np
+import pandas as pd
+
+class CircuitAnalyzer:
+    """
+    A class to build and analyze attribution circuits using an AttributionEngine.
+    """
+    def __init__(self, attribution_engine):
+        """
+        Args:
+            attribution_engine: An instance of AttributionEngine.
+        """
+        self.engine = attribution_engine
+
+    def compute_connection_matrices(self, backprop_config, layers=None):
+        """
+        Step 1: Compute dense connection matrices for all layer transitions.
+        Returns a list of dictionaries containing matrix data for each transition.
+        """
+        model = self.engine.manager.get_model()
+        model_layers = model.model.layers
+        n_layers = len(model_layers)
+
+        if layers is None:
+            # Auto-detect if mid activations are available (captured in compute_logits)
+            check_layer = model_layers[0]
+            has_mid = False
+            # Check for mid_activation attribute on the normative layer where we hook
+            if hasattr(check_layer, 'post_attention_layernorm') and hasattr(check_layer.post_attention_layernorm, 'mid_activation'):
+                has_mid = True
+                print("Auto-detected mid activations: Using fine-grained circuit (Attn/MLP separation).")
+
+            nodes = [-1]
+            for i in range(n_layers):
+                if has_mid:
+                    nodes.append((i, 'mid'))
+                nodes.append(i) # integers imply 'post'
+
+            # Sort/Sequence is crucial.
+            # [-1, (0,'mid'), 0, (1,'mid'), 1, ...]
+            layer_pairs = list(zip(nodes[:-1], nodes[1:]))
+        else:
+            # Use provided layers list
+            # Ensure it's sorted? User responsibility if custom list.
+            nodes = layers
+            layer_pairs = list(zip(nodes[:-1], nodes[1:]))
+
+        # 1. Run Backward Pass
+        print("Running backward pass...")
+        self.engine.run_backward_pass(backprop_config)
+
+        connection_data = []
+
+        print(f"Computing matrices for {len(nodes)} nodes ({len(layer_pairs)} transitions)...")
+
+        for i, (src_layer, tgt_layer) in enumerate(layer_pairs):
+            print(f"Computing transition: {src_layer} -> {tgt_layer}...")
+
+            gen = self.engine.compute_connection_matrix_gen(
+                source=src_layer,
+                target=tgt_layer
+            )
+
+            final_res = None
+            for item in gen:
+                if item['type'] == 'result':
+                    final_res = item['payload']
+
+            if final_res:
+                # Store the dense matrix and relevance vectors
+                connection_data.append({
+                    'src_layer': src_layer,
+                    'tgt_layer': tgt_layer,
+                    'matrix': final_res['matrix'], # Dense numpy array
+                    'real_source_rel': final_res['real_source_rel'],
+                    'real_target_rel': final_res['real_target_rel']
+                })
+
+                # Basic cleanup of engine internals, but we keep the matrix in connection_data
+                torch.cuda.empty_cache()
+
+        return connection_data
+
+    def build_graph_from_matrices(self, connection_data, edge_rel_threshold=0.01, pruning_mode="by_global_threshold", top_p=0.9):
+        """
+        Step 2: Prune based on threshold and build NetworkX graph.
+
+        Args:
+            connection_data: List of dictionaries containing matrix data.
+            edge_rel_threshold: Threshold for 'by_global_threshold' mode.
+            pruning_mode: Pruning strategy. Options:
+                          1. "by_global_threshold" (default): Prune globally using edge_rel_threshold.
+                          2. "by_per_layer_cum_mass_percentile": Prune per layer to keep top_p mass.
+            top_p: The cumulative mass percentile (0.0-1.0) for "by_per_layer_cum_mass_percentile". Default 0.9.
+
+        Returns:
+            G: The built NetworkX graph.
+            pruning_details: A list of dictionaries containing pruning stats per layer pair.
+        """
+        G = nx.DiGraph()
+        print(f"Building graph from {len(connection_data)} transitions. Mode: {pruning_mode}...")
+
+        pruning_details = []
+
+        for data in connection_data:
+            src_layer = data['src_layer']
+            tgt_layer = data['tgt_layer']
+            matrix = data['matrix']
+            real_source_rel = data['real_source_rel']
+            real_target_rel = data['real_target_rel']
+            abs_matrix = np.abs(matrix)
+
+            # Count only active edges (non-zero) as total_elements for percentage calculation
+            nonzero_mask = abs_matrix > 1e-9
+            total_elements = np.sum(nonzero_mask)
+
+            used_threshold = edge_rel_threshold
+
+            # Determine mask for edges to keep based on mode
+            if pruning_mode == "by_per_layer_cum_mass_percentile":
+                # Dynamic thresholding per layer
+                flattened = np.sort(abs_matrix.flatten())[::-1]
+                total_mass = flattened.sum()
+
+                dynamic_threshold = 1.0 # Default High if empty
+
+                if total_mass > 1e-12:
+                     cumsum = np.cumsum(flattened)
+                     cutoff_mass = total_mass * top_p
+                     # Find first index where cumsum >= cutoff_mass
+                     # searchsorted returns insertion point index i such that a[i-1] < v <= a[i]
+                     # If we want elements up to index k such sum(0..k) >= target
+                     cutoff_idx = np.searchsorted(cumsum, cutoff_mass)
+
+                     if cutoff_idx >= len(flattened):
+                         cutoff_idx = len(flattened) - 1
+
+                     dynamic_threshold = flattened[cutoff_idx]
+                     # Ensure we do not include effective zeros
+                     if dynamic_threshold < 1e-9:
+                         dynamic_threshold = 1e-9
+
+                used_threshold = dynamic_threshold
+                rows, cols = np.where(abs_matrix >= dynamic_threshold)
+            else:
+                # "by_global_threshold"
+                rows, cols = np.where(abs_matrix > edge_rel_threshold)
+
+            # Record pruning details
+            num_kept = len(rows)
+            percentage = (num_kept / total_elements * 100) if total_elements > 0 else 0
+
+            pruning_details.append({
+                'src_layer': src_layer,
+                'tgt_layer': tgt_layer,
+                'threshold': float(used_threshold),
+                'kept_edges': int(num_kept),
+                'total_edges': int(total_elements),
+                'kept_percentage': percentage
+            })
+
+            # Add Source Nodes
+            for t_idx in np.where(np.abs(real_source_rel) > 0)[0]:
+                    src_node_id = (src_layer, t_idx)
+                    if not G.has_node(src_node_id):
+                        G.add_node(src_node_id, layer=src_layer, token=t_idx, relevance=real_source_rel[t_idx])
+
+            # Add Target Nodes
+            for t_idx in np.where(np.abs(real_target_rel) > 0)[0]:
+                    tgt_node_id = (tgt_layer, t_idx)
+                    if not G.has_node(tgt_node_id):
+                        G.add_node(tgt_node_id, layer=tgt_layer, token=t_idx, relevance=real_target_rel[t_idx])
+
+            # Add Edges (Sparse)
+            # rows, cols calculated above
+            weights = matrix[rows, cols]
+
+            edges_to_add = []
+            for r, c, w in zip(rows, cols, weights):
+                # Edge direction: Source (col/c) -> Target (row/r)
+                u = (src_layer, c)
+                v = (tgt_layer, r)
+                # Ensure nodes exist (redundant safety check)
+                if not G.has_node(u): G.add_node(u, layer=src_layer, token=c, relevance=real_source_rel[c])
+                if not G.has_node(v): G.add_node(v, layer=tgt_layer, token=r, relevance=real_target_rel[r])
+
+                edges_to_add.append((u, v, {'weight': float(w)}))
+
+            G.add_edges_from(edges_to_add)
+
+        print(f"Graph construction complete. Nodes: {G.number_of_nodes()}, Edges: {G.number_of_edges()}")
+        return G, pruning_details
+
+    def build_graph(self, backprop_config, layers=None, edge_rel_threshold=0.01, pruning_mode="by_global_threshold", top_p=0.9):
+        """
+        Wrapper that performs both steps: compute matrices and build graph.
+        """
+        connection_data = self.compute_connection_matrices(backprop_config, layers)
+        return self.build_graph_from_matrices(connection_data, edge_rel_threshold, pruning_mode=pruning_mode, top_p=top_p)
+
+    def get_connected_subgraph(self, G, target_node=None):
+        """
+        Extracts the subgraph connected to the target node (backward reachability).
+        If target_node is None, it tries to infer the last token at the last layer.
+
+        Args:
+            G: The full attribution graph
+            target_node: tuple (layer, token_idx) of the target.
+
+        Returns:
+            (subgraph, target_node): The connected subgraph and the resolved target node.
+        """
+        # 1. Identify Target Node
+        if target_node is None:
+            if G.number_of_nodes() == 0:
+                print("Graph is empty.")
+                return None, None
+
+            # Infer: Max layer
+            try:
+                max_layer = max([n[0] for n in G.nodes()])
+            except ValueError:
+                 print("Error finding max layer.")
+                 return None, None
+
+            # Check nodes in max_layer
+            nodes_in_last = [n for n in G.nodes() if n[0] == max_layer]
+            if not nodes_in_last:
+                 print(f"No nodes found in layer {max_layer}.")
+                 return None, None
+
+            # Max token index
+            max_token = max([n[1] for n in nodes_in_last])
+            target_node = (max_layer, max_token)
+
+        print(f"Extracting connected component for Target Node: {target_node}")
+
+        if not G.has_node(target_node):
+            print(f"Target node {target_node} not found in graph (maybe thresholded out?).")
+            return None, target_node
+
+        # 2. Get Ancestors (Backward Reachability)
+        ancestors = nx.ancestors(G, target_node)
+        ancestors.add(target_node) # Include self
+
+        subgraph = G.subgraph(ancestors).copy()
+        print(f"Connected Subgraph: {subgraph.number_of_nodes()} nodes, {subgraph.number_of_edges()} edges")
+
+        return subgraph, target_node

backend/core.py

@@ -0,0 +1,812 @@
+import torch
+import numpy as np
+from .batch_config import get_batch_chunk_size
+
+class AttributionEngine:
+    def __init__(self, model_manager):
+        self.manager = model_manager
+        self.hook_handles = []
+        self.outputs = None
+        self.input_ids = None
+
+    def _parse_node(self, node):
+        if isinstance(node, int):
+            return node, 'post'
+        if isinstance(node, (tuple, list)) and len(node) == 2:
+            return node[0], node[1]
+        raise ValueError(f"Invalid node format: {node}")
+
+    def _forward_part1(self, layer_module, hidden_states, position_embeddings=None, attention_mask=None):
+        """
+        Executes: Norm -> Attn -> Residual Add
+        Returns: resid_mid
+        """
+        return self.manager.decomposer.forward_part1(layer_module, hidden_states, position_embeddings, attention_mask)
+
+    def _forward_part2(self, layer_module, hidden_states):
+        """
+        Executes: Norm -> MLP -> Residual Add
+        Returns: resid_post
+        """
+        return self.manager.decomposer.forward_part2(layer_module, hidden_states)
+
+    def _hook_hidden_activation(self, module, input, output):
+        """
+        Hook to save activation and enable gradient retention.
+        """
+        if isinstance(output, tuple):
+            output = output[0]
+
+        # Save output to the module for later access
+        module.output = output
+        if module.output.requires_grad:
+            module.output.retain_grad()
+
+    def _hook_mid_activation(self, module, input, output):
+        """
+        Hook to capture resid_mid at the input of post_attention_layernorm.
+        """
+        # input is a tuple (tensor,)
+        val = input[0]
+        # Attach to the module (which is the Norm layer)
+        module.mid_activation = val
+        if module.mid_activation.requires_grad:
+            module.mid_activation.retain_grad()
+
+    def register_hooks(self, capture_mid=False):
+        """
+        Register forward hooks on all model layers.
+        """
+        self.remove_hooks() # Clear existing
+        model = self.manager.get_model()
+        if not model:
+            raise ValueError("Model not loaded yet.")
+
+        # Assuming generic structure model.model.layers (common in HF Qwen, Llama, etc.)
+        if hasattr(model, "model") and hasattr(model.model, "layers"):
+            layers = model.model.layers
+        elif hasattr(model, "layers"): # Some archs
+            layers = model.layers
+        else:
+            layers = []
+
+        for layer in layers:
+            # Hook Output (Resid Post)
+            handle = layer.register_forward_hook(self._hook_hidden_activation)
+            self.hook_handles.append(handle)
+
+            # Hook Mid (Resid Mid) - Pre-hook on Post-Attn Norm
+            if capture_mid:
+                # Use decomposer to find the correct module for mid activation
+                mid_module = self.manager.decomposer.get_mid_activation_module(layer)
+
+                if mid_module:
+                    # Use forward hook to get input?
+                    # register_forward_hook receives (module, input, output)
+                    # input is (resid_mid,)
+                    handle_mid = mid_module.register_forward_hook(self._hook_mid_activation)
+                    self.hook_handles.append(handle_mid)
+                else:
+                    print(f"Warning: Could not identify mid-activation module for layer {layer}. Skipping mid hook.")
+
+    def remove_hooks(self):
+        for handle in self.hook_handles:
+            handle.remove()
+        self.hook_handles = []
+
+    def reset(self):
+        """
+        Clears all internal state and specific temporary data from model layers.
+        """
+        self.remove_hooks()
+        self.outputs = None
+        self.input_ids = None
+
+        # Manually clear output tensors attached to layers to free graph
+        model = self.manager.get_model()
+        if model:
+            layers = []
+            if hasattr(model, "model") and hasattr(model.model, "layers"):
+                layers = model.model.layers
+            elif hasattr(model, "layers"):
+                layers = model.layers
+
+            for layer in layers:
+                if hasattr(layer, 'output'):
+                    del layer.output
+                if hasattr(layer, 'post_attention_layernorm') and hasattr(layer.post_attention_layernorm, 'mid_activation'):
+                    del layer.post_attention_layernorm.mid_activation
+
+        torch.cuda.empty_cache()
+
+
+    def compute_logits(self, prompt, is_append_bos=False, topk=10, extra_token_ids=None, extra_token_strs=None, capture_mid=False):
+        """
+        Section 1: Forward pass to get logits and top-k predictions.
+        """
+        model = self.manager.get_model()
+        tokenizer = self.manager.get_tokenizer()
+
+        self.register_hooks(capture_mid=capture_mid)
+
+        # Prepare input
+        # We tokenize with add_special_tokens=False to manually control the BOS/Start token
+        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
+        input_ids = inputs.input_ids.to(model.device)
+
+        if is_append_bos:
+            # 1. Try explicit BOS
+            bos_id = tokenizer.bos_token_id
+
+            # 2. Try CLS (BERT-like)
+            if bos_id is None:
+                bos_id = tokenizer.cls_token_id
+
+            # 3. Fallback: EOS (Often used as BOS/Separator in Llama/decoder-only models if BOS is missing)
+            if bos_id is None:
+                bos_id = tokenizer.eos_token_id
+
+            if bos_id is not None:
+                prefix = torch.tensor([[bos_id]], device=model.device)
+                input_ids = torch.cat([prefix, input_ids], dim=1)
+                print(f"Appended start token ID: {bos_id}")
+            else:
+                print("Warning: Append BOS requested but no suitable start token (BOS/CLS/EOS) found.")
+
+        self.input_ids = input_ids
+
+        # Embedding with gradients required for LRP base
+        # We detach and enable gradients so we can compute attribution w.r.t input embeddings
+        # even if the model is frozen/quantized.
+        self.input_embeddings = model.get_input_embeddings()(self.input_ids).detach()
+        self.input_embeddings.requires_grad_(True)
+
+        # Forward pass
+        # output_hidden_states=True is crucial for some LRP methods,
+        # though we use hooks for "efficient" method mostly.
+        self.outputs = model(
+            inputs_embeds=self.input_embeddings,
+            use_cache=False
+        )
+
+        output_logits = self.outputs.logits
+        last_logits = output_logits[0, -1, :]
+
+        # Get Top-K
+        sorted_logits, sorted_indices = torch.sort(last_logits, dim=-1, descending=True)
+
+        # Formatted output
+        topk_data = []
+        for i in range(topk):
+            idx = sorted_indices[i].item()
+            token_str = tokenizer.decode([idx])
+            logit_val = sorted_logits[i].item()
+            topk_data.append({
+                "rank": i + 1,
+                "token_id": idx,
+                "token_str": token_str,
+                "logit": logit_val
+            })
+
+        # Handle Extra Tokens (if requested)
+        if extra_token_ids or extra_token_strs:
+            # Helper to find rank
+            def get_rank(val, sorted_vals):
+                # tensor search for rank
+                # val is float, sorted_vals is tensor
+                # find first index where sorted_vals < val is NOT true?
+                # sorted_vals is descending
+                # we want count of items > val
+                return (sorted_vals > val).sum().item() + 1
+
+            processed_ids = set()
+
+            # Process IDs
+            if extra_token_ids:
+                for tid in extra_token_ids:
+                    if tid < 0 or tid >= len(last_logits): continue
+                    if tid in processed_ids: continue
+
+                    logit_val = last_logits[tid].item()
+                    rank = get_rank(logit_val, sorted_logits)
+                    token_str = tokenizer.decode([tid])
+
+                    topk_data.append({
+                        "rank": rank,
+                        "token_id": tid,
+                        "token_str": token_str,
+                        "logit": logit_val,
+                        "is_extra": True
+                    })
+                    processed_ids.add(tid)
+
+            # Process Strings
+            if extra_token_strs:
+                print(f"DEBUG: Processing extra strs: {extra_token_strs}")
+                for tstr in extra_token_strs:
+                    # Encode
+                    try:
+                        # Ensure we get list of ints
+                        encoded = tokenizer.encode(tstr, add_special_tokens=False)
+                        print(f"DEBUG: Encoded '{tstr}' -> {encoded} (Type: {type(encoded)})")
+
+                        if hasattr(encoded, 'tolist'): encoded = encoded.tolist()
+
+                        if len(encoded) == 0:
+                            print(f"DEBUG: Empty encoding for '{tstr}'")
+                            continue
+
+                        # Take first token
+                        tid = encoded[0]
+                        print(f"DEBUG: Using TID {tid} for '{tstr}'")
+
+                        if tid in processed_ids:
+                            print(f"DEBUG: TID {tid} already processed")
+                            continue
+
+                        logit_val = last_logits[tid].item()
+                        rank = get_rank(logit_val, sorted_logits)
+                        real_str = tokenizer.decode([tid])
+
+                        print(f"DEBUG: Added extra token: {real_str} (ID: {tid}, Rank: {rank})")
+
+                        topk_data.append({
+                            "rank": rank,
+                            "token_id": tid,
+                            "token_str": real_str,
+                            "logit": logit_val,
+                            "is_extra": True
+                        })
+                        processed_ids.add(tid)
+                    except Exception as e:
+                        print(f"DEBUG: Error processing extra str '{tstr}': {e}")
+                        import traceback
+                        traceback.print_exc()
+
+        # Sort combined data by rank for display consistency?
+        # Or keep extras at the bottom? User request: "add more tokens in the existing top-50 table"
+        # If we sort, they mix in. If they are rank 1000, they go to bottom.
+        # But if they are rank 5 (and we showed top 10), they mix in.
+        # Let's sort.
+        topk_data.sort(key=lambda x: x['rank'])
+
+        if self.input_ids is None:
+             raise ValueError("Input IDs not found. Ensure compute_logits was run.")
+
+        # Get input tokens for visualization
+        # Robust token reconstruction ensuring spaces are preserved
+        input_tokens = []
+        # convert_ids_to_tokens usually preserves the special characters (like Ġ or  )
+        raw_tokens = tokenizer.convert_ids_to_tokens(self.input_ids[0])
+
+        for t in raw_tokens:
+            # Handle bytes (common in tiktoken-based tokenizers like Qwen)
+            if isinstance(t, bytes):
+                try:
+                    t = t.decode('utf-8')
+                except:
+                    # Fallback for weird bytes behavior
+                    t = str(t)
+
+            # If it's a string, it might still have the special whitespace characters
+            if isinstance(t, str):
+                # Replace SentencePiece underline (U+2581)
+                t = t.replace('\u2581', ' ')
+                # Replace GPT-2/RoBERTa G-dot (U+0120)
+                t = t.replace('\u0120', ' ')
+                # Replace Newline char (U+010A)
+                t = t.replace('\u010A', '\n')
+                # Replace generic replacement char just in case
+                t = t.replace('', '')
+
+            input_tokens.append(t)
+
+        return topk_data, last_logits, input_tokens
+
+    def get_target_score(self, backprop_config):
+        """
+        Calculates the target scalar score (e.g. logit diff) based on config.
+        Returns the score tensor (attached to graph).
+        """
+        if self.outputs is None:
+             raise ValueError("Model outputs not computed. Call compute_logits first.")
+
+        mode = backprop_config.get("mode", "max_logit")
+        last_logits = self.outputs.logits[0, -1, :]
+        sorted_logits, sorted_indices = torch.sort(last_logits, dim=-1, descending=True)
+
+        target_token_id = backprop_config.get("target_token_id")
+        if target_token_id is not None:
+             target_logit = last_logits[target_token_id]
+        else:
+             target_logit = sorted_logits[0] # Default to Top 1
+
+        if mode == "max_logit":
+            return target_logit
+
+        elif mode == "logit_diff":
+            strategy = backprop_config.get("strategy", "by_topk_avg")
+            top_logit = target_logit
+
+            if strategy == "by_ref_token":
+                ref_id = backprop_config.get("ref_token_id")
+                if ref_id is None:
+                    raise ValueError("ref_token_id required for strategy 'by_ref_token'")
+                contrast_logit = last_logits[ref_id]
+                target_logit = top_logit - contrast_logit
+
+            elif strategy == "demean":
+                target_logit = top_logit - last_logits.mean()
+
+            elif strategy == "by_topk_avg":
+                k = backprop_config.get("k", 10) # default K=10
+                k = min(k, len(sorted_logits))
+                contrast_logit = sorted_logits[:k].mean()
+                target_logit = top_logit - contrast_logit
+
+        return target_logit
+
+    def run_backward_pass(self, backprop_config):
+        """
+        Section 2 Part A: execute backward pass based on configuration.
+        """
+        target_logit = self.get_target_score(backprop_config)
+
+        if target_logit is None:
+            raise ValueError(f"Invalid backprop configuration: {backprop_config}")
+
+        # Clear previous gradients
+        model = self.manager.get_model()
+        model.zero_grad()
+
+        # Also clear gradients on input embeddings if they exist
+        if hasattr(self, 'input_embeddings') and self.input_embeddings is not None:
+            if self.input_embeddings.grad is not None:
+                self.input_embeddings.grad.zero_()
+
+        # Run backward
+        # We need to retain grad on hidden states often?
+        # In notebook: h = outputs.hidden_states[-1]; h.retain_grad(); target_logit.backward()
+        # But we act on layer.output.grad which is captured by hook + activation
+
+        target_logit.backward(retain_graph=True) # retain_graph needed for interactive exploration where we run backward multiple times
+
+    def compute_input_attribution(self, backprop_config):
+        """
+        Compute input attribution (Input * Gradient).
+        """
+        # Ensure correct LRP rule is active
+        # The forward pass must have been run with the correct rule.
+        # If we detect a mismatch, we must force a reload and ask user to re-run forward.
+        # Use the currently loaded LRP rule as default (not hardcoded "Attn-LRP")
+        # to avoid false mismatch when the frontend omits lrp_rule from backprop_config.
+        default_rule = self.manager.current_lrp_rule or "Attn-LRP"
+        required_rule = backprop_config.get("lrp_rule", default_rule)
+
+        if self.manager.current_lrp_rule and self.manager.current_lrp_rule != required_rule:
+             print(f"LRP Rule Mismatch detected (Current: {self.manager.current_lrp_rule}, Requested: {required_rule})")
+             print(f"Reloading model {self.manager.current_model_path} with rule={required_rule}...")
+
+             old_rule = self.manager.current_lrp_rule
+
+             self.manager.load_model(
+                model_path=self.manager.current_model_path,
+                dtype=self.manager.current_dtype,
+                lrp_rule=required_rule
+             )
+
+             # Since the forward pass graph (self.outputs) was built with the OLD rule,
+             # we cannot proceed. The user must re-run compute_logits.
+             raise RuntimeError(
+                 f"LRP rule changed from '{old_rule}' to '{required_rule}'. "
+                 "The model has been reloaded. You MUST re-run 'compute_logits()' to rebuild the computation graph with the new rule, "
+                 "then call 'compute_input_attribution()' again."
+             )
+        # backprop_config['target_token_id'] = 2877
+        self.run_backward_pass(backprop_config)
+
+        # Calculate relevance: (input * grad).sum(-1)
+        # self.input_embeddings is [Batch, Seq, Dim]
+        if self.input_embeddings.grad is None:
+             raise RuntimeError("No gradient found on input embeddings. Ensure compute_logits was run.")
+
+        relevance = (self.input_embeddings * self.input_embeddings.grad).float().sum(-1).detach().cpu()[0]
+        print(f"Computed input attribution with shape: {relevance}")
+        # Return raw relevance
+        return relevance.tolist()
+
+    def compute_input_attribution_gradient(self, backprop_config):
+        """
+        Compute input attribution using vanilla gradient method (Input * Gradient).
+        This does NOT require LRP monkey-patching - uses standard PyTorch autograd.
+        The gradient flows through normal attention/MLP without LRP decomposition rules.
+        """
+        if self.outputs is None:
+            raise RuntimeError("No forward pass found. Run compute_logits first.")
+
+        # Run backward pass (works on vanilla model without LRP)
+        self.run_backward_pass(backprop_config)
+
+        # Calculate relevance: (input * grad).sum(-1)
+        if self.input_embeddings.grad is None:
+            raise RuntimeError("No gradient found on input embeddings. Ensure compute_logits was run.")
+
+        relevance = (self.input_embeddings * self.input_embeddings.grad).float().sum(-1).detach().cpu()[0]
+        print(f"Computed GRADIENT input attribution with shape: {relevance.shape}")
+        return relevance.tolist()
+
+    def compute_perturbation_eval(self, attribution_scores, k_values, target_token_id):
+        """
+        Evaluate attribution quality by perturbing top-attributed tokens.
+
+        For each k in k_values:
+        1. Sort tokens by |attribution score| descending
+        2. Take top-k token indices
+        3. Clone input embeddings, zero out those k tokens' embeddings
+        4. Run forward pass with perturbed embeddings
+        5. Check if the error token (target_token_id) is still top-1
+
+        Args:
+            attribution_scores: list of floats (one per input token)
+            k_values: list of int (e.g., [1, 3, 5, 10])
+            target_token_id: int - the error token ID to check
+
+        Returns:
+            list of result dicts for each k
+        """
+        if self.input_ids is None or self.input_embeddings is None:
+            raise RuntimeError("No forward pass found. Run compute_logits first.")
+
+        model = self.manager.get_model()
+        tokenizer = self.manager.get_tokenizer()
+        device = model.device
+
+        # Get original top-1 prediction for reference
+        with torch.no_grad():
+            original_logits = model(
+                inputs_embeds=self.input_embeddings.detach(),
+                use_cache=False
+            ).logits[0, -1, :]
+            original_top1_id = original_logits.argmax().item()
+            original_target_logit = original_logits[target_token_id].item()
+
+        # Sort tokens by |attribution score| descending
+        scores = torch.tensor(attribution_scores, dtype=torch.float32)
+        sorted_indices = torch.argsort(scores.abs(), descending=True)
+
+        seq_len = self.input_embeddings.shape[1]
+
+        results = []
+        for k in k_values:
+            k_clamped = min(k, seq_len)
+            top_k_indices = sorted_indices[:k_clamped].tolist()
+
+            # Get the token strings being perturbed
+            perturbed_token_strs = []
+            for idx in top_k_indices:
+                if idx < len(self.input_ids[0]):
+                    tid = self.input_ids[0][idx].item()
+                    perturbed_token_strs.append(tokenizer.decode([tid]))
+                else:
+                    perturbed_token_strs.append("?")
+
+            # Clone embeddings and zero out top-k tokens
+            perturbed_embeddings = self.input_embeddings.detach().clone()
+            for idx in top_k_indices:
+                perturbed_embeddings[0, idx, :] = 0.0
+
+            # Forward pass with perturbed embeddings
+            with torch.no_grad():
+                perturbed_logits = model(
+                    inputs_embeds=perturbed_embeddings,
+                    use_cache=False
+                ).logits[0, -1, :]
+
+            perturbed_top1_id = perturbed_logits.argmax().item()
+            perturbed_top1_str = tokenizer.decode([perturbed_top1_id])
+            perturbed_target_logit = perturbed_logits[target_token_id].item()
+
+            # Error is "fixed" if the target token is no longer top-1
+            error_fixed = (perturbed_top1_id != target_token_id)
+
+            # Compute logit change
+            logit_change = perturbed_target_logit - original_target_logit
+
+            # Compute rank of target token after perturbation
+            sorted_perturbed, sorted_perturbed_idx = torch.sort(perturbed_logits, descending=True)
+            target_rank_after = (sorted_perturbed_idx == target_token_id).nonzero(as_tuple=True)[0].item() + 1
+
+            results.append({
+                "k": k,
+                "perturbed_tokens": perturbed_token_strs,
+                "perturbed_indices": top_k_indices,
+                "new_top1_token_id": perturbed_top1_id,
+                "new_top1_token_str": perturbed_top1_str,
+                "error_fixed": error_fixed,
+                "original_target_logit": round(original_target_logit, 4),
+                "perturbed_target_logit": round(perturbed_target_logit, 4),
+                "logit_change": round(logit_change, 4),
+                "target_rank_after": target_rank_after
+            })
+
+            print(f"Perturbation k={k}: error_fixed={error_fixed}, "
+                  f"new_top1='{perturbed_top1_str}' (ID={perturbed_top1_id}), "
+                  f"logit_change={logit_change:.4f}, target_rank={target_rank_after}")
+
+            del perturbed_embeddings
+
+        torch.cuda.empty_cache()
+        return results
+
+    def compute_perturbation_manual(self, perturb_indices, target_token_id):
+        """
+        Evaluate attribution by perturbing manually selected token positions.
+
+        Args:
+            perturb_indices: list of int - token position indices to zero out
+            target_token_id: int - the error token ID to check
+
+        Returns:
+            dict with perturbation result
+        """
+        if self.input_ids is None or self.input_embeddings is None:
+            raise RuntimeError("No forward pass found. Run compute_logits first.")
+
+        model = self.manager.get_model()
+        tokenizer = self.manager.get_tokenizer()
+        seq_len = self.input_embeddings.shape[1]
+
+        # Validate indices
+        valid_indices = [idx for idx in perturb_indices if 0 <= idx < seq_len]
+        if len(valid_indices) == 0:
+            raise ValueError("No valid token indices provided.")
+
+        # Get original top-1 prediction for reference
+        with torch.no_grad():
+            original_logits = model(
+                inputs_embeds=self.input_embeddings.detach(),
+                use_cache=False
+            ).logits[0, -1, :]
+            original_top1_id = original_logits.argmax().item()
+            original_target_logit = original_logits[target_token_id].item()
+
+        # Get the token strings being perturbed
+        perturbed_token_strs = []
+        for idx in valid_indices:
+            if idx < len(self.input_ids[0]):
+                tid = self.input_ids[0][idx].item()
+                perturbed_token_strs.append(tokenizer.decode([tid]))
+            else:
+                perturbed_token_strs.append("?")
+
+        # Clone embeddings and zero out selected tokens
+        perturbed_embeddings = self.input_embeddings.detach().clone()
+        for idx in valid_indices:
+            perturbed_embeddings[0, idx, :] = 0.0
+
+        # Forward pass with perturbed embeddings
+        with torch.no_grad():
+            perturbed_logits = model(
+                inputs_embeds=perturbed_embeddings,
+                use_cache=False
+            ).logits[0, -1, :]
+
+        perturbed_top1_id = perturbed_logits.argmax().item()
+        perturbed_top1_str = tokenizer.decode([perturbed_top1_id])
+        perturbed_target_logit = perturbed_logits[target_token_id].item()
+
+        # Error is "fixed" if the target token is no longer top-1
+        error_fixed = (perturbed_top1_id != target_token_id)
+
+        # Compute logit change
+        logit_change = perturbed_target_logit - original_target_logit
+
+        # Compute rank of target token after perturbation
+        sorted_perturbed, sorted_perturbed_idx = torch.sort(perturbed_logits, descending=True)
+        target_rank_after = (sorted_perturbed_idx == target_token_id).nonzero(as_tuple=True)[0].item() + 1
+
+        result = {
+            "k": len(valid_indices),
+            "perturbed_tokens": perturbed_token_strs,
+            "perturbed_indices": valid_indices,
+            "new_top1_token_id": perturbed_top1_id,
+            "new_top1_token_str": perturbed_top1_str,
+            "error_fixed": error_fixed,
+            "original_target_logit": round(original_target_logit, 4),
+            "perturbed_target_logit": round(perturbed_target_logit, 4),
+            "logit_change": round(logit_change, 4),
+            "target_rank_after": target_rank_after
+        }
+
+        print(f"Manual Perturbation ({len(valid_indices)} tokens): error_fixed={error_fixed}, "
+              f"new_top1='{perturbed_top1_str}' (ID={perturbed_top1_id}), "
+              f"logit_change={logit_change:.4f}, target_rank={target_rank_after}")
+
+        del perturbed_embeddings
+        torch.cuda.empty_cache()
+        return result
+
+    def compute_connection_matrix_gen(self, source, target, node_threshold=None):
+        """
+        Section 2 Part B: Compute Token-to-Token interaction matrix between two nodes.
+        Generator version that yields progress.
+        source, target: int (layer idx) or tuple (layer_idx, 'mid'/'post')
+        """
+        source_layer_idx, source_type = self._parse_node(source)
+        target_layer_idx, target_type = self._parse_node(target)
+
+        model = self.manager.get_model()
+        layers = model.model.layers
+
+        target_layer = layers[target_layer_idx]
+
+        # 1. Identify Source Tensor
+        if source_layer_idx == -1:
+            source_tensor = self.input_embeddings
+        else:
+            layer = layers[source_layer_idx]
+            if source_type == 'mid':
+                # Use Decomposer to get module
+                mid_mod = self.manager.decomposer.get_mid_activation_module(layer)
+                if not mid_mod:
+                    raise ValueError(f"Decomposer could not identify mid-activation module for layer {source_layer_idx}")
+
+                source_tensor = getattr(mid_mod, 'mid_activation', None)
+                if source_tensor is None:
+                    raise ValueError(f"Mid activation for layer {source_layer_idx} not captured. Enable capture_mid in compute_logits.")
+            else:
+                source_tensor = layer.output
+
+        # 2. Identify Target Tensor and Gradient
+        if target_type == 'mid':
+            # We need the gradient at the mid point (input to post_attn_norm)
+            mid_mod = self.manager.decomposer.get_mid_activation_module(target_layer)
+            if not mid_mod:
+                 raise ValueError(f"Decomposer could not identify mid-activation module for target {target_layer_idx}")
+
+            target_tensor = getattr(mid_mod, 'mid_activation', None)
+            if target_tensor is None:
+                raise ValueError(f"Mid activation for target {target_layer_idx} not captured.")
+        else:
+            target_tensor = target_layer.output
+
+        target_grad = target_tensor.grad
+
+        # Disable gradient checkpointing temporarily
+        was_checkpointing = model.is_gradient_checkpointing
+        if was_checkpointing:
+            model.gradient_checkpointing_disable()
+
+        try:
+            # Prepare Input
+            target_layer_input = source_tensor.detach()
+            batch_size, seq_len, hidden_dim = target_layer_input.shape
+
+            # Target Real Relevance
+            if target_grad is not None:
+                real_target_rel = (target_tensor * target_grad).sum(dim=-1)[0]
+            else:
+                real_target_rel = torch.zeros(seq_len, device=model.device)
+
+            # Filter Indices
+            total_params = model.num_parameters()
+            if node_threshold is None: node_threshold = 0.01
+
+            if node_threshold > 0:
+                indices_to_compute = torch.nonzero(real_target_rel.abs() > node_threshold).squeeze(-1).tolist()
+                if isinstance(indices_to_compute, int): indices_to_compute = [indices_to_compute]
+                print(f"DEBUG: Node Threshold {node_threshold}. Computing for {len(indices_to_compute)}/{seq_len} nodes.")
+            else:
+                indices_to_compute = list(range(seq_len))
+
+            # Fixed Position IDs (for Rotary)
+            position_ids = torch.arange(0, seq_len, dtype=torch.long, device=model.device).unsqueeze(0)
+
+            # Construct Operation Sequence
+            ops = []
+            if source_layer_idx != -1:
+                 if source_type == 'mid':
+                      ops.append(('part2', layers[source_layer_idx]))
+
+            # Intermediate Layers
+            for i in range(source_layer_idx + 1, target_layer_idx):
+                 ops.append(('part1', layers[i]))
+                 ops.append(('part2', layers[i]))
+
+            # Target Layer
+            if target_layer_idx > source_layer_idx:
+                 ops.append(('part1', layers[target_layer_idx]))
+                 if target_type == 'post':
+                      ops.append(('part2', layers[target_layer_idx]))
+            elif target_layer_idx == source_layer_idx:
+                 pass # Already handled or identity
+            elif source_layer_idx == -1:
+                 # Special case: source is embeddings, target is 0
+                 # Range was (0,0) empty.
+                 # Need to add target 0 parts
+                 ops.append(('part1', layers[target_layer_idx]))
+                 if target_type == 'post':
+                      ops.append(('part2', layers[target_layer_idx]))
+
+            # Pre-calc Rotary Embedding (using dummy execution or helper)
+            # We assume rotary depends only on position_ids and shape
+            rotary_emb = None
+            if hasattr(model.model, 'rotary_emb'):
+                 rotary_emb = model.model.rotary_emb(target_layer_input, position_ids)
+            elif hasattr(model.model, 'rotary_embs') and 'full_attention' in model.model.rotary_embs:
+                 rotary_emb = model.model.rotary_embs['full_attention'](target_layer_input, position_ids)
+            elif hasattr(model.model, 'rotary_embs') and len(model.model.rotary_embs) > 0:
+                 rotary_emb = list(model.model.rotary_embs.values())[0](target_layer_input, position_ids)
+
+            # Chunk Processing
+            current_dtype = target_layer_input.dtype
+            BATCH_CHUNK_SIZE = get_batch_chunk_size(total_params, current_dtype)
+            token_interaction = torch.zeros(seq_len, seq_len, device=model.device)
+            target_grad_full = target_grad # Alias
+
+            total_items = len(indices_to_compute)
+            print(f"DEBUG: BATCH_CHUNK_SIZE={BATCH_CHUNK_SIZE}, total_items={total_items}, seq_len={seq_len}, params={total_params/1e9:.2f}B, dtype={current_dtype}")
+
+            for i in range(0, total_items, BATCH_CHUNK_SIZE):
+                yield {"type": "progress", "current": i, "total": total_items}
+
+                chunk_indices = indices_to_compute[i : i + BATCH_CHUNK_SIZE]
+                current_batch_size = len(chunk_indices)
+
+                expanded_input = target_layer_input.expand(current_batch_size, seq_len, hidden_dim).clone().requires_grad_(True)
+
+                # Execute Ops
+                hidden_states = expanded_input
+                for op_type, layer_mod in ops:
+                    if op_type == 'part1':
+                        hidden_states = self._forward_part1(layer_mod, hidden_states, position_embeddings=rotary_emb)
+                    else:
+                        hidden_states = self._forward_part2(layer_mod, hidden_states)
+
+                reconstructed_output = hidden_states
+
+                # Backward
+                grad_output_chunk = torch.zeros(current_batch_size, seq_len, hidden_dim, dtype=reconstructed_output.dtype, device=model.device)
+
+                for batch_idx, global_idx in enumerate(chunk_indices):
+                    if target_grad_full is not None:
+                        grad_output_chunk[batch_idx, global_idx, :] = target_grad_full[0, global_idx, :]
+
+                grad_input = torch.autograd.grad(outputs=reconstructed_output, inputs=expanded_input, grad_outputs=grad_output_chunk, retain_graph=False)[0]
+
+                chunk_relevance = (grad_input * expanded_input).sum(dim=-1)
+                token_interaction[chunk_indices, :] = chunk_relevance.detach().to(token_interaction.dtype)
+
+                del expanded_input, hidden_states, reconstructed_output, grad_output_chunk, grad_input, chunk_relevance
+                torch.cuda.empty_cache()
+
+            # Source Real Relevance
+            if source_layer_idx == -1:
+                if self.input_embeddings.grad is not None:
+                    real_source_rel = (self.input_embeddings * self.input_embeddings.grad).sum(dim=-1)[0]
+                else:
+                    real_source_rel = torch.zeros(seq_len, device=model.device)
+            else:
+                if source_tensor.grad is not None:
+                    real_source_rel = (source_tensor * source_tensor.grad).sum(dim=-1)[0]
+                else:
+                    real_source_rel = torch.zeros(seq_len, device=model.device)
+
+            yield {
+                "type": "result",
+                "payload": {
+                    "matrix": token_interaction.detach().float().cpu().numpy(),
+                    "real_target_rel": real_target_rel.detach().float().cpu().numpy(),
+                    "real_source_rel": real_source_rel.detach().float().cpu().numpy()
+                }
+            }
+
+        finally:
+            if was_checkpointing:
+                model.gradient_checkpointing_enable()
+
+    def compute_connection_matrix(self, source, target):
+        for item in self.compute_connection_matrix_gen(source, target):
+            if item.get("type") == "result":
+                return item["payload"]
+        return None

backend/error_token_location.py

@@ -0,0 +1,312 @@
+import logging
+logging.basicConfig(level=logging.INFO,
+                    format='%(asctime)s | %(levelname)-6s | %(name)-40s || %(message)s',
+                    datefmt='%m-%d %H:%M:%S')
+logger = logging.getLogger(__name__)
+
+import os
+import sys
+from pathlib import Path
+
+import yaml
+import json
+from typing import Optional
+from jinja2 import Template
+
+# Use env-var based OpenAI client instead of Azure-specific endpoint
+_openai_client = None
+
+def _get_openai_client():
+    """Get or create OpenAI client using environment variables."""
+    global _openai_client
+    if _openai_client is not None:
+        return _openai_client
+
+    try:
+        import openai
+    except ImportError:
+        logger.warning("openai package not installed. LLM-based error token localization will not be available.")
+        return None
+
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        logger.warning("OPENAI_API_KEY not set. LLM-based error token localization will not be available.")
+        return None
+
+    base_url = os.environ.get("OPENAI_BASE_URL", None)
+    kwargs = {"api_key": api_key}
+    if base_url:
+        kwargs["base_url"] = base_url
+
+    _openai_client = openai.OpenAI(**kwargs)
+    return _openai_client
+
+
+def _get_default_models():
+    """Get the list of LLM models to use as validators from env or defaults."""
+    models_str = os.environ.get("LLM_MODELS", "gpt-4o-mini")
+    return [m.strip() for m in models_str.split(",") if m.strip()]
+
+
+class ErrorTokenLocator:
+    def __init__(self, model, tokenizer, prompt_template_yaml=None):
+        """
+        Initialize the error token locator
+
+        Args:
+            model: The language model to use
+            tokenizer: The corresponding tokenizer for tokenizing text
+            prompt_template_yaml (str, optional): Path to the prompt template YAML file, uses default template when None
+        """
+        self.model = model
+        self.tokenizer = tokenizer
+
+        self.client = _get_openai_client()
+
+        self.endpoint_list = _get_default_models()
+
+        if prompt_template_yaml is None:
+            self.system_prompt = self.load_general_prompt_template()
+        else:
+            self.system_prompt = self.load_general_prompt_template(prompt_template_yaml)
+
+    def prompt_constructor(self, query, completion, ground_truth: str=None):
+        """
+        Construct prompts for error token location
+
+        Args:
+            query (str): The input query/prompt
+            completion (str): The completion text generated by the model
+            ground_truth (str, optional): The correct answer/ground truth, defaults to None
+
+        Returns:
+            tuple: (msg, tokens)
+                - msg (list): The constructed conversation message list containing system and user messages
+                - tokens (list): List of token IDs for the completion
+        """
+        tokens = self.tokenizer(completion).input_ids
+        indexed_completion = ""
+        for i, tok in enumerate(tokens):
+            indexed_completion += f"{self.tokenizer.decode([tok])}[{i}] "
+
+        user_msg_content = {
+            "prompt": query,
+            "completion": completion,
+            "indexed_completion": indexed_completion,
+            "ground_truth": ground_truth
+        }
+
+        msg = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": json.dumps(user_msg_content, indent=2)}
+        ]
+        return msg, tokens
+
+    def load_general_prompt_template(self, prompt_template_yaml=None):
+        """
+        Load the general prompt template
+
+        Args:
+            prompt_template_yaml (str, optional): Path to the YAML template file, uses default path when None
+
+        Returns:
+            str: The rendered system prompt template string
+        """
+        if prompt_template_yaml is None:
+            # Use path relative to this file's location
+            prompt_template_yaml = os.path.join(
+                os.path.dirname(os.path.abspath(__file__)),
+                "token_locator_prompts",
+                "err_token_localization.yaml"
+            )
+
+        with open(prompt_template_yaml, "r", encoding='utf-8') as f:
+            system_prompt_temp = yaml.safe_load(f)
+        system_prompt = Template(system_prompt_temp['system']).render(
+            dataset_description="No dataset description provided.",
+            dataset_specific_instructions="-No dataset specific instructions provided."
+        )
+
+        return system_prompt
+
+    def call_validator(self, msg, tokens, endpoint_list=None):
+        """
+        Call the validator for error token location validation
+
+        Args:
+            msg (list): The constructed conversation message list
+            tokens (list): List of token IDs
+            endpoint_list (list): List of validator endpoint names
+
+        Returns:
+            tuple: (completion_before_err, explanation, vote_details)
+                - completion_before_err (str): Completion text truncated before the error token
+                - explanation (str): Explanation of why this token is the error
+                - vote_details (dict): Detailed voting information for each validator
+        """
+        if self.client is None:
+            raise RuntimeError(
+                "OpenAI client not available. Please set OPENAI_API_KEY environment variable "
+                "or use manual truncation instead."
+            )
+
+        responses = []
+        if endpoint_list is None:
+            endpoint_list = self.endpoint_list
+        for model_name in endpoint_list:
+            response = self.client.chat.completions.create(
+                model=model_name,
+                messages=msg,
+                temperature=0,
+                seed=42,
+                top_p=1,
+                frequency_penalty=0,
+                presence_penalty=0,
+                response_format={"type": "json_object"}
+            )
+            responses.append(response)
+
+        votes = {}
+        first_response_for_token = {}
+        validator_votes = {}  # Track each validator's vote
+
+        for idx, response in enumerate(responses):
+            model_name = endpoint_list[idx]
+            try:
+                res_json = json.loads(response.choices[0].message.content)
+                token_index = int(res_json["token_index"])
+                explanation = res_json.get("explanation", "")
+
+                # Track votes
+                votes[token_index] = votes.get(token_index, 0) + 1
+                if token_index not in first_response_for_token:
+                    first_response_for_token[token_index] = (response, idx)
+
+                # Store each validator's vote details
+                validator_votes[model_name] = {
+                    "token_index": token_index,
+                    "error_token": self.tokenizer.decode([tokens[token_index]]) if token_index < len(tokens) else "N/A",
+                    "explanation": explanation
+                }
+            except Exception as e:
+                logger.error(f"Error processing response from {model_name}: {e}")
+                validator_votes[model_name] = {
+                    "token_index": -1,
+                    "error_token": "Error",
+                    "explanation": f"Failed to parse response: {str(e)}"
+                }
+
+        if not votes:
+            return responses[0] if responses else None, "", {}
+
+        max_votes = max(votes.values())
+        candidates = [t for t, c in votes.items() if c == max_votes]
+        # pick the candidate whose first corresponding response appeared earliest
+        winner_token = min(candidates, key=lambda t: first_response_for_token[t][1])
+
+        response = first_response_for_token[winner_token][0].choices[0].message.content
+        failure_rca = json.loads(response)
+        completion_before_err = ""
+        for idx, tok in enumerate(tokens):
+            if idx == failure_rca["token_index"]:
+                break
+            completion_before_err += f"{self.tokenizer.decode([tok])}"
+        explanation = failure_rca["explanation"]
+
+        # Add vote summary to vote_details
+        vote_summary = {
+            "winner_token_index": winner_token,
+            "winner_votes": max_votes,
+            "total_validators": len(endpoint_list),
+            "vote_distribution": votes
+        }
+
+        return completion_before_err, explanation, {
+            "validators": validator_votes,
+            "summary": vote_summary
+        }
+
+    def locate_error_token(self, prompt: str, completion: str, ground_truth: str = None,
+                           validators: Optional[list] = None,
+                           use_llm: bool = True,
+                           manual_chunks: Optional[list] = None):
+        """
+        Main method to locate the error token in a completion
+
+        Args:
+            prompt (str): The input prompt
+            completion (str): The completion text to analyze
+            ground_truth (str, optional): The correct answer/ground truth, defaults to None
+
+        Returns:
+            dict: Dictionary containing:
+                - status (str): "success" or "error"
+                - truncated_text (str): Prompt + completion truncated before error token
+                - explanation (str): Explanation of the error
+                - error_token_index (int): Index of the error token
+        """
+        try:
+            # Construct prompt messages
+            msg, tokens = self.prompt_constructor(prompt, completion, ground_truth)
+
+            # If user requests to skip LLM search, use manual chunks if provided
+            if not use_llm:
+                if manual_chunks and len(manual_chunks) > 0:
+                    manual_chunk = manual_chunks[0]
+                    completion_before_err = manual_chunk
+                    explanation = "Manual chunk provided by user (LLM search skipped)."
+                    error_token_index = len(self.tokenizer(manual_chunk).input_ids)
+                    truncated_text = prompt + completion_before_err
+                    return {
+                        "status": "success",
+                        "truncated_text": truncated_text,
+                        "explanation": explanation,
+                        "error_token_index": error_token_index
+                    }
+                else:
+                    return {
+                        "status": "error",
+                        "message": "LLM search disabled but no manual chunk provided.",
+                        "truncated_text": "",
+                        "explanation": ""
+                    }
+
+            # Check if OpenAI client is available
+            if self.client is None:
+                return {
+                    "status": "error",
+                    "message": "OpenAI API key not configured. Please set OPENAI_API_KEY environment variable or use manual truncation (disable LLM search).",
+                    "truncated_text": "",
+                    "explanation": ""
+                }
+
+            # If validators provided, use them for this call
+            endpoint_list = validators if (validators and isinstance(validators, list) and len(validators) > 0) else None
+
+            # Call validator to get error token location with vote details
+            completion_before_err, explanation, vote_details = self.call_validator(msg, tokens, endpoint_list)
+
+            # Combine prompt with truncated completion
+            truncated_text = prompt + completion_before_err
+
+            # Calculate error token index
+            error_token_index = len(self.tokenizer(completion_before_err).input_ids)
+
+            return {
+                "status": "success",
+                "truncated_text": truncated_text,
+                "explanation": explanation,
+                "error_token_index": error_token_index,
+                "vote_details": vote_details
+            }
+
+        except Exception as e:
+            logger.error(f"Error in locate_error_token: {e}")
+            import traceback
+            traceback.print_exc()
+            return {
+                "status": "error",
+                "message": str(e),
+                "truncated_text": "",
+                "explanation": ""
+            }

backend/graph_metrics.py

@@ -0,0 +1,423 @@
+import networkx as nx
+import numpy as np
+import pandas as pd
+from scipy.stats import entropy
+
+class GraphMetrics:
+    """
+    A unified class for calculating metrics on attribution graphs.
+    Metrics are organized into three groups:
+    1. Connectivity & Sparsity
+    2. Information Flow & Concentration
+    3. Temporal Dynamics
+
+    Each metric provides:
+    - Global Summary (single scalar)
+    - Layer-wise Profile (DataFrame/Series)
+    """
+    def __init__(self, G, tokens=None, total_seq_len=None):
+        self.G = G
+        self.tokens = tokens
+
+        if total_seq_len is None and tokens is not None:
+            self.total_seq_len = len(tokens)
+        else:
+            self.total_seq_len = total_seq_len
+
+        self.nodes_by_layer = self._group_nodes_by_layer()
+        self.sorted_layers = sorted(self.nodes_by_layer.keys())
+
+    def _get_token_str(self, index):
+        if self.tokens and 0 <= index < len(self.tokens):
+            return self.tokens[index]
+        return str(index)
+
+    def _group_nodes_by_layer(self):
+        groups = {}
+        for n in self.G.nodes():
+            layer = n[0]
+            if layer not in groups:
+                groups[layer] = []
+            groups[layer].append(n)
+        return groups
+
+    def _calculate_gini(self, array):
+        """Auxiliary function to calculate Gini coefficient."""
+        array = np.abs(array)
+        if np.sum(array) == 0:
+            return 0.0
+        array = np.sort(array)
+        index = np.arange(1, array.shape[0] + 1)
+        n = array.shape[0]
+        return ((np.sum((2 * index - n  - 1) * array)) / (n * np.sum(array)))
+
+    def _calculate_top_mass(self, values, fraction=0.9):
+        """Calculates count and percentage of items needed to reach mass fraction."""
+        values = np.abs(values)
+        total = values.sum()
+        if total == 0:
+            return 0, 0.0
+
+        sorted_vals = np.sort(values)[::-1]
+        cumsum = np.cumsum(sorted_vals)
+        cutoff = total * fraction
+
+        # Find index where cumsum >= cutoff
+        idx = np.searchsorted(cumsum, cutoff)
+        count = idx + 1
+        pct = (count / len(values)) * 100.0
+        return count, pct
+
+    def _get_top_nodes_by_mass(self, node_val_pairs, fraction=0.9):
+        """
+        Calculates top nodes making up mass fraction.
+        Args:
+            node_val_pairs: list of (node, value)
+        Returns:
+            count, pct, list of (token_idx, token_str)
+        """
+        if not node_val_pairs:
+            return 0, 0.0, []
+
+        values = np.array([abs(v) for _, v in node_val_pairs])
+        total = values.sum()
+        if total == 0:
+            return 0, 0.0, []
+
+        # Sort indices by value descending
+        sorted_indices = np.argsort(values)[::-1]
+        sorted_vals = values[sorted_indices]
+
+        cumsum = np.cumsum(sorted_vals)
+        cutoff = total * fraction
+
+        idx_cutoff = np.searchsorted(cumsum, cutoff)
+        count = idx_cutoff + 1
+        pct = (count / len(values)) * 100.0
+
+        # Extract top nodes
+        top_indices = sorted_indices[:count]
+        top_nodes = []
+        for i in top_indices:
+            node, _ = node_val_pairs[i]
+            # node is (layer, token_idx)
+            token_idx = int(node[1])
+            token_str = self._get_token_str(token_idx)
+            top_nodes.append((token_idx, token_str))
+
+        return count, pct, top_nodes
+
+    # ==========================================
+    # Group 1: Connectivity & Sparsity
+    # ==========================================
+
+    def get_connectivity_stats(self):
+        """Calculates Scale, Complexity, and Branching metrics."""
+        if self.G.number_of_nodes() == 0:
+            return pd.DataFrame()
+
+        layer_stats = []
+
+        # --- Layer-wise ---
+        for layer in self.sorted_layers:
+            nodes = self.nodes_by_layer[layer]
+
+            # Fan-In (Inputs from prev layers)
+            in_degrees = [self.G.in_degree(n) for n in nodes]
+            avg_in = np.mean(in_degrees) if in_degrees else 0
+            # std_in = np.std(in_degrees) if in_degrees else 0
+
+            # Fan-Out (Outputs to next layers)
+            out_degrees = [self.G.out_degree(n) for n in nodes]
+            avg_out = np.mean(out_degrees) if out_degrees else 0
+            # std_out = np.std(out_degrees) if out_degrees else 0
+
+            # Derived Layer Metrics
+            node_count = len(nodes)
+            edge_count = np.sum(out_degrees)
+
+            # Edge/Node Ratio
+            edge_node_ratio = edge_count / node_count if node_count > 0 else 0
+
+            stats = {
+                "Layer": layer,
+                "Node_Count": node_count,
+                "Edge_Count": edge_count,
+                "Edge_Node_Ratio": edge_node_ratio,
+                "Avg_Fan_In": avg_in,
+                # "Std_Fan_In": std_in,
+                "Avg_Fan_Out": avg_out,
+                # "Std_Fan_Out": std_out
+            }
+
+            if self.total_seq_len:
+                stats["Sparsity"] = 1.0 - (node_count / self.total_seq_len)
+
+            layer_stats.append(stats)
+
+        df_layer = pd.DataFrame(layer_stats)
+
+        return df_layer
+
+    # ==========================================
+    # Group 2: Information Flow & Concentration
+    # ==========================================
+
+    def get_flow_stats(self):
+        """Calculates Effective Branching, Gini, and Verticality."""
+        if self.G.number_of_nodes() == 0:
+            return pd.DataFrame()
+
+        layer_stats = []
+
+        for layer in self.sorted_layers:
+            nodes = self.nodes_by_layer[layer]
+            # eff_degrees_in = []
+            # eff_degrees_out = []
+            node_ginis_in = []
+            node_ginis_out = []
+
+            # New Collectors
+            layer_node_rels = []
+            layer_out_edge_rels = []
+
+            layer_vertical_mass = 0.0
+            layer_total_mass = 0.0
+
+            for n in nodes:
+                # Node Relevance
+                rel = abs(self.G.nodes[n].get('relevance', 0.0))
+                layer_node_rels.append(rel)
+
+                # Incoming Edges analysis
+                in_edges = self.G.in_edges(n, data=True)
+                weights_in = np.array([abs(d.get('weight', 0.0)) for u, v, d in in_edges])
+
+                # Flow Mass
+                if len(weights_in) > 0:
+                    w_sum = weights_in.sum()
+                    layer_total_mass += w_sum
+
+                    # Verticality check
+                    for u, v, d in in_edges:
+                        if u[1] == v[1]: # Same token index
+                            layer_vertical_mass += abs(d.get('weight', 0.0))
+
+
+
+                # Outgoing Edges analysis
+                out_edges = self.G.out_edges(n, data=True)
+                weights_out = np.array([abs(d.get('weight', 0.0)) for u, v, d in out_edges])
+
+                # Collect Out Edge Relevances
+                if len(weights_out) > 0:
+                    layer_out_edge_rels.extend(weights_out)
+
+
+                # Gini (On Inputs)
+                if len(weights_in) > 1:
+                    g = self._calculate_gini(weights_in)
+                    node_ginis_in.append(g)
+                elif len(weights_in) == 1:
+                    node_ginis_in.append(1.0)
+
+                # Gini (On Outputs)
+                if len(weights_out) > 1:
+                    g = self._calculate_gini(weights_out)
+                    node_ginis_out.append(g)
+                elif len(weights_out) == 1:
+                    node_ginis_out.append(1.0)
+
+            # Node Statistics
+            layer_node_rels = np.array(layer_node_rels)
+            avg_node_rel = np.mean(layer_node_rels) if len(layer_node_rels) > 0 else 0
+            node_rel_gini = self._calculate_gini(layer_node_rels) if len(layer_node_rels) > 0 else 0
+
+            # Node Mass Fractions (Moved to Group 3)
+
+            # Out Edge Statistics
+            layer_out_edge_rels = np.array(layer_out_edge_rels)
+            avg_out_edge_rel = np.mean(layer_out_edge_rels) if len(layer_out_edge_rels) > 0 else 0
+            out_edge_rel_gini = self._calculate_gini(layer_out_edge_rels) if len(layer_out_edge_rels) > 0 else 0
+
+            layer_stats.append({
+                "Layer": layer,
+
+                # Node Relevance Stats
+                "Avg_Node_Rel": avg_node_rel,
+                "Node_Rel_Gini": node_rel_gini,
+
+                # Out Edge Relevance Stats
+                "Avg_Out_Edge_Rel": avg_out_edge_rel,
+                "Gini_Out_Edge_Rel": out_edge_rel_gini,
+
+                # by node Gini Stats
+                "Avg_Edge_Gini_In_by_Node": np.mean(node_ginis_in) if node_ginis_in else 0,
+                "Avg_Edge_Gini_Out_by_Node": np.mean(node_ginis_out) if node_ginis_out else 0,
+
+                "Verticality_Ratio_by_Node": layer_vertical_mass / layer_total_mass if layer_total_mass > 0 else 0
+            })
+
+        df_layer = pd.DataFrame(layer_stats)
+
+        return df_layer
+
+    # ==========================================
+    # Group 3: Node Hubs & Significant Tokens
+    # ==========================================
+
+    def get_node_hub_stats(self):
+        """Calculates Top Mass Nodes, Degree Hubs, and Signed Relevance Hubs."""
+        if self.G.number_of_nodes() == 0:
+            return pd.DataFrame()
+
+        layer_stats = []
+
+        for layer in self.sorted_layers:
+            nodes = self.nodes_by_layer[layer]
+
+            # 1. Gather Data
+            node_rels = [] # Pairs of (node, rel)
+            in_degrees = [] # Pairs of (node, deg)
+            out_degrees = [] # Pairs of (node, deg)
+
+            pos_node_rels = []
+            neg_node_rels = []
+
+            for n in nodes:
+                # Relevance
+                rel = self.G.nodes[n].get('relevance', 0.0)
+                node_rels.append((n, rel))
+
+                if rel >= 0:
+                    pos_node_rels.append((n, rel))
+                else:
+                    neg_node_rels.append((n, rel))
+
+                # Degree
+                in_deg = self.G.in_degree(n)
+                in_degrees.append((n, in_deg))
+
+                out_deg = self.G.out_degree(n)
+                out_degrees.append((n, out_deg))
+
+            # 2. General Top Mass (Abs)
+            n_90, pct_90, top_90_nodes = self._get_top_nodes_by_mass(node_rels, fraction=0.9)
+
+            # 3. Top Mass Positive & Negative
+            n_pos, pct_pos, top_pos_nodes = self._get_top_nodes_by_mass(pos_node_rels, fraction=0.9)
+            n_neg, pct_neg, top_neg_nodes = self._get_top_nodes_by_mass(neg_node_rels, fraction=0.9)
+
+            # 4. Degree Hubs (Mean + Std)
+            # In-Degree
+            in_deg_vals = [d for _, d in in_degrees]
+            if len(in_deg_vals) > 0:
+                avg_in = np.mean(in_deg_vals)
+                std_in = np.std(in_deg_vals)
+                thresh_in = avg_in + std_in
+                hub_in_nodes = []
+                for n, deg in in_degrees:
+                    if deg > thresh_in:
+                        token_idx = int(n[1])
+                        token_str = self._get_token_str(token_idx)
+                        hub_in_nodes.append((token_idx, token_str))
+            else:
+                hub_in_nodes = []
+
+            # Out-Degree
+            out_deg_vals = [d for _, d in out_degrees]
+            if len(out_deg_vals) > 0:
+                avg_out = np.mean(out_deg_vals)
+                std_out = np.std(out_deg_vals)
+                thresh_out = avg_out + std_out
+                hub_out_nodes = []
+                for n, deg in out_degrees:
+                    if deg > thresh_out:
+                         token_idx = int(n[1])
+                         token_str = self._get_token_str(token_idx)
+                         hub_out_nodes.append((token_idx, token_str))
+            else:
+                hub_out_nodes = []
+
+            layer_stats.append({
+                "Layer": layer,
+
+                # Abs Relevance Hubs
+                "Node_Rel_Top_90_Pct_Count": n_90,
+                "Node_Rel_Top_90_Pct": pct_90,
+                # "Top_Rel_Nodes": top_90_nodes,
+
+                # Signed Relevance Hubs
+                "Pos_Rel_Top_90_Pct_Count": n_pos,
+                "Pos_Rel_Top_90_Pct": pct_pos,
+                # "Top_Pos_Nodes": top_pos_nodes,
+
+                "Neg_Rel_Top_90_Pct_Count": n_neg,
+                "Neg_Rel_Top_90_Pct": pct_neg,
+                # "Top_Neg_Nodes": top_neg_nodes,
+
+                # Structural Hubs
+                "Hub_In_Count": len(hub_in_nodes),
+                # "Hub_In_Nodes": hub_in_nodes,
+                "Hub_Out_Count": len(hub_out_nodes),
+                # "Hub_Out_Nodes": hub_out_nodes
+            })
+
+        return pd.DataFrame(layer_stats)
+
+    # ==========================================
+    # Group 4: Temporal Dynamics
+    # ==========================================
+
+    def get_temporal_stats(self):
+        """Calculates Lookback, Drift, and Locality."""
+        if self.G.number_of_nodes() == 0:
+            return pd.DataFrame()
+
+        layer_stats = []
+
+        for layer in self.sorted_layers:
+            nodes = self.nodes_by_layer[layer]
+
+            # --- Center of Mass (Drift) ---
+            positions = []
+            relevances = []
+
+            # For CoM, we look at the nodes THEMSELVES, not edges
+            for n in nodes:
+                # We need node relevance. If not stored, default to 1
+                rel = self.G.nodes[n].get('relevance', 1.0)
+                positions.append(n[1]) # Token Index
+                relevances.append(abs(rel))
+
+            positions = np.array(positions)
+            relevances = np.array(relevances)
+
+            if relevances.sum() > 0:
+                com = np.average(positions, weights=relevances)
+            else:
+                com = np.mean(positions) if len(positions) > 0 else 0
+
+            # --- Lookback & Locality (Edges entering this layer) ---
+            layer_lookbacks = []
+
+            for n in nodes:
+                in_edges = self.G.in_edges(n)
+                for u, v in in_edges:
+                    # u is source, v is target (n)
+                    # Lookback = v_pos - u_pos
+                    dist = v[1] - u[1]
+                    layer_lookbacks.append(dist)
+
+            layer_lookbacks = np.array(layer_lookbacks)
+
+            layer_stats.append({
+                "Layer": layer,
+                "Center_Of_Mass_Idx": com,
+                "Mean_Lookback": np.mean(layer_lookbacks) if len(layer_lookbacks) > 0 else 0,
+                "Local_Processing_Ratio": np.mean(layer_lookbacks == 0) if len(layer_lookbacks) > 0 else 0
+            })
+
+        df_layer = pd.DataFrame(layer_stats)
+
+        return df_layer

backend/metrics.py

@@ -0,0 +1,143 @@
+import numpy as np
+from scipy.stats import entropy
+
+def calculate_normalized_entropy(relevance):
+    """
+    Calculates Shannon Entropy and Normalized Entropy of the relevance distribution.
+    Use absolute values of relevance to treat as a probability distribution.
+
+    Returns:
+        tuple: (entropy, normalized_entropy)
+    """
+    rel = np.array(relevance)
+    abs_rel = np.abs(rel)
+    n = len(abs_rel)
+
+    if n == 0:
+        return 0.0, 0.0
+
+    total_mass = abs_rel.sum()
+    if total_mass > 0:
+        probs = abs_rel / total_mass
+    else:
+        # Uniform if sum is 0
+        probs = np.ones_like(abs_rel) / n
+
+    # Shannon Entropy
+    ent = entropy(probs)
+
+    # Normalized Entropy (0 to 1)
+    # max_ent = ln(N)
+    max_ent = np.log(n)
+    norm_ent = ent / max_ent if max_ent > 0 else 0
+
+    return ent, norm_ent
+
+def calculate_gini_coefficient(relevance):
+    """
+    Calculates the Gini Coefficient of the relevance absolute values.
+    Values near 1 indicate very unequal distribution (high concentration/sparsity).
+    Values near 0 indicate uniform distribution.
+    """
+    rel = np.array(relevance)
+    abs_rel = np.abs(rel)
+    n = len(abs_rel)
+
+    if n == 0 or abs_rel.mean() == 0:
+        return 0.0
+
+    # Sort ascending for Gini calculation formula
+    sorted_asc = np.sort(abs_rel)
+    index = np.arange(1, n + 1)
+
+    # Gini = (2 * sum(i * xi) - (n + 1) * sum(xi)) / (n * sum(xi))
+    gini_coeff = (2 * np.sum(index * sorted_asc) - (n + 1) * np.sum(sorted_asc)) / (n * np.sum(sorted_asc))
+
+    return gini_coeff
+
+def calculate_top_mass_fraction(relevance, fraction=0.9):
+    """
+    Calculates how many tokens (and percentage) account for a specific fraction of the total attribution mass.
+
+    Returns:
+        tuple: (count, percentage, cdf_array, sorted_indices)
+    """
+    rel = np.array(relevance)
+    abs_rel = np.abs(rel)
+    n = len(abs_rel)
+
+    if n == 0:
+        return 0, 0.0, np.array([]), np.array([])
+
+    total_mass = abs_rel.sum()
+    if total_mass == 0:
+        # Avoid division by zero
+        # Return as if uniform or none
+        return n, 100.0, np.linspace(0, 1, n), np.arange(n)
+
+    # Sort descending
+    sorted_indices = np.argsort(abs_rel)[::-1]
+    sorted_mass = abs_rel[sorted_indices]
+    cumulative_mass = np.cumsum(sorted_mass)
+    cdf = cumulative_mass / total_mass
+
+    # Find how many tokens account for 'fraction' mass
+    count = np.searchsorted(cdf, fraction) + 1
+    percentage = (count / n) * 100
+
+    return count, percentage, cdf, sorted_indices, sorted_mass
+
+def calculate_center_of_mass(relevance):
+    """
+    Calculates the center of mass (expected position) of the attribution.
+
+    Returns:
+        tuple: (absolute_center_index, relative_center_0_to_1)
+    """
+    rel = np.array(relevance)
+    abs_rel = np.abs(rel)
+    total_mass = abs_rel.sum()
+    n = len(abs_rel)
+
+    if n == 0 or total_mass == 0:
+        return 0.0, 0.0
+
+    indices = np.arange(n)
+    # E[index] = sum(p_i * i)
+    center_of_mass = np.sum(indices * abs_rel) / total_mass
+
+    # Relative center: 0 = start, 1 = end
+    # Use (n-1) as denominator because indices are 0 to n-1
+    relative_center = center_of_mass / (n - 1) if n > 1 else 0.5
+
+    return center_of_mass, relative_center
+
+def calculate_early_late_ratio(relevance, split_ratio=0.5):
+    """
+    Calculates the ratio of mass in the first part vs the second part of the sequence.
+
+    Args:
+        relevance: The attribution array.
+        split_ratio: The point to split early vs late (0.5 = middle).
+
+    Returns:
+        float: Ratio (Early Mass / Late Mass).
+               Returns infinity if Late Mass is 0.
+               Returns 0 if Early Mass is 0.
+    """
+    rel = np.array(relevance)
+    abs_rel = np.abs(rel)
+    n = len(abs_rel)
+
+    if n == 0:
+        return 0.0
+
+    split_index = int(n * split_ratio)
+
+    early_mass = np.sum(abs_rel[:split_index])
+    late_mass = np.sum(abs_rel[split_index:])
+
+    if late_mass == 0:
+        return float('inf') if early_mass > 0 else 0.0
+
+    return early_mass / late_mass

backend/models/__init__.py

@@ -0,0 +1,4 @@
+from .manager import ModelManager
+from .base import LayerDecomposer
+from .qwen import QwenDecomposer
+from .factory import get_decomposer

backend/models/base.py

@@ -0,0 +1,32 @@
+from abc import ABC, abstractmethod
+import torch.nn as nn
+
+class LayerDecomposer(ABC):
+    """
+    Abstract base class for decomposing Transformer layers into
+    Attention (Part 1) and MLP (Part 2) components.
+    """
+
+    @abstractmethod
+    def get_mid_activation_module(self, layer_module):
+        """
+        Returns the module whose input corresponds to 'resid_mid'.
+        This is typically the Post-Attention LayerNorm.
+        """
+        pass
+
+    @abstractmethod
+    def forward_part1(self, layer_module, hidden_states, position_embeddings=None, attention_mask=None):
+        """
+        Executes: Norm -> Attn -> Residual Add
+        Returns: resid_mid
+        """
+        pass
+
+    @abstractmethod
+    def forward_part2(self, layer_module, hidden_states):
+        """
+        Executes: Norm -> MLP -> Residual Add
+        Returns: resid_post
+        """
+        pass

backend/models/factory.py

@@ -0,0 +1,27 @@
+from .base import LayerDecomposer
+from .qwen import QwenDecomposer
+
+# Registry mapping model type string (or class) to Decomposer
+DECOMPOSER_REGISTRY = {
+    # Keys should match what we expect in model config or name logic
+    "qwen2": QwenDecomposer,
+    "qwen3": QwenDecomposer,
+    "qwen": QwenDecomposer, # Generic fallback
+    "llama": QwenDecomposer, # Llama usually identical structure (PreNorm, RMS, MLP)
+}
+
+def get_decomposer(model_name_or_obj) -> LayerDecomposer:
+    """
+    Factory to return appropriate decomposer.
+    """
+    # Simple logic based on string for now
+    name = str(model_name_or_obj).lower()
+
+    if "qwen" in name:
+        return QwenDecomposer()
+    if "llama" in name:
+        return QwenDecomposer() # Re-use for now as structure is same
+
+    # Default fallback (hope compatibility)
+    print(f"Warning: No specific decomposer for {name}. Using Qwen/Llama default.")
+    return QwenDecomposer()

backend/models/manager.py

@@ -0,0 +1,196 @@
+import torch
+import importlib
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers.models.qwen3 import modeling_qwen3
+# Import other models as needed via conditional imports or a mapping
+try:
+    from lxt.efficient import monkey_patch
+except ImportError:
+    monkey_patch = None
+    print("Warning: lxt package not available. LRP attribution methods will be disabled.")
+import gc
+from .factory import get_decomposer
+
+class ModelManager:
+    """
+    Manages model loading, quantization, and patching.
+    """
+    def __init__(self):
+        self.model = None
+        self.tokenizer = None
+        self.model_name = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.decomposer = None
+
+        # Track active configuration for reloading
+        self.current_model_path = None
+        self.current_dtype = None
+        self.current_lrp_rule = None
+        self.current_quantization = False
+        self.current_revision = None
+
+    def load_model(self, model_path="Qwen/Qwen3-0.6B", quantization_4bit=False, dtype="auto", revision=None, lrp_rule=None):
+        """
+        Loads the model and tokenizer, applies monkey patches for LRP if lrp_rule is specified.
+        lrp_rule: None (no LRP), "Attn-LRP", or "CP-LRP" (Conservative Propagation)
+        """
+        if revision == "" or revision == "null":
+            revision = None
+
+        print(f"Loading model from {model_path} with revision={revision} and rule={lrp_rule}...")
+
+        # Store active configuration
+        self.current_model_path = model_path
+        self.current_dtype = dtype
+        self.current_lrp_rule = lrp_rule
+        self.current_quantization = quantization_4bit
+        self.current_revision = revision
+
+        # Free up memory if reloading
+        if self.model is not None:
+            del self.model
+            del self.tokenizer
+            torch.cuda.empty_cache()
+            gc.collect()
+
+        self.model_name = model_path.split('/')[-1]
+
+        # Initialize Decomposer
+        self.decomposer = get_decomposer(self.model_name)
+
+        # Apply Monkey Patch for Efficient LRP (only if lrp_rule is specified)
+        if lrp_rule is not None:
+            if monkey_patch is None:
+                print("Warning: lxt package not available. Cannot apply LRP patches. Loading model without LRP.")
+            else:
+                target_module = None
+                patch_map = None
+
+                lower_path = model_path.lower()
+                if "qwen3" in lower_path:
+                    importlib.reload(modeling_qwen3) # Reset to original classes to remove previous patches
+                    target_module = modeling_qwen3
+                    try:
+                        import lxt.efficient.models.qwen3 as lxt_qwen3
+                        importlib.reload(lxt_qwen3) # Reload to update class references from new modeling_qwen3
+                        patch_map = lxt_qwen3.cp_LRP if lrp_rule == "CP-LRP" else lxt_qwen3.attnLRP
+                    except ImportError as e:
+                        print(f"Warning: Could not import lxt.efficient.models.qwen3: {e}")
+
+                elif "olmo" in lower_path:
+                    try:
+                        from transformers.models.olmo3 import modeling_olmo3
+                        importlib.reload(modeling_olmo3)
+                        target_module = modeling_olmo3
+                        import lxt.efficient.models.olmo3 as lxt_olmo3
+                        importlib.reload(lxt_olmo3)
+                        patch_map = lxt_olmo3.cp_LRP if lrp_rule == "CP-LRP" else lxt_olmo3.attnLRP
+                    except ImportError as e:
+                        print(f"Warning: Could not import modeling_olmo3 or lxt module. LRP might fail. Error: {e}")
+
+                elif "qwen2" in lower_path:
+                     try:
+                         from transformers.models.qwen2 import modeling_qwen2
+                         importlib.reload(modeling_qwen2)
+                         target_module = modeling_qwen2
+                         import lxt.efficient.models.qwen2 as lxt_qwen2
+                         importlib.reload(lxt_qwen2)
+                         patch_map = lxt_qwen2.cp_LRP if lrp_rule == "CP-LRP" else lxt_qwen2.attnLRP
+                     except ImportError as e:
+                         print(f"Warning: Could not import qwen2 or lxt: {e}")
+
+                if target_module:
+                    if patch_map:
+                        monkey_patch(target_module, patch_map=patch_map, verbose=True)
+                        print(f"Applied LRP patches with rule: {lrp_rule}")
+                    else:
+                        monkey_patch(target_module, verbose=True) # Fallback to default
+                        print("Applied default LRP patches")
+        else:
+            # LRP not enabled - reload modules to remove any previous monkey patches
+            # so that the model is loaded with vanilla (unpatched) classes.
+            lower_path = model_path.lower()
+            if "qwen3" in lower_path:
+                importlib.reload(modeling_qwen3)
+                print("Reloaded modeling_qwen3 to remove LRP patches")
+            elif "olmo" in lower_path:
+                try:
+                    from transformers.models.olmo3 import modeling_olmo3
+                    importlib.reload(modeling_olmo3)
+                    print("Reloaded modeling_olmo3 to remove LRP patches")
+                except ImportError:
+                    pass
+            elif "qwen2" in lower_path:
+                try:
+                    from transformers.models.qwen2 import modeling_qwen2
+                    importlib.reload(modeling_qwen2)
+                    print("Reloaded modeling_qwen2 to remove LRP patches")
+                except ImportError:
+                    pass
+            print("LRP not enabled - model loaded without attribution patches")
+
+        # Add else if for other models supported by lxt
+
+        # Map string dtype to torch dtype
+        torch_dtype = "auto"
+        bnb_dtype = torch.bfloat16 # Default for 4bit compute
+
+        if dtype == "float16":
+            torch_dtype = torch.float16
+            bnb_dtype = torch.float16
+        elif dtype == "bfloat16":
+            torch_dtype = torch.bfloat16
+            bnb_dtype = torch.bfloat16
+        elif dtype == "float32":
+            torch_dtype = torch.float32
+            bnb_dtype = torch.float32
+
+        # Quantization Config
+        quantization_config = None
+        if quantization_4bit:
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=bnb_dtype,
+            )
+
+        # Load Model
+        if "qwen3" in model_path.lower():
+            self.model = modeling_qwen3.Qwen3ForCausalLM.from_pretrained(
+                model_path,
+                device_map=self.device,
+                torch_dtype=torch_dtype,
+                quantization_config=quantization_config,
+                revision=revision
+            )
+        else:
+             # Fallback for generic loading if specific class fails
+             self.model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                device_map=self.device,
+                torch_dtype=torch_dtype,
+                quantization_config=quantization_config,
+                revision=revision
+            )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, revision=revision)
+
+        # Prepare model for LRP
+        self.model.eval() # Use eval usually, but test.ipynb uses train() + gradients
+        # test.ipynb: model.train(), gradient_checkpointing_enable(), requires_grad=False
+
+        # "model.train()" is often needed for Gradient Checkpointing to work in HF
+        self.model.train()
+        self.model.gradient_checkpointing_enable()
+
+        # Deactivate gradients on parameters
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        print(f"Model {self.model_name} loaded successfully on {self.device}")
+        return self.model_name
+
+    def get_model(self):
+        return self.model
+
+    def get_tokenizer(self):
+        return self.tokenizer

backend/models/qwen.py

@@ -0,0 +1,55 @@
+import torch.nn as nn
+from .base import LayerDecomposer
+
+class QwenDecomposer(LayerDecomposer):
+    """
+    Decomposer for Qwen2/Qwen3 models.
+    Assumes Pre-Norm architecture:
+    Part 1: x -> InputNorm -> Attn -> Add (resid_mid)
+    Part 2: resid_mid -> PostAttnNorm -> MLP -> Add (resid_post)
+    """
+
+    def get_mid_activation_module(self, layer_module):
+        # Qwen models usually have 'post_attention_layernorm'
+        return getattr(layer_module, "post_attention_layernorm", None)
+
+    def forward_part1(self, layer_module, hidden_states, position_embeddings=None, attention_mask=None):
+        norm = getattr(layer_module, "input_layernorm", None)
+        # Try both generic 'self_attn' and Qwen specific naming if needed,
+        # though HF implementation usually maps to 'self_attn'
+        attn = getattr(layer_module, "self_attn", None)
+
+        if not norm or not attn:
+            # Fallback for some versions or wrapped modules
+            # Check named children
+            children = dict(layer_module.named_modules())
+            attn = attn or children.get("self_attn")
+
+        if not norm or not attn:
+             raise AttributeError(f"QwenDecomposer: Layer module {type(layer_module)} missing input_layernorm or self_attn")
+
+        # Norm
+        norm_out = norm(hidden_states)
+
+        # Attn
+        if position_embeddings is not None:
+             # Qwen/Llama usually accept position_embeddings
+             attn_out = attn(norm_out, position_embeddings=position_embeddings, attention_mask=attention_mask)
+        else:
+             attn_out = attn(norm_out, attention_mask=attention_mask)
+
+        if isinstance(attn_out, tuple): attn_out = attn_out[0]
+
+        return hidden_states + attn_out
+
+    def forward_part2(self, layer_module, hidden_states):
+        norm = getattr(layer_module, "post_attention_layernorm", None)
+        mlp = getattr(layer_module, "mlp", None)
+
+        if not norm or not mlp:
+             raise AttributeError("QwenDecomposer: Layer module missing post_attention_layernorm or mlp")
+
+        norm_out = norm(hidden_states)
+        mlp_out = mlp(norm_out)
+
+        return hidden_states + mlp_out

backend/token_locator_prompts/err_token_localization.yaml

@@ -0,0 +1,75 @@
+version: v1
+system: |-
+  - You are an expert in localizing error tokens in model-generated text completions.
+  - You will be provided with a prompt and a model-generated completion, and your task is to localize and output the FIRST token in the completion that is erroneous or causes the error.
+  - As the prompt and completion are taken from a particular dataset, refer to Section "Dataset Description" for a brief overview of the dataset.
+  - Refer to Section "On the Input Format" for details on how the prompt and completion will be presented to you.
+  - Follow the "General Guidance (Dataset Agnostic)" section for general instructions on how to approach the task of error token localization.
+  - Refer to the "Dataset Specific Instructions" section for any additional instructions that are specific to this dataset.
+  - Finally, adhere to the "On the Output Format" section to ensure your response is structured correctly.
+  - Note that the dataset specific sections (i.e., "Dataset Description" and "Dataset Specific Instructions") can be missing for some datasets. In such cases, simply skip those sections and proceed with the rest of the instructions.
+
+  # Dataset Description
+  {{dataset_description}}
+  
+  # On the Input Format
+  - The input will be provided in the following JSON format:
+    {
+      "prompt": "<The prompt text here>",
+      "completion": "<The model-generated completion text here>",
+      "indexed_completion": "<The model-generated completion text with each token indexed, i.e., `token_1[index_1]token_2[index_2]`>"
+      "ground_truth": "<The ground truth answer text here>"  (This field may be absent in some cases)
+    }
+  - An example of the input format is as follows:
+    {
+      "prompt": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n",
+      "completion": "<think> okay, let me see. The capital of France is Lyon. </think>\n\Lyon.",
+      "indexed_completion": "<think>[0] okay,[1] let[2] me[3] see.[4] The[5] capital[6] of[7] France[8] is[9] Lyon.[10] </think>[11]\n\n[12] Lyon[13] .[14]",
+      "ground_truth": "Paris, the capital of France."
+    }
+  - The "indexed_completion" field provides a tokenized version of the completion, where each token is followed by its index in square brackets. This will help you localize the position of tokens in the completion.
+  - Both the "prompt" and "completion" fields may contain special tokens, e.g., "<|im_start|>", "<|im_end|>", "<think>", "</think>", etc. These tokens are part of the model's output format, specifically,
+    - "<|im_start|>" and "<|im_end|>" denote the start and end of a message in a multi-turn conversation, e.g., "<|im_start|>user\n xxx <|im_end|>" indicates the start of a user message, while "<|im_start|>assistant\n xxx <|im_end|>" indicates the start of an assistant message.
+    - "<think>" and "</think>" denote the start and end of the model's internal reasoning process.
+  - More on the internal thought process of the model:
+    - The internal thought process of the model is represented by the reasoning tokens between the "<think>" and "</think>" tokens.
+    - This part reflects the model's reasoning steps before arriving at the final answer.
+    - The final answer to the prompt is the text that comes AFTER the "</think>" token.
+
+  # General Guidance (Dataset Agnostic)
+  - Refer to the following steps for localizing an erroneous token in the model-generated completion:
+    - Step 1: Carefully read dataset description and the provided prompt to understand the context and requirements.
+    - Step 2: If "completion" and "indexed_completion" fields contains reasoning tokens (i.e., tokens between "<think>" and "</think>"), seperate the reasoning part from the final answer part.
+    - Step 3: Examine the final answer part of the completion first to identify any errors. If no errors are found in this part, skip the rest of the steps and output "<CORRECT>" as the final answer.
+    - Step 4: No matther whether the final answer part contains errors, proceed to examine the reasoning part for potential errors.
+    - Step 5: Once an erroneous token is identified, use the "indexed_completion" field to find its index.
+  - Addiational strict requirements for localizing the erroneous token:
+    - Always report the FIRST erroneous token in the "completion" or "indexed_completion" field.
+    - The localized erroneous token MUST NOT be the special tokens used for formatting (e.g., "<|im_start|>", "<|im_end|>", "<think>", "</think>").
+    - NEVER treat a token is erroneous just because it is a subword or punctuation mark, as the indexed token str does NOT have to be a single word; it can be a subword or punctuation mark as per the tokenization used by the model.
+  - On the difference between examining the final answer part and the reasoning part:
+    - For the final answer part (i.e., the part AFTER "</think>"):
+      - ALWAYS examine whether the tokens in this part fulfill the format/style requirements specified in the prompt. If you identify a token in this part that violates the format/style requirements, you SHOULD consider it as a potential erroneous token.
+      - ALSO examine the factual correctness and logical consistency of the tokens in this part. If you identify a token in this part that is factually incorrect or logically inconsistent with the prompt, you SHOULD consider it as a potential erroneous token.
+    - For the reasoning part (i.e., the part BETWEEN "<think>" and "</think>"):
+      - NEVER examine whether the reasoning tokens have fulfilled the format/style requirements specified in the prompt, as the reasoning tokens between "<think>" and "</think>" are often in a free-form text format. For example, if the prompt requires the model to "respond in a poem format", you MUST NOT consider a token in the reasoning part as erroneous just because the reasoning tokens are not in a poem format. 
+      - Focus on the factual correctness and logical consistency of the reasoning tokens. If you identify a token in this part that is factually incorrect or logically inconsistent with the prompt, you SHOULD consider it as a potential erroneous token.
+
+  # Dataset Specific Instructions
+  {{dataset_specific_instructions}}
+
+  # On the Output Format
+  - The output MUST be in the following JSON format:
+    {
+      "error_token": "<The localized erroneous token>",
+      "token_index": <The index of the erroneous token in the completion>,
+      "explanation": "<A brief explanation of why this token is considered erroneous. Less than 50 words.>"
+    }
+  - An example of the output format is as follows:
+    {
+      "error_token": "Lyon",
+      "token_index": 10,
+      "explanation": "The token 'Lyon' is erroneous because the correct capital of France is Paris, not Lyon. This indicates a factual error in the model's completion."
+    }
+  - Ensure that the "token_index" corresponds to the index provided in the "indexed_completion" field.
+  - The "explanation" should be concise yet informative, providing enough context to justify the localization of the error token.

frontend/css/style.css

@@ -0,0 +1,274 @@
+body {
+    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+    margin: 0;
+    padding: 0;
+    background-color: #f5f7fa;
+    color: #333;
+}
+
+header {
+    background-color: #2c3e50;
+    color: white;
+    padding: 1rem 2rem;
+}
+
+header h1 {
+    margin: 0;
+    font-size: 1.5rem;
+}
+
+main {
+    padding: 2rem;
+    max-width: 1200px;
+    margin: 0 auto;
+}
+
+.panel {
+    background: white;
+    border-radius: 8px;
+    padding: 1.5rem;
+    margin-bottom: 2rem;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+}
+
+.hidden {
+    display: none;
+}
+
+.control-group {
+    margin-bottom: 1rem;
+    display: flex;
+    align-items: center;
+    gap: 1rem;
+    flex-wrap: wrap;
+}
+
+label {
+    font-weight: 500;
+}
+
+input[type="text"], input[type="number"], textarea, select {
+    padding: 0.5rem;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    font-size: 1rem;
+}
+
+textarea {
+    width: 100%;
+    max-width: 600px;
+}
+
+button {
+    background-color: #3498db;
+    color: white;
+    border: none;
+    padding: 0.5rem 1rem;
+    border-radius: 4px;
+    cursor: pointer;
+    font-size: 1rem;
+    transition: background 0.2s;
+}
+
+button:hover {
+    background-color: #2980b9;
+}
+
+button:disabled {
+    background-color: #bdc3c7;
+    cursor: not-allowed;
+}
+
+.status {
+    font-style: italic;
+    color: #7f8c8d;
+}
+
+.status.success {
+    color: #27ae60;
+}
+.status.error {
+    color: #c0392b;
+}
+
+/* Logits Table */
+.scrollable-table {
+    max-height: 300px;
+    overflow-y: auto;
+    border: 1px solid #eee;
+}
+
+table {
+    width: 100%;
+    border-collapse: collapse;
+}
+
+th, td {
+    padding: 0.35rem 0.75rem; /* Reduced vertical padding */
+    text-align: left;
+    border-bottom: 1px solid #eee;
+}
+
+th {
+    background-color: #f8f9fa;
+    position: sticky;
+    top: 0;
+}
+
+tr:hover {
+    background-color: #f1f1f1;
+}
+
+.btn-select-ref {
+    padding: 0.25rem 0.5rem;
+    font-size: 0.8rem;
+    background-color: #9b59b6;
+}
+
+.btn-select-ref.selected {
+    background-color: #27ae60;
+}
+
+/* Circuit Controls */
+.controls-row {
+    display: flex;
+    gap: 2rem;
+    margin-bottom: 1.5rem;
+    border-bottom: 1px solid #eee;
+    padding-bottom: 1rem;
+}
+
+.control-col {
+    flex: 1;
+}
+
+.sub-opt {
+    margin-left: 1rem;
+    margin-top: 0.5rem;
+    font-size: 0.9rem;
+}
+
+.setting-item {
+    margin-bottom: 0.8rem;
+}
+
+.input-row {
+    display: flex;
+    align-items: center;
+    gap: 10px;
+    margin-top: 4px;
+}
+
+.input-row input[type="range"] {
+    flex: 1;
+}
+
+/* Canvas */
+#canvas-container {
+    position: relative;
+    border: 1px solid #ddd;
+    background: white;
+    border-radius: 4px;
+    overflow-x: auto;
+}
+
+canvas {
+    display: block;
+    cursor: crosshair;
+}
+
+#tooltip {
+    position: fixed;
+    background: rgba(0, 0, 0, 0.85);
+    color: white;
+    padding: 8px;
+    border-radius: 4px;
+    font-size: 12px;
+    pointer-events: none;
+    opacity: 0;
+    transition: opacity 0.2s;
+    z-index: 1000;
+}
+
+/* Legend */
+#vis-legend {
+    margin-top: 1rem;
+    font-size: 0.9rem;
+}
+.legend-box {
+    display: inline-block;
+    width: 15px;
+    height: 15px;
+    margin-right: 5px;
+    margin-left: 10px;
+    vertical-align: middle;
+}
+.legend-box.pos { background: rgba(255, 0, 0, 0.6); }
+.legend-box.neg { background: rgba(0, 0, 255, 0.6); }
+
+/* Input Attribution Styles */
+.attribution-text-box {
+    line-height: 1.6;
+    font-size: 1.1em;
+    padding: 15px;
+    border: 1px solid #ddd;
+    border-radius: 4px;
+    background-color: #fff;
+    white-space: pre-wrap; /* Preserve spacing */
+    font-family: monospace;
+}
+
+.token-span {
+    padding: 2px 0;
+    margin: 0;
+    border-radius: 2px;
+    transition: background-color 0.2s;
+    cursor: default;
+}
+
+.token-span:hover {
+    outline: 1px solid #333;
+    position: relative;
+    z-index: 10;
+}
+
+/* Loading Overlay */
+#loading-overlay {
+    position: fixed;
+    top: 0; 
+    left: 0; 
+    width: 100%; 
+    height: 100%;
+    background: rgba(255, 255, 255, 0.8);
+    display: flex;
+    flex-direction: column;
+    justify-content: center;
+    align-items: center;
+    z-index: 9999;
+    backdrop-filter: blur(2px);
+}
+
+#loading-overlay.hidden {
+    display: none !important;
+}
+
+.spinner {
+    border: 6px solid #f3f3f3;
+    border-top: 6px solid #3498db;
+    border-radius: 50%;
+    width: 50px;
+    height: 50px;
+    animation: spin 1s linear infinite;
+    margin-bottom: 15px;
+}
+
+@keyframes spin {
+    0% { transform: rotate(0deg); }
+    100% { transform: rotate(360deg); }
+}
+
+#loading-message {
+    font-size: 1.2rem;
+    font-weight: 500;
+    color: #2c3e50;
+}

frontend/index.html

@@ -0,0 +1,702 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>NeuralPostmortem: An Interactive Framework for LLM Failures Analysis</title>
+    <link rel="stylesheet" href="css/style.css">
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/jspdf/2.5.1/jspdf.umd.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
+    <style>
+        /* Additional styles for step-by-step UI */
+        .step-container {
+            margin-bottom: 30px;
+            border: 2px solid #e0e0e0;
+            border-radius: 8px;
+            overflow: hidden;
+        }
+        
+        .step-header {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 15px 20px;
+            cursor: pointer;
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+            font-weight: bold;
+            font-size: 1.1em;
+        }
+        
+        .step-header.completed {
+            background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
+        }
+        
+        .step-header.active {
+            background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
+        }
+        
+        .step-header.disabled {
+            background: #ccc;
+            cursor: not-allowed;
+            opacity: 0.6;
+        }
+        
+        .step-number {
+            background: rgba(255,255,255,0.3);
+            width: 35px;
+            height: 35px;
+            border-radius: 50%;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            font-size: 1.2em;
+            margin-right: 15px;
+        }
+        
+        .step-title {
+            flex: 1;
+        }
+        
+        .step-status {
+            font-size: 0.9em;
+            padding: 5px 15px;
+            background: rgba(255,255,255,0.2);
+            border-radius: 15px;
+        }
+        
+        .step-body {
+            padding: 20px;
+            background: #f9f9f9;
+        }
+        
+        .step-body.collapsed {
+            display: none;
+        }
+        
+        .subsection {
+            margin: 15px 0;
+            padding: 15px;
+            background: white;
+            border-radius: 5px;
+            border-left: 4px solid #667eea;
+        }
+        
+        .subsection h4 {
+            margin-top: 0;
+            color: #667eea;
+        }
+        
+        .action-buttons {
+            margin-top: 20px;
+            display: flex;
+            gap: 10px;
+        }
+        
+        .btn-primary {
+            background: #667eea;
+            color: white;
+            border: none;
+            padding: 10px 20px;
+            border-radius: 5px;
+            cursor: pointer;
+            font-size: 1em;
+        }
+        
+        .btn-primary:hover {
+            background: #5568d3;
+        }
+        
+        .btn-primary:disabled {
+            background: #ccc;
+            cursor: not-allowed;
+        }
+        
+        .btn-secondary {
+            background: #6c757d;
+            color: white;
+            border: none;
+            padding: 10px 20px;
+            border-radius: 5px;
+            cursor: pointer;
+            font-size: 1em;
+        }
+        
+        .vote-details {
+            margin-top: 10px;
+            padding: 10px;
+            background: #f0f0f0;
+            border-radius: 5px;
+        }
+        
+        .vote-item {
+            padding: 5px;
+            margin: 5px 0;
+            background: white;
+            border-radius: 3px;
+        }
+
+        /* Method Selector Styles */
+        .method-btn {
+            background: rgba(255,255,255,0.15);
+            color: #ccc;
+            border: 1px solid rgba(255,255,255,0.25);
+            padding: 6px 16px;
+            border-radius: 20px;
+            cursor: pointer;
+            font-size: 0.9em;
+            font-weight: 500;
+            transition: all 0.2s ease;
+        }
+        .method-btn:hover {
+            background: rgba(255,255,255,0.25);
+            color: white;
+        }
+        .method-btn.active {
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            border-color: #764ba2;
+            box-shadow: 0 2px 8px rgba(102,126,234,0.4);
+        }
+        .method-btn:disabled {
+            opacity: 0.4;
+            cursor: not-allowed;
+        }
+
+        /* Method-specific visibility */
+        .attnlrp-only { display: block; }
+        .nnsight-only { display: none; }
+        .pyvene-only { display: none; }
+
+        body[data-method="nnsight"] .attnlrp-only { display: none; }
+        body[data-method="nnsight"] .nnsight-only { display: block; }
+
+        body[data-method="pyvene"] .attnlrp-only { display: none; }
+        body[data-method="pyvene"] .pyvene-only { display: block; }
+    </style>
+</head>
+<body>
+    <div id="app">
+        <header>
+            <h1>🔍 NeuralPostmortem</h1>
+            <p style="margin: 10px 0; color: #666;">An Interactive Framework for LLM Failure Analysis</p>
+            <!-- Method Selector Bar -->
+            <div id="method-bar" style="margin-top: 12px; display: flex; align-items: center; gap: 15px;">
+                <span style="font-weight: 600; color: #ccc; font-size: 0.95em;">Attribution Method:</span>
+                <div id="method-selector" style="display: flex; gap: 6px;">
+                    <button class="method-btn active" data-method="attnlrp" data-port="8000">Attn-LRP</button>
+                    <button class="method-btn" data-method="nnsight" data-port="8001">NNsight</button>
+                    <button class="method-btn" data-method="pyvene" data-port="8002">Pyvene</button>
+                </div>
+                <span id="method-status" style="font-size: 0.85em; color: #8e8e8e; margin-left: 10px;">
+                    Backend: <code id="current-api-base" style="color: #38ef7d;">localhost:8000</code>
+                </span>
+            </div>
+        </header>
+
+        <main style="max-width: 1400px; margin: 0 auto; padding: 20px;">
+            
+            <!-- Step 1: Case Setup & Model Configuration -->
+            <div class="step-container" id="step-1">
+                <div class="step-header active" onclick="toggleStep(1)">
+                    <div style="display: flex; align-items: center;">
+                        <span class="step-number">1</span>
+                        <span class="step-title">Case Setup & Model Configuration</span>
+                    </div>
+                    <span class="step-status" id="status-1">Pending</span>
+                </div>
+                <div class="step-body" id="body-1">
+                    <div class="subsection">
+                        <h4>📁 Load from Trace (Auto-fills all fields)</h4>
+                        <div class="control-group">
+                            <label>Dataset: 
+                                <select id="trace-dataset">
+                                    <option value="">-- Select --</option>
+                                    <!-- Populated by JS -->
+                                </select>
+                            </label>
+                            <label>Trace File: 
+                                <select id="trace-file" disabled>
+                                    <option value="">-- Select Dataset First --</option>
+                                </select>
+                            </label>
+                        </div>
+                    </div>
+
+                    <div class="subsection">
+                        <h4>✍️ Manual Input</h4>
+                        <div class="control-group">
+                            <label for="input-prompt"><strong>Prompt:</strong></label>
+                            <textarea id="input-prompt" rows="4" placeholder="Enter your prompt here..."></textarea>
+                        </div>
+                        <div class="control-group">
+                            <label for="input-completion"><strong>Completion:</strong></label>
+                            <textarea id="input-completion" rows="4" placeholder="Enter the completion to analyze..."></textarea>
+                        </div>
+                        <div class="control-group">
+                            <label for="input-ground-truth"><strong>Ground Truth (Optional):</strong></label>
+                            <textarea id="input-ground-truth" rows="2" placeholder="Expected correct answer..."></textarea>
+                        </div>
+                    </div>
+
+                    <div class="subsection">
+                        <h4>🤖 Model Configuration</h4>
+                        <div class="control-group">
+                            <div style="display: flex; gap: 15px; flex-wrap: wrap;">
+                                <label>Series: 
+                                    <select id="model-series">
+                                        <option value="Qwen3">Qwen3</option>
+                                        <option value="Qwen2">Qwen2</option>
+                                        <option value="OLMo3">OLMo3</option>
+                                    </select>
+                                </label>
+                                <label>Model: 
+                                    <select id="model-select">
+                                        <!-- Populated dynamically -->
+                                    </select>
+                                </label>
+                                <label>DType:
+                                    <select id="model-dtype">
+                                        <option value="float16">float16</option>
+                                        <option value="bfloat16" selected>bfloat16</option>
+                                        <option value="float32">float32</option>
+                                    </select>
+                                </label>
+                                <label>
+                                    <input type="checkbox" id="quant-4bit"> 4-bit Quantization
+                                </label>
+                            </div>
+                        </div>
+                        <div class="control-group">
+                            <label>Model Path:
+                                <input type="text" id="model-path" value="Qwen/Qwen3-0.6B" style="width: 400px;">
+                            </label>
+                        </div>
+                    </div>
+
+                    <div class="action-buttons">
+                        <button class="btn-primary" id="btn-load-model">Load Model & Proceed</button>
+                        <span id="model-load-status" style="margin-left: 15px; line-height: 40px; color: #666;"></span>
+                    </div>
+                </div>
+            </div>
+
+            <!-- Step 2: Contrast Token Pair Identification -->
+            <div class="step-container" id="step-2">
+                <div class="step-header disabled" onclick="toggleStep(2)">
+                    <div style="display: flex; align-items: center;">
+                        <span class="step-number">2</span>
+                        <span class="step-title">Contrast Token Pair Identification</span>
+                    </div>
+                    <span class="step-status" id="status-2">Locked</span>
+                </div>
+                <div class="step-body collapsed" id="body-2">
+                    
+                    <!-- Subsection 3.1: Error Token Localization -->
+                    <div class="subsection">
+                        <h4>🎯 Target Token Localization & Text Truncation</h4>
+                        <p style="color: #666; margin-bottom: 15px;">
+                            Use multiple LLM validators to identify the first erroneous token in the completion.
+                        </p>
+                        
+                        <div class="control-group">
+                            <label style="display: flex; align-items: center; gap: 10px;">
+                                <input type="checkbox" id="use-llm-search" checked> Use LLM search (validators)
+                            </label>
+                        </div>
+                        
+                        <div class="control-group" style="margin-top: 10px;" id="manual-chunk-group" hidden>
+                            <label><strong>Manual Localization(s):</strong> (Include both prompt and completion)</label>
+                            <textarea id="manual-chunks" rows="3" style="width: 100%; font-family: monospace;"></textarea>
+                        </div>
+                        
+                        <div class="control-group" style="margin-top: 10px;">
+                            <label><strong>Validator Models:</strong></label>
+                            <div style="display: flex; gap: 10px; flex-wrap: wrap; margin: 10px 0;">
+                                <label><input type="checkbox" id="validator-gpt52" checked> GPT-5.2</label>
+                                <label><input type="checkbox" id="validator-gptoss" checked> GPT-OSS-120B</label>
+                                <label><input type="checkbox" id="validator-deepseek" checked> DeepSeek-V3.2</label>
+                                <label><input type="checkbox" id="validator-kimi"> Kimi-K2-Thinking</label>
+                            </div>
+                        </div>
+
+                        <div class="action-buttons">
+                            <button class="btn-primary" id="btn-locate-error">Locate Target Token</button>
+                            <span id="error-locate-status" style="margin-left: 15px; line-height: 40px; color: #666;"></span>
+                        </div>
+
+                        <!-- Results -->
+                        <div id="error-location-results" class="hidden" style="margin-top: 20px;">
+                            <h5 style="color: #667eea;">📊 Voting Results:</h5>
+                            <div id="vote-details" class="vote-details">
+                                <!-- Populated by JS -->
+                            </div>
+
+                            <h5 style="color: #667eea; margin-top: 15px;">✂️ Truncated Text:</h5>
+                            <div class="control-group">
+                                <textarea id="truncated-text" rows="5" style="width: 100%; font-family: monospace;"></textarea>
+                            </div>
+
+                            <h5 style="color: #667eea;">💡 Explanation:</h5>
+                            <div id="error-explanation" style="padding: 10px; background: #fff3cd; border-radius: 5px; white-space: pre-wrap;"></div>
+
+                            <div class="action-buttons" style="margin-top: 15px;">
+                                <button class="btn-primary" id="btn-accept-truncation">Accept & Use Truncated Text</button>
+                                <button class="btn-secondary" id="btn-edit-truncation">Edit Truncation</button>
+                                <span id="accept-status" style="margin-left: 15px; line-height: 40px; color: #28a745; font-weight: bold;"></span>
+                            </div>
+                        </div>
+                    </div>
+
+                    <!-- Subsection 3.2: Contrast Token Selection -->
+                    <div class="subsection" id="contrast-selection-section" class="hidden">
+                        <h4>🔄 Contrast Token Selection</h4>
+                        <p style="color: #666; margin-bottom: 15px;">
+                            After error token localization, identify the contrast token (correct alternative).
+                        </p>
+
+                        <div class="control-group">
+                            <label><strong>Selection Method:</strong></label>
+                            <div style="margin: 10px 0;">
+                                <label style="display: block; margin: 5px 0;">
+                                    <input type="radio" name="contrast-method" value="manual" checked> 
+                                    Manual Selection (from Top-K predictions)
+                                </label>
+                                <label style="display: block; margin: 5px 0;">
+                                    <input type="radio" name="contrast-method" value="beam"> 
+                                    Beam Search (automatic exploration)
+                                </label>
+                            </div>
+                        </div>
+
+                        <div id="manual-contrast-section">
+                            <div class="action-buttons">
+                                <button class="btn-primary" id="btn-compute-topk">Compute Top-K Predictions</button>
+                                <label style="margin-left: 15px;">
+                                    K = <input type="number" id="topk-value" value="10" min="1" max="100" style="width: 60px;">
+                                </label>
+                            </div>
+
+                            <!-- Top-K Table -->
+                            <div id="topk-results" class="hidden" style="margin-top: 20px;">
+                                <h5>Top-K Token Predictions:</h5>
+                                <div class="scrollable-table">
+                                    <table id="topk-table">
+                                        <thead>
+                                            <tr>
+                                                <th>Rank</th>
+                                                <th>Token</th>
+                                                <th>Logit</th>
+                                                <th>Action</th>
+                                            </tr>
+                                        </thead>
+                                        <tbody>
+                                            <!-- Populated by JS -->
+                                        </tbody>
+                                    </table>
+                                </div>
+                            </div>
+                        </div>
+
+                        <div id="beam-contrast-section" class="hidden">
+                            <p style="color: #999; font-style: italic;">Beam search functionality - Coming soon...</p>
+                            <div class="action-buttons">
+                                <button class="btn-primary" disabled>Run Beam Search</button>
+                            </div>
+                        </div>
+
+                        <!-- Selected Contrast Token Display -->
+                        <div id="contrast-selected" class="hidden" style="margin-top: 20px; padding: 15px; background: #d4edda; border-radius: 5px;">
+                            <h5 style="margin-top: 0; color: #155724;">✓ Contrast Token Pair Identified:</h5>
+                            <p><strong>Error Token:</strong> <span id="display-error-token">-</span></p>
+                            <p><strong>Contrast Token:</strong> <span id="display-contrast-token">-</span></p>
+                            <p><strong>Logit Difference:</strong> <span id="display-logit-diff">-</span></p>
+                        </div>
+                    </div>
+
+                    <div class="action-buttons" style="margin-top: 20px;">
+                        <button class="btn-primary" id="btn-proceed-to-lrp" disabled>Proceed to Input Attribution Analysis</button>
+                    </div>
+                </div>
+            </div>
+
+            <!-- Step 3: Input Attribution -->
+            <div class="step-container" id="step-3">
+                <div class="step-header disabled" onclick="toggleStep(3)">
+                    <div style="display: flex; align-items: center;">
+                        <span class="step-number">3</span>
+                        <span class="step-title">Input Attribution Analysis</span>
+                    </div>
+                    <span class="step-status" id="status-3">Locked</span>
+                </div>
+                <div class="step-body collapsed" id="body-3">
+                    
+                    <div class="subsection attnlrp-only">
+                        <h4>⚙️ Input Attribution Configuration</h4>
+                        <div class="control-group">
+                            <label>Rule:
+                                <select id="lrp-rule">
+                                    <option value="Attn-LRP" selected>Attn-LRP</option>
+                                    <option value="CP-LRP">CP-LRP</option>
+                                    <option value="Gradient">Gradient (Vanilla)</option>
+                                </select>
+                            </label>
+                            <label style="margin-left: 20px;">
+                                <input type="checkbox" id="capture-mid"> Attn/MLP Split
+                            </label>
+                        </div>
+                    </div>
+
+                    <div class="subsection nnsight-only">
+                        <h4>⚙️ NNsight Attribution Configuration</h4>
+                        <p style="color: #666;">NNsight uses neural network surgery to trace attributions through the model's computation graph.</p>
+                    </div>
+
+                    <div class="subsection pyvene-only">
+                        <h4>⚙️ Pyvene Attribution Configuration</h4>
+                        <p style="color: #666;">Pyvene uses causal interventions to identify token attributions.</p>
+                    </div>
+
+                    <div class="subsection">
+                        <h4>🎨 Input Attribution Visualization</h4>
+                        <div class="control-group">
+                            <label>Mode:
+                                <select id="attr-mode">
+                                    <option value="max_logit">Max Logit</option>
+                                    <option value="logit_diff" selected>Logit Difference</option>
+                                </select>
+                            </label>
+                            <label style="margin-left: 15px;">Strategy:
+                                <select id="attr-strategy">
+                                    <option value="by_topk_avg">By Top-K Avg</option>
+                                    <option value="demean">Demean</option>
+                                    <option value="by_ref_token" selected>By Reference Token</option>
+                                </select>
+                            </label>
+                        </div>
+
+                        <div class="action-buttons">
+                            <button class="btn-primary" id="btn-compute-attribution">Compute Input Attribution</button>
+                        </div>
+
+                        <div id="attribution-results" class="hidden" style="margin-top: 20px;">
+                            <h5>Attribution Map:</h5>
+                            <div id="attribution-display" style="padding: 15px; background: white; border: 1px solid #ddd; border-radius: 5px; font-family: monospace; line-height: 1.8;">
+                                <!-- Populated by JS -->
+                            </div>
+                            <div class="action-buttons" style="margin-top: 10px;">
+                                <button class="btn-secondary" id="btn-save-attr-png" style="background-color: #2ecc71;">Save PNG</button>
+                                <button class="btn-secondary" id="btn-save-attr-pdf" style="background-color: #e74c3c;">Save PDF</button>
+                            </div>
+                            <!-- Positive Relevance Sorted Indices -->
+                            <div id="positive-relevance-sorted" class="hidden" style="margin-top: 18px;">
+                                <h5 style="margin: 0 0 8px 0; color: #667eea;">Positive Relevance Tokens (Sorted by Relevance, Descending)</h5>
+                                <div id="pos-rel-indices" style="padding: 10px; background: #f8f9fa; border: 1px solid #ddd; border-radius: 4px; font-family: monospace; font-size: 0.95em; word-break: break-all; line-height: 1.8;"></div>
+                                <h5 style="margin: 12px 0 8px 0; color: #7f8c8d; font-size: 0.9em;">Corresponding Tokens (for verification)</h5>
+                                <div id="pos-rel-tokens" style="padding: 10px; background: #fefefe; border: 1px solid #eee; border-radius: 4px; font-family: monospace; font-size: 0.9em; word-break: break-all; line-height: 1.8; color: #555;"></div>
+                            </div>
+                        </div>
+                    </div>
+
+                    <!-- Perturbation Evaluation Subsection -->
+                    <div class="subsection">
+                        <h4>🧪 Perturbation Evaluation (Ground-Truth Comparison)</h4>
+                        <p style="color: #666; margin-bottom: 15px;">
+                            Evaluate attribution quality by zeroing out top-attributed tokens and checking if the model's error is fixed.
+                            A better attribution method should fix the error at smaller K values.
+                        </p>
+                        <div class="control-group">
+                            <label>K values (comma-separated):
+                                <input type="text" id="perturbation-k-values" value="1, 3, 5, 10" style="width: 200px;">
+                            </label>
+                        </div>
+                        <div class="action-buttons">
+                            <button class="btn-primary" id="btn-run-perturbation" disabled>Run Perturbation Evaluation</button>
+                            <span id="perturbation-status" style="margin-left: 15px; line-height: 40px; color: #666;"></span>
+                        </div>
+                        <div id="perturbation-results" class="hidden" style="margin-top: 20px;">
+                            <h5>Perturbation Results:</h5>
+                            <div class="scrollable-table">
+                                <table id="perturbation-table">
+                                    <thead>
+                                        <tr>
+                                            <th>K</th>
+                                            <th>Perturbed Tokens</th>
+                                            <th>New Top-1 Token</th>
+                                            <th>Error Fixed?</th>
+                                            <th>Target Rank After</th>
+                                            <th>Logit Change</th>
+                                        </tr>
+                                    </thead>
+                                    <tbody>
+                                        <!-- Populated by JS -->
+                                    </tbody>
+                                </table>
+                            </div>
+                        </div>
+
+                        <!-- Manual Perturbation Sub-section -->
+                        <div style="margin-top: 25px; padding-top: 20px; border-top: 2px dashed #dee2e6;">
+                            <h5 style="color: #667eea; margin-top: 0;">🖱️ Manual Token Perturbation</h5>
+                            <p style="color: #666; font-size: 0.9em; margin-bottom: 12px;">
+                                Click tokens in the <strong>Attribution Map</strong> above to select/deselect which tokens to mask.
+                                Selected tokens will be highlighted with a dashed border.
+                            </p>
+                            <div style="margin-bottom: 12px;">
+                                <strong style="font-size: 0.9em;">Selected Tokens:</strong>
+                                <div id="manual-perturb-selected" style="display: flex; flex-wrap: wrap; gap: 6px; margin-top: 8px; min-height: 32px; padding: 8px; background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 5px;">
+                                    <span style="color: #999; font-size: 0.85em; line-height: 24px;" id="manual-perturb-placeholder">No tokens selected</span>
+                                </div>
+                            </div>
+                            <div class="action-buttons">
+                                <button class="btn-secondary" id="btn-clear-manual-selection" style="background: #95a5a6;">Clear Selection</button>
+                                <button class="btn-primary" id="btn-run-manual-perturbation" disabled>Run Manual Perturbation</button>
+                                <span id="manual-perturbation-status" style="margin-left: 15px; line-height: 40px; color: #666;"></span>
+                            </div>
+                            <div id="manual-perturbation-result" class="hidden" style="margin-top: 15px;">
+                                <!-- Populated by JS -->
+                            </div>
+                        </div>
+                    </div>
+
+                    <div class="action-buttons" style="margin-top: 20px;">
+                        <button class="btn-primary" id="btn-proceed-to-graph">Proceed to Attribution Graph</button>
+                    </div>
+                </div>
+            </div>
+
+            <!-- Step 4: Attribution Graph -->
+            <div class="step-container" id="step-4">
+                <div class="step-header disabled" onclick="toggleStep(4)">
+                    <div style="display: flex; align-items: center;">
+                        <span class="step-number">4</span>
+                        <span class="step-title">Attribution Graph Visualization</span>
+                    </div>
+                    <span class="step-status" id="status-4">Locked</span>
+                </div>
+                <div class="step-body collapsed" id="body-4">
+                    
+                    <div class="subsection">
+                        <h4>🔧 Graph Configuration</h4>
+                        
+                        <div class="control-group">
+                            <label><strong>Layer Selection:</strong></label>
+                            <input type="text" id="layers-list" placeholder="e.g., 0, 10, 20, 27" style="width: 300px;">
+                            <button class="btn-secondary" id="btn-layer-preset-default">Recommended</button>
+                            <button class="btn-secondary" id="btn-layer-preset-all">All Layers</button>
+                        </div>
+
+                        <div class="control-group" style="margin-top: 15px;">
+                            <label>Backprop Mode:
+                                <select id="graph-bp-mode">
+                                    <option value="max_logit">Max Logit</option>
+                                    <option value="logit_diff" selected>Logit Difference</option>
+                                </select>
+                            </label>
+                            <label style="margin-left: 15px;">Strategy:
+                                <select id="graph-bp-strategy">
+                                    <option value="by_topk_avg" selected>By Top-K Avg</option>
+                                    <option value="demean">Demean</option>
+                                    <option value="by_ref_token">By Reference Token</option>
+                                </select>
+                            </label>
+                        </div>
+
+                        <div class="control-group" style="margin-top: 15px;">
+                            <label>Pruning Mode:
+                                <select id="pruning-mode">
+                                    <option value="by_per_layer_cum_mass_percentile" selected>By Top-P Mass (Layer)</option>
+                                    <option value="by_global_threshold">By Global Threshold</option>
+                                </select>
+                            </label>
+                        </div>
+                        
+                        <div id="ctrl-top-p" class="control-group" style="margin-top: 10px;">
+                            <label>Top-P Mass (0-1):
+                                <input type="range" id="top-p-slider" min="0" max="1" step="0.01" value="0.85" style="width: 200px; margin-left: 10px;">
+                                <input type="number" id="top-p-value" value="0.85" min="0" max="1" step="0.01" style="width: 70px; margin-left: 10px;">
+                            </label>
+                        </div>
+                        
+                        <div id="ctrl-global-thresh" class="control-group hidden" style="margin-top: 10px;">
+                            <label>Global Edge Threshold:
+                                <input type="number" id="global-thresh-value" value="0.01" min="0" max="1" step="0.001" style="width: 100px; margin-left: 10px;">
+                            </label>
+                        </div>
+
+                        <div class="action-buttons">
+                            <button class="btn-primary" id="btn-compute-graph">Compute & Visualize Graph</button>
+                        </div>
+                    </div>
+
+                    <div id="graph-results" class="hidden" style="margin-top: 20px;">
+                        <h4>📊 Circuit Visualization</h4>
+                        
+                        <!-- Visualization Controls -->
+                        <div style="display: flex; gap: 15px; margin-bottom: 15px; flex-wrap: wrap; align-items: center;">
+                            <label><input type="checkbox" id="show-all-tokens"> Show All Tokens</label>
+                            <label><input type="checkbox" id="hide-bos"> Hide BOS Token</label>
+                            <label><input type="checkbox" id="show-values" checked> Show Node Values</label>
+                        </div>
+                        
+                        <div style="display: flex; gap: 20px; margin-bottom: 15px; flex-wrap: wrap; align-items: center;">
+                            <label style="display: flex; align-items: center; gap: 8px;">
+                                Line Strength: 
+                                <input type="range" id="line-strength" min="2" max="20" value="6" step="0.1" style="width: 120px;">
+                                <span id="line-strength-val" style="min-width: 35px; font-weight: bold;">6.0</span>
+                            </label>
+                            <label style="display: flex; align-items: center; gap: 8px;">
+                                Layer Spacing (px): 
+                                <input type="range" id="layer-spacing-slider" min="50" max="400" value="100" step="10" style="width: 120px;">
+                                <input type="number" id="layer-spacing-value" min="50" max="400" value="100" step="10" style="width: 60px;">
+                            </label>
+                        </div>
+                        
+                        <!-- Threshold Adjustment (Post-Compute) -->
+                        <div style="padding: 10px; background: #f0f0f0; border-radius: 5px; margin-bottom: 15px;">
+                            <strong>🎛️ Adjust Graph Sparsity:</strong>
+                            <div id="post-ctrl-top-p" style="margin-top: 8px;">
+                                <label>Top-P Mass: 
+                                    <input type="range" id="post-top-p-slider" min="0" max="1" step="0.01" value="0.85" style="width: 200px; margin-left: 10px;">
+                                    <input type="number" id="post-top-p-value" value="0.85" min="0" max="1" step="0.01" style="width: 70px; margin-left: 10px;">
+                                    <button class="btn-secondary" id="btn-recompute-threshold" style="margin-left: 10px;">Recompute</button>
+                                </label>
+                            </div>
+                            <div id="post-ctrl-global-thresh" class="hidden" style="margin-top: 8px;">
+                                <label>Global Threshold: 
+                                    <input type="number" id="post-global-thresh-value" value="0.01" min="0" max="1" step="0.001" style="width: 100px; margin-left: 10px;">
+                                    <button class="btn-secondary" id="btn-recompute-threshold-global" style="margin-left: 10px;">Recompute</button>
+                                </label>
+                            </div>
+                        </div>
+
+                        <div id="canvas-container" style="border: 1px solid #ddd; border-radius: 5px; overflow: hidden;">
+                            <canvas id="graph-canvas"></canvas>
+                            <div id="graph-tooltip"></div>
+                        </div>
+
+                        <div class="action-buttons" style="margin-top: 15px;">
+                            <button class="btn-secondary" id="btn-save-graph">Save as PNG</button>
+                            <button class="btn-secondary" id="btn-save-graph-pdf">Save as PDF</button>
+                        </div>
+                    </div>
+
+                    <div style="margin-top: 30px; padding: 20px; background: #d4edda; border-radius: 8px;">
+                        <h3 style="margin-top: 0; color: #155724;">✅ Analysis Complete!</h3>
+                        <p>You have completed the full step-by-step analysis pipeline. Review the results above or export your findings.</p>
+                    </div>
+                </div>
+            </div>
+
+        </main>
+    </div>
+
+    <!-- Loading Overlay -->
+    <div id="loading-overlay" class="hidden">
+        <div class="spinner"></div>
+        <div id="loading-message">Processing...</div>
+    </div>
+
+    <script src="js/main_new.js?v=20260330"></script>
+</body>
+</html>

frontend/index_old.html

@@ -0,0 +1,384 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>LLM Insider for Failure Debugging</title>
+    <link rel="stylesheet" href="css/style.css">
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/jspdf/2.5.1/jspdf.umd.min.js"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.min.js"></script>
+</head>
+<body>
+    <div id="app">
+        <header>
+            <h1>LLM Insider for Failure Debugging</h1>
+        </header>
+
+        <main>
+            <!-- Section 1: Build Circuit (Model & Inputs) -->
+            <section id="section-setup" class="panel">
+                <h2>1. Setup & Inputs</h2>
+                
+                <!-- Trace Loading (Top) -->
+                <div class="control-group" style="background: #f8f9fa; padding: 10px; border-radius: 4px; border: 1px dashed #ced4da;">
+                    <label><strong>Load from Trace:</strong></label>
+                    <label>Dataset: 
+                        <select id="trace-dataset">
+                            <option value="">-- Select --</option>
+                            <!-- Populated by JS -->
+                        </select>
+                    </label>
+                    <label>Trace: 
+                        <select id="trace-file" disabled>
+                            <option value="">-- Select Dataset First --</option>
+                        </select>
+                    </label>
+                </div>
+
+                <div class="control-group">
+                    <div style="display: flex; flex-direction: column; gap: 5px;">
+                        <div style="display: flex; gap: 10px; align-items: center;">
+                            <label>Series: 
+                                <select id="model-series">
+                                    <option value="Qwen3">Qwen3</option>
+                                    <option value="Qwen2">Qwen2</option>
+                                    <option value="OLMo3">OLMo3</option>
+                                </select>
+                            </label>
+                            <label>Model: 
+                                <select id="model-select">
+                                    <!-- Populated dynamically -->
+                                </select>
+                            </label>
+                        </div>
+                        <div style="display: flex; gap: 10px; align-items: center;">
+                            <label>Revision: 
+                                <select id="model-revision" disabled>
+                                    <option value="">Latest (Default)</option>
+                                </select>
+                            </label>
+                            <label>Path:
+                                <input type="text" id="model-path" value="Qwen/Qwen3-0.6B" style="width: 250px;">
+                            </label>
+                        </div>
+                    </div>
+
+                    <label>DType:
+                        <select id="model-dtype">
+                            <option value="float16">float16</option>
+                            <option value="bfloat16" selected>bfloat16</option>
+                            <option value="float32">float32</option>
+                            <option value="auto">auto</option>
+                        </select>
+                    </label>
+                    <label>LRP Rule:
+                        <select id="model-lrp-rule">
+                            <option value="Attn-LRP" selected>Attn-LRP</option>
+                            <option value="CP-LRP">CP-LRP</option>
+                        </select>
+                    </label>
+                    <label>
+                        <input type="checkbox" id="quant-4bit"> 4-bit Quantization
+                    </label>
+                    <button id="btn-load-model">Load Model</button>
+                    <span id="model-status" class="status">Not Loaded</span>
+                </div>
+
+                <div class="control-group">
+                    <label for="prompt-input">Prompt:</label>
+                    <textarea id="prompt-input" rows="3">After John and Mary went to the store, John gave a bottle of milk to</textarea>
+                </div>
+
+                <div class="control-group">
+                    <label for="completion-input">Completion:</label>
+                    <textarea id="completion-input" rows="3" placeholder="Enter the completion text here..."></textarea>
+                </div>
+
+                <div class="control-group" style="display: flex; gap: 10px; align-items: center;">
+                    <button id="btn-locate-error" disabled style="background-color: #e74c3c;">Locate Error Token</button>
+                    <span id="error-locate-status" style="color: #666;"></span>
+                </div>
+            </section>
+
+            <!-- Section: Error Token Location Result -->
+            <section id="section-error-result" class="panel hidden">
+                <h2>Error Token Location Result</h2>
+                <div class="control-group">
+                    <label for="error-explanation"><strong>Error Analysis:</strong></label>
+                    <div id="error-explanation" style="padding: 10px; background: #fff3cd; border: 1px solid #ffc107; border-radius: 4px; margin-bottom: 10px; white-space: pre-wrap;"></div>
+                </div>
+
+                <div class="control-group">
+                    <label for="truncated-text"><strong>Truncated Text (Editable):</strong></label>
+                    <textarea id="truncated-text" rows="5" style="width: 100%; font-family: monospace;"></textarea>
+                </div>
+
+                <div class="control-group" style="display: flex; gap: 10px;">
+                    <button id="btn-use-truncated" style="background-color: #28a745;">Use This Text</button>
+                    <button id="btn-cancel-error" style="background-color: #6c757d;">Cancel</button>
+                </div>
+            </section>
+
+            <!-- Section for continuing analysis -->
+            <section id="section-continue-analysis" class="panel hidden">
+                <h2>Continue Analysis</h2>
+                <div class="control-group">
+                    <label for="prompt-orig-completion">Current Text for Analysis:</label>
+                    <textarea id="prompt-orig-completion" rows="3" readonly style="background-color: #f8f9fa;"></textarea>
+                </div>
+                
+                <div class="control-group">
+                    <label style="margin-right: 15px;">
+                        <input type="checkbox" id="append-bos"> Append BOS Token
+                    </label>
+                    <label style="margin-right: 15px;">
+                        <input type="checkbox" id="capture-mid"> Attn/MLP Split
+                    </label>
+                    <button id="btn-compute-logits" disabled>Compute Logits</button>
+                    <span id="token-count-display" style="margin-left: 15px; color: #666;"></span>
+                </div>
+            </section>
+
+            <!-- Section: Trace Exploration Data (Hidden by default, shown when trace loaded) -->
+            <section id="section-trace-explore" class="panel hidden">
+                <details>
+                    <summary style="cursor: pointer; font-weight: bold; font-size: 1.2em; margin-bottom: 10px;">Trace Candidate Exploration (Correctness Check)</summary>
+                    <div style="display: flex; gap: 20px; flex-wrap: wrap; margin-top: 10px;">
+                        <div style="flex: 1; min-width: 400px;">
+                            <h3 style="margin-top: 0;">Current Model Candidates</h3>
+                            <div id="container-explore-original" style="max-height: 400px; overflow-y: auto; border: 1px solid #ccc; background: white;">
+                                <!-- Table will be injected here -->
+                            </div>
+                        </div>
+                        <div style="flex: 1; min-width: 400px;">
+                            <h3 id="header-explore-other" style="margin-top: 0;">Other Model Candidates</h3>
+                            <div id="container-explore-4b" style="max-height: 400px; overflow-y: auto; border: 1px solid #ccc; background: white;">
+                                <!-- Table will be injected here -->
+                            </div>
+                        </div>
+                    </div>
+                </details>
+            </section>
+
+            <!-- Section 1 Output: Logits -->
+            <section id="section-logits" class="panel hidden">
+                <div style="display: flex; justify-content: space-between; align-items: center; flex-wrap: wrap; gap: 15px;">
+                    <h2 style="margin: 0;">Logit Analysis (Top-K)</h2>
+                    <div style="display: flex; gap: 10px; align-items: center;">
+                        <input type="text" id="add-token-input" placeholder="Token Str or ID" style="padding: 4px; border: 1px solid #ccc; border-radius: 4px; width: 120px;">
+                        <button id="btn-add-token" style="padding: 4px 8px; font-size: 0.9em;">Add Token</button>
+                        <span style="border-left: 1px solid #ccc; height: 20px; margin: 0 5px;"></span>
+                        <label style="font-size: 0.9rem;">Gen Length: <input type="number" id="gen-max-tokens" value="30" min="1" max="200" style="width: 50px; padding: 4px; border: 1px solid #ccc; border-radius: 4px;"></label>
+                    </div>
+                </div>
+                <!-- Table -->
+                <div class="scrollable-table" style="margin-top: 10px;">
+                    <table id="logits-table">
+                        <thead>
+                            <tr>
+                                <th>Rank</th>
+                                <th>Token ID</th>
+                                <th>Token</th>
+                                <th>Logit</th>
+                                <th>Action</th>
+                                <th>Continual Generation</th>
+                            </tr>
+                        </thead>
+                        <tbody>
+                            <!-- Populated by JS -->
+                        </tbody>
+                    </table>
+                </div>
+                <!-- Selection State for Circuit -->
+                <div id="circuit-selection-info">
+                    <p><strong>Top Predicted:</strong> <span id="top-token-display">-</span></p>
+                    <p><strong>Contrast With:</strong> <span id="contrast-token-display">None</span></p>
+                </div>
+            </section>
+
+            <!-- Section 1.5: Input Attribution -->
+            <section id="section-input-attr" class="panel hidden">
+                <h2>2. Input Attribution Visualization</h2>
+                <div class="controls-row">
+                    <div class="control-col">
+                         <h3>Attribution Config</h3>
+                         <label>Mode:
+                            <select id="attr-bp-mode">
+                                <option value="max_logit">Max Logit</option>
+                                <option value="logit_diff" selected>Logit Difference</option>
+                            </select>
+                         </label>
+                         <div id="attr-diff-strategy-opts">
+                            <label>Strategy:
+                                <select id="attr-bp-strategy">
+                                    <option value="by_topk_avg">By Top-K Avg</option>
+                                    <option value="demean">Demean (Avg All)</option>
+                                    <option value="by_ref_token" selected>By Reference Token</option>
+                                </select>
+                            </label>
+                            <label id="attr-opt-ref">Ref ID: <input type="number" id="attr-bp-ref-id" placeholder="ID"></label>
+                         </div>
+                         <label style="margin-top: 10px; display: block; font-size: 0.9em;">
+                            <input type="checkbox" id="attr-hide-bos" checked> Hide First Token
+                         </label>
+                         <label style="margin-top: 5px; display: block; font-size: 0.9em;">
+                            <input type="checkbox" id="attr-hide-special" checked> Hide Special Tokens (&lt;|im_start|&gt;, &lt;think&gt;, etc.)
+                         </label>
+                         <label style="margin-top: 5px; display: block; font-size: 0.9em;">
+                            <input type="checkbox" id="attr-reverse-sign"> Reverse Sign (Red &harr; Blue)
+                         </label>
+                         <button id="btn-compute-attr">Compute Attribution</button>
+                    </div>
+                    <div class="control-col" style="flex: 2;">
+                        <h3>Attribution Map</h3>
+                        <div id="input-attribution-display" class="attribution-text-box">
+                            <!-- Tokens will be rendered here -->
+                             Click "Compute Attribution" to see results.
+                        </div>
+                        <div style="margin-top: 10px; display: flex; gap: 10px;">
+                            <button id="btn-save-attr-png" style="background-color: #2ecc71; font-size: 0.9em; padding: 4px 10px;">Save PNG</button>
+                            <button id="btn-save-attr-pdf" style="background-color: #e74c3c; font-size: 0.9em; padding: 4px 10px;">Save PDF</button>
+                        </div>
+                    </div>
+                </div>
+            </section>
+
+            <!-- Section 2: Circuit Visualization -->
+            <section id="section-circuit" class="panel hidden">
+                <h2>3. Circuit Visualization</h2>
+                
+                <!-- PROGRESS BAR CONTAINER -->
+                <div id="compute-progress-container" style="display: none; margin-bottom: 1rem; border: 1px solid #ddd; padding: 10px; border-radius: 4px; background: #f0f8ff;">
+                     <div style="display: flex; justify-content: space-between; margin-bottom: 5px;">
+                        <strong id="progress-status-text">Computation Status...</strong>
+                        <span id="progress-percent-text">0%</span>
+                     </div>
+                     <progress id="compute-progress-bar" value="0" max="100" style="width: 100%; height: 20px;"></progress>
+                </div>
+                
+                <div class="controls-row">
+                    <!-- Column 1: Backprop Config -->
+                    <div class="control-col">
+                        <h3>Backprop Strategy</h3>
+                        <label>Mode:
+                            <select id="bp-mode">
+                                <option value="max_logit">Max Logit (Target Only)</option>
+                                <option value="logit_diff" selected>Logit Difference</option>
+                            </select>
+                        </label>
+                        
+                        <div id="diff-strategy-opts">
+                            <label>Strategy:
+                                <select id="bp-strategy">
+                                    <option value="by_topk_avg" selected>By Top-K Avg</option>
+                                    <option value="demean">Demean (Avg All)</option>
+                                    <option value="by_ref_token">By Reference Token</option>
+                                </select>
+                            </label>
+                            
+                            <div id="opt-topk" class="sub-opt">
+                                <label>K: <input type="number" id="bp-k" value="10" style="width: 50px;"></label>
+                            </div>
+                            
+                            <div id="opt-ref" class="sub-opt hidden">
+                                <label>Ref Token ID: <input type="number" id="bp-ref-id" placeholder="ID"></label>
+                            </div>
+                        </div>
+                    </div>
+
+                    <!-- Column 2: Layer Config -->
+                    <div class="control-col">
+                        <h3>Layer Connection</h3>
+                        <label>Layers (Comma-sep, increasing): 
+                            <input type="text" id="layers-list" placeholder="e.g. 0, 10, 27" style="width: 100%;">
+                        </label>
+                        <div style="margin-top: 5px;">
+                            <small>Presets:</small>
+                            <button id="btn-layer-default" class="btn-micro">Recommended</button>
+                            <button id="btn-layer-all" class="btn-micro">All Layers</button>
+                        </div>
+                        <button id="btn-compute-circuit" style="margin-top: 10px;">Visualize Connection</button>
+                    </div>
+
+                    <!-- Column 3: Vis Controls -->
+                    <div class="control-col">
+                        <h3>Vis Settings</h3>
+                        <div class="setting-item">
+                            <label><input type="checkbox" id="show-all-tokens"> Show All Tokens</label>
+                        </div>
+                        <div class="setting-item">
+                            <label><input type="checkbox" id="hide-bos-node"> Hide First Token (BOS)</label>
+                        </div>
+                        <div class="setting-item">
+                            <label><input type="checkbox" id="show-node-values" checked> Show Node Values</label>
+                        </div>
+                        
+                        <!-- New Pruning Controls -->
+                        <div class="setting-item">
+                            <label>Pruning Mode:</label>
+                            <select id="pruning-mode" style="width: 100%;">
+                                <option value="by_per_layer_cum_mass_percentile" selected>By Top-P Mass (Layer)</option>
+                                <option value="by_global_threshold">By Global Threshold</option>
+                            </select>
+                        </div>
+                        
+                        <div id="ctrl-top-p" class="setting-item">
+                            <label>Top-P Mass (0-1):</label>
+                            <div class="input-row">
+                                <input type="range" id="vis-top-p" min="0" max="1" step="0.01" value="0.85">
+                                <input type="number" id="val-vis-top-p" min="0" max="1" step="0.01" value="0.85" style="width: 60px;">
+                            </div>
+                        </div>
+
+                        <div id="ctrl-global-thresh" class="setting-item hidden">
+                            <label>Global Threshold:</label>
+                            <div class="input-row">
+                                <input type="number" id="val-global-thresh" value="0.01" step="0.001" style="width: 100%;">
+                            </div>
+                        </div>
+
+                        <div class="setting-item">
+                            <label>Line Strength:</label>
+                            <div class="input-row">
+                                <input type="range" id="vis-strength" min="2" max="20" step="0.1" value="6">
+                                <input type="number" id="val-vis-strength" min="2" max="20" step="0.1" value="6" style="width: 60px;">
+                            </div>
+                        </div>
+                        <div class="setting-item">
+                            <label>Layer Spacing (px):</label>
+                            <div class="input-row">
+                                <input type="range" id="vis-layer-spacing" min="10" max="200" step="10" value="50">
+                                <input type="number" id="val-vis-layer-spacing" min="10" max="200" step="10" value="50" style="width: 60px;">
+                            </div>
+                        </div>
+
+                        <div class="setting-item" style="margin-top: 15px;">
+                            <button id="btn-save-graph" style="width: 48%; background-color: #2ecc71;">Save PNG</button>
+                            <button id="btn-save-pdf" style="width: 48%; background-color: #e74c3c;">Save PDF</button>
+                        </div>
+                    </div>
+                </div>
+
+                <div id="canvas-container">
+                    <canvas id="circuit-canvas"></canvas>
+                    <div id="tooltip"></div>
+                </div>
+                
+                <div id="vis-legend">
+                     <span class="legend-box pos"></span> Pos Relevance
+                     <span class="legend-box neg"></span> Neg Relevance
+                </div>
+            </section>
+        </main>
+    </div>
+
+    <script src="https://cdn.jsdelivr.net/npm/html2canvas@1.4.1/dist/html2canvas.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/jspdf@2.5.1/dist/jspdf.umd.min.js"></script>
+
+    <script src="js/main.js"></script>
+    <div id="loading-overlay" class="hidden">
+        <div class="spinner"></div>
+        <div id="loading-message">Processing...</div>
+    </div>
+</body>
+</html>

frontend/js/main.js

@@ -0,0 +1,2231 @@
+const API_BASE = "/api";
+
+// Loading Spinner Helpers
+const elLoadingOverlay = document.getElementById('loading-overlay');
+const elLoadingMessage = document.getElementById('loading-message');
+
+function showLoading(msg="Processing...") {
+    if(elLoadingMessage) elLoadingMessage.textContent = msg;
+    if(elLoadingOverlay) {
+        elLoadingOverlay.classList.remove('hidden');
+        // Force browser repaint
+        void elLoadingOverlay.offsetWidth;
+    }
+}
+
+function hideLoading() {
+    if(elLoadingOverlay) elLoadingOverlay.classList.add('hidden');
+}
+
+// State
+let appState = {
+    tokens: [],
+    selectedTargetToken: null,
+    selectedContrastToken: null,
+    circuitData: null,
+    mouseX: -1,
+    mouseY: -1,
+    extraTokens: [], // Stores objects: { type: 'id'|'str', value: ... }
+    lastRelevance: null // Cache for attribution map re-rendering
+};
+
+// Selection State for Graph
+let selectedNode = null;
+let lastRenderedNodes = [];
+
+// DOM Elements
+const elModelSeries = document.getElementById('model-series');
+const elModelSelect = document.getElementById('model-select');
+const elModelRevision = document.getElementById('model-revision');
+const elModelPath = document.getElementById('model-path');
+const elModelDtype = document.getElementById('model-dtype');
+const elQuant = document.getElementById('quant-4bit');
+const elBtnLoad = document.getElementById('btn-load-model');
+const elStatus = document.getElementById('model-status');
+
+const elTraceDataset = document.getElementById('trace-dataset');
+const elTraceFile = document.getElementById('trace-file');
+const elSectionTraceExplore = document.getElementById('section-trace-explore');
+const elContainerExploreOriginal = document.getElementById('container-explore-original');
+const elContainerExplore4b = document.getElementById('container-explore-4b');
+const elHeaderExploreOther = document.getElementById('header-explore-other');
+
+const elPrompt = document.getElementById('prompt-input');
+const elCompletion = document.getElementById('completion-input');
+const elPromptOrig = document.getElementById('prompt-orig-completion');
+const elAppendBos = document.getElementById('append-bos');
+const elBtnLogits = document.getElementById('btn-compute-logits');
+const elBtnLocateError = document.getElementById('btn-locate-error');
+const elErrorLocateStatus = document.getElementById('error-locate-status');
+
+const elSectionErrorResult = document.getElementById('section-error-result');
+const elSectionContinueAnalysis = document.getElementById('section-continue-analysis');
+const elErrorExplanation = document.getElementById('error-explanation');
+const elTruncatedText = document.getElementById('truncated-text');
+const elBtnUseTruncated = document.getElementById('btn-use-truncated');
+const elBtnCancelError = document.getElementById('btn-cancel-error');
+
+const elSectionLogits = document.getElementById('section-logits');
+const elSectionInputAttr = document.getElementById('section-input-attr');
+const elLogitsTableBody = document.querySelector('#logits-table tbody');
+const elTopTokenDisplay = document.getElementById('top-token-display');
+const elContrastTokenDisplay = document.getElementById('contrast-token-display');
+
+const elSectionCircuit = document.getElementById('section-circuit');
+const elBpMode = document.getElementById('bp-mode');
+const elBpStrategy = document.getElementById('bp-strategy');
+const elBpK = document.getElementById('bp-k');
+const elBpRefId = document.getElementById('bp-ref-id');
+const elOptTopk = document.getElementById('opt-topk');
+const elOptRef = document.getElementById('opt-ref');
+
+const elTargetLayer = document.getElementById('target-layer');
+const elSourceLayer = document.getElementById('source-layer');
+const elLayersList = document.getElementById('layers-list'); // NEW INPUT
+const elBtnCircuit = document.getElementById('btn-compute-circuit');
+
+const elVisStrength = document.getElementById('vis-strength');
+const elValVisStrength = document.getElementById('val-vis-strength');
+
+// New Controls
+const elPruningMode = document.getElementById('pruning-mode');
+const elVisTopP = document.getElementById('vis-top-p');
+const elValVisTopP = document.getElementById('val-vis-top-p');
+const elValGlobalThresh = document.getElementById('val-global-thresh');
+const elCtrlTopP = document.getElementById('ctrl-top-p');
+const elCtrlGlobalThresh = document.getElementById('ctrl-global-thresh');
+
+const elShowAllTokens = document.getElementById('show-all-tokens');
+const elBtnLayerDefault = document.getElementById('btn-layer-default');
+const elBtnLayerAll = document.getElementById('btn-layer-all');
+
+// Input Attribution Elements
+const elAttrBpMode = document.getElementById('attr-bp-mode');
+const elAttrBpStrategy = document.getElementById('attr-bp-strategy');
+const elAttrBpRefId = document.getElementById('attr-bp-ref-id');
+const elAttrOptRef = document.getElementById('attr-opt-ref');
+const elAttrHideBos = document.getElementById('attr-hide-bos');
+const elAttrHideSpecial = document.getElementById('attr-hide-special');
+const elAttrReverseSign = document.getElementById('attr-reverse-sign');
+const elBtnComputeAttr = document.getElementById('btn-compute-attr');
+const elBtnSaveGraph = document.getElementById('btn-save-graph'); // Save Button
+const elInputAttributionDisplay = document.getElementById('input-attribution-display');
+const elAttrDiffStrategyOpts = document.getElementById('attr-diff-strategy-opts'); // Container for strategy ops
+
+const elTokenCountDisplay = document.getElementById('token-count-display');
+const elVisLayerSpacing = document.getElementById('vis-layer-spacing');
+const elValVisLayerSpacing = document.getElementById('val-vis-layer-spacing');
+
+const elCanvas = document.getElementById('circuit-canvas');
+const elTooltip = document.getElementById('tooltip');
+
+// Event Listeners
+console.log("Attaching event listeners...");
+try {
+    if(elBtnLoad) elBtnLoad.addEventListener('click', loadModel);
+    else console.error("elBtnLoad not found");
+    
+    if(elBtnLocateError) elBtnLocateError.addEventListener('click', locateErrorToken);
+    else console.warn("elBtnLocateError not found");
+    
+    if(elBtnUseTruncated) elBtnUseTruncated.addEventListener('click', useTruncatedText);
+    else console.warn("elBtnUseTruncated not found");
+    
+    if(elBtnCancelError) elBtnCancelError.addEventListener('click', cancelErrorLocation);
+    else console.warn("elBtnCancelError not found");
+    
+    if(elBtnLogits) elBtnLogits.addEventListener('click', () => computeLogits(false));
+
+    // Add Token Feature
+    const elBtnAddToken = document.getElementById('btn-add-token');
+    const elAddTokenInput = document.getElementById('add-token-input');
+    
+    if(elBtnAddToken && elAddTokenInput) {
+        elBtnAddToken.addEventListener('click', () => {
+            const val = elAddTokenInput.value.trim();
+            if(!val) return;
+            
+            // Parse: is it ID or String?
+            // If numeric, treat as ID.
+            if (/^\d+$/.test(val)) {
+                const id = parseInt(val);
+                appState.extraTokens.push({ type: 'id', value: id });
+            } else {
+                appState.extraTokens.push({ type: 'str', value: val });
+            }
+            
+            elAddTokenInput.value = ""; // Clear input
+            computeLogits(true); // Keep extras!
+        });
+        
+        // Allow Enter key
+        elAddTokenInput.addEventListener('keypress', (e) => {
+            if (e.key === 'Enter') elBtnAddToken.click();
+        });
+    }
+    if(elBpMode) elBpMode.addEventListener('change', updateBpOptions);
+    if(elBpStrategy) elBpStrategy.addEventListener('change', updateBpOptions);
+    if(elBtnCircuit) elBtnCircuit.addEventListener('click', computeCircuit);
+    if(elShowAllTokens) elShowAllTokens.addEventListener('change', drawCircuit);
+    
+    if(elBtnLayerDefault) elBtnLayerDefault.addEventListener('click', () => setLayerPreset('default'));
+    if(elBtnLayerAll) elBtnLayerAll.addEventListener('click', () => setLayerPreset('all'));
+
+    // Input Attribution Listeners
+    if(elAttrBpMode) elAttrBpMode.addEventListener('change', updateAttrBpOptions);
+    if(elAttrBpStrategy) elAttrBpStrategy.addEventListener('change', updateAttrBpOptions);
+    // Re-render when checkboxes change (if data exists)
+    const reRenderAttr = () => { if(appState.lastRelevance) renderAttributionMap(appState.lastRelevance, appState.tokens); };
+    if(elAttrHideBos) elAttrHideBos.addEventListener('change', reRenderAttr);
+    if(elAttrHideSpecial) elAttrHideSpecial.addEventListener('change', reRenderAttr);
+    if(elAttrReverseSign) elAttrReverseSign.addEventListener('change', reRenderAttr);
+
+    if(elBtnComputeAttr) elBtnComputeAttr.addEventListener('click', computeInputAttribution);
+    if(elBtnSaveGraph) elBtnSaveGraph.addEventListener('click', saveCircuitGraph);
+    const elBtnSavePdf = document.getElementById('btn-save-pdf');
+    if(elBtnSavePdf) elBtnSavePdf.addEventListener('click', saveCircuitGraphPDF);
+    
+    // Attrib Save
+    const elBtnSaveAttrPng = document.getElementById('btn-save-attr-png');
+    const elBtnSaveAttrPdf = document.getElementById('btn-save-attr-pdf');
+    if(elBtnSaveAttrPng) elBtnSaveAttrPng.addEventListener('click', saveAttributionMapPNG);
+    if(elBtnSaveAttrPdf) elBtnSaveAttrPdf.addEventListener('click', saveAttributionMapPDF);
+    
+    
+    if(elPruningMode) elPruningMode.addEventListener('change', updatePruningControls);
+
+    console.log("Event listeners attached.");
+} catch(e) {
+    console.error("Error attaching listeners:", e);
+}
+
+// Visualization Controls (Slider <-> Input Sync)
+function bindControl(slider, input, callback, debounceMs = 0) {
+    if (!slider || !input) {
+        // console.error("bindControl missing elements", slider, input);
+        return;
+    }
+    
+    let debounceTimer = null;
+    const trigger = () => {
+        if (debounceMs > 0) {
+            if (debounceTimer) clearTimeout(debounceTimer);
+            debounceTimer = setTimeout(() => {
+                if(callback) callback();
+            }, debounceMs);
+        } else {
+            if(callback) callback();
+        }
+    };
+
+    // Slider updates Input
+    slider.addEventListener('input', () => {
+        input.value = slider.value;
+        trigger();
+    });
+    // Input updates Slider
+    input.addEventListener('input', () => {
+        let val = parseFloat(input.value);
+        if (!isNaN(val)) {
+            // Clamp to slider limits
+            val = Math.max(parseFloat(slider.min), Math.min(parseFloat(slider.max), val));
+            slider.value = val;
+            trigger();
+        }
+    });
+}
+console.log("Binding controls...");
+// Bind with specific update functions for efficiency
+// Layout updates (heavy): Edge/Node Thresholds, Spacing - Debounced
+// bindControl(elVisThreshold, elValVisThreshold, updateCircuitLayout, 300); 
+// bindControl(elVisNodeThreshold, elValVisNodeThreshold, updateCircuitLayout, 300);
+
+// NEW Bindings for TopP (Callback: computeCircuit, Debounced)
+bindControl(elVisTopP, elValVisTopP, computeCircuit, 600);
+
+// Layout Geometry updates (moved to light): Spacing
+bindControl(elVisLayerSpacing, elValVisLayerSpacing, drawCircuit, 0);
+// Visual updates (light): Strength
+bindControl(elVisStrength, elValVisStrength, drawCircuit, 0);
+
+function updatePruningControls() {
+    const mode = elPruningMode.value;
+    if (mode === 'by_per_layer_cum_mass_percentile') {
+        elCtrlTopP.classList.remove('hidden');
+        elCtrlGlobalThresh.classList.add('hidden');
+    } else {
+        elCtrlTopP.classList.add('hidden');
+        elCtrlGlobalThresh.classList.remove('hidden');
+    }
+}
+// Init State
+updatePruningControls();
+
+
+
+// Functions
+
+async function loadModel() {
+    console.log("loadModel called");
+    updateStatus('Loading...', 'normal');
+    showLoading("Loading Model...");
+    elBtnLoad.disabled = true;
+    
+    try {
+        const response = await fetch(`${API_BASE}/load_model`, {
+            method: 'POST',
+            headers: {'Content-Type': 'application/json'},
+            body: JSON.stringify({
+                model_path: elModelPath.value,
+                quantization_4bit: elQuant.checked,
+                dtype: elModelDtype.value,
+                revision: elModelRevision ? elModelRevision.value : null,
+                lrp_rule: document.getElementById('model-lrp-rule').value
+            })
+        });
+        
+        const data = await response.json();
+        if (response.ok) {
+            updateStatus('Loaded', 'success');
+            if(elBtnLogits) elBtnLogits.disabled = false;
+            if(elBtnLocateError) elBtnLocateError.disabled = false;
+            if(data.n_layers) {
+                appState.n_layers = data.n_layers;
+                // Update default presets
+                generateLayerPresets();
+            }
+        } else {
+            throw new Error(data.detail || 'Failed to load model');
+        }
+    } catch (e) {
+        updateStatus(`Error: ${e.message}`, 'error');
+        alert(e.message);
+    } finally {
+        elBtnLoad.disabled = false;
+        hideLoading();
+    }
+}
+
+async function locateErrorToken() {
+    const prompt = elPrompt.value.trim();
+    const completion = elCompletion.value.trim();
+    
+    if (!prompt) {
+        alert("Please enter a prompt.");
+        return;
+    }
+    
+    if (!completion) {
+        alert("Please enter a completion.");
+        return;
+    }
+    
+    elBtnLocateError.disabled = true;
+    elErrorLocateStatus.textContent = "Locating error token...";
+    showLoading("Analyzing error token...");
+    
+    try {
+        const res = await fetch(`${API_BASE}/locate_err_token`, {
+            method: 'POST',
+            headers: { 'Content-Type': 'application/json' },
+            body: JSON.stringify({
+                prompt: prompt,
+                completion: completion
+            })
+        });
+        
+        const data = await res.json();
+        
+        if (res.ok && data.status === "success") {
+            elErrorLocateStatus.textContent = "Error token located successfully.";
+            
+            // Display results
+            elErrorExplanation.textContent = data.explanation || "No explanation provided.";
+            elTruncatedText.value = data.truncated_text || "";
+            
+            // Show result section
+            elSectionErrorResult.classList.remove('hidden');
+            
+        } else {
+            elErrorLocateStatus.textContent = "Failed to locate error token.";
+            alert(`Error: ${data.message || data.detail || 'Unknown error'}`);
+        }
+    } catch (e) {
+        elErrorLocateStatus.textContent = "Request failed.";
+        alert(`Network error: ${e.message}`);
+    } finally {
+        hideLoading();
+        elBtnLocateError.disabled = false;
+    }
+}
+
+function useTruncatedText() {
+    const truncatedText = elTruncatedText.value.trim();
+    
+    if (!truncatedText) {
+        alert("Truncated text is empty.");
+        return;
+    }
+    
+    // Update the prompt field with truncated text
+    elPrompt.value = truncatedText;
+    elPromptOrig.value = truncatedText;
+    
+    // Hide error result section and show continue analysis section
+    elSectionErrorResult.classList.add('hidden');
+    elSectionContinueAnalysis.classList.remove('hidden');
+    
+    // Clear error status
+    elErrorLocateStatus.textContent = "";
+}
+
+function cancelErrorLocation() {
+    // Hide error result section
+    elSectionErrorResult.classList.add('hidden');
+    
+    // Clear error status
+    elErrorLocateStatus.textContent = "";
+}
+
+async function computeLogits(keepExtras = false) {
+    elBtnLogits.disabled = true;
+    showLoading("Computing Logits...");
+    
+    // If this is a fresh run (not adding a token), clear extras
+    if (!keepExtras) {
+        appState.extraTokens = [];
+    }
+
+    elSectionLogits.classList.add('hidden');
+    if(elSectionInputAttr) elSectionInputAttr.classList.add('hidden');
+    elSectionCircuit.classList.add('hidden');
+    
+    // Prepare extras
+    const extraIds = [];
+    const extraStrs = [];
+    
+    appState.extraTokens.forEach(item => {
+        if(item.type === 'id') extraIds.push(item.value);
+        else extraStrs.push(item.value);
+    });
+    
+    try {
+        const response = await fetch(`${API_BASE}/compute_logits`, {
+            method: 'POST',
+            headers: {'Content-Type': 'application/json'},
+            body: JSON.stringify({
+                prompt: elPrompt.value,
+                is_append_bos: elAppendBos.checked,
+                topk: 50,
+                extra_token_ids: extraIds,
+                extra_token_strs: extraStrs,
+                capture_mid: document.getElementById('capture-mid') ? document.getElementById('capture-mid').checked : false
+            })
+        });
+        
+        const res = await response.json();
+        if (response.ok) {
+            renderLogitsTable(res.data);
+            appState.tokens = res.tokens;
+            
+            // Display token count
+            if (elTokenCountDisplay) {
+                elTokenCountDisplay.textContent = `Total Input Tokens: ${res.tokens.length}`;
+            }
+            
+            elSectionLogits.classList.remove('hidden');
+            if(elSectionInputAttr) elSectionInputAttr.classList.remove('hidden');
+            elSectionCircuit.classList.remove('hidden');
+            
+            // Set default selections
+            if (res.data.length > 0) {
+                appState.selectedTargetToken = res.data[0];
+                elTopTokenDisplay.textContent = `${appState.selectedTargetToken.token_str} (ID: ${appState.selectedTargetToken.token_id})`;
+            }
+        } else {
+            alert(`Error: ${res.detail}`);
+        }
+    } catch (e) {
+        alert(`Error: ${e.message}`);
+    } finally {
+        elBtnLogits.disabled = false;
+        hideLoading();
+    }
+}
+
+function renderLogitsTable(data) {
+    elLogitsTableBody.innerHTML = '';
+    data.forEach(item => {
+        // Robust escaping for the onclick handler string
+        const escapedTokenStr = formatTokenForDisplay(item.token_str, 'data');
+
+        // VISUAL ESCAPE for the table cell
+        let visualTokenStr = formatTokenForDisplay(item.token_str, 'visual');
+
+        const row = document.createElement('tr');
+        row.innerHTML = `
+            <td>${item.rank}</td>
+            <td>${item.token_id}</td>
+            <td>${visualTokenStr}</td>
+            <td>${item.logit.toFixed(4)}</td>
+            <td><button id="btn-select-${item.token_id}" class="btn-select-ref" onclick="selectContrast(${item.token_id}, '${escapedTokenStr}', ${item.logit})">Select for Contrast</button></td>
+            <td>
+                <button class="btn-generate-cont" onclick="generateContinuation(${item.token_id}, this)">Generate</button>
+                <div class="gen-result" style="font-size: 0.85em; color: #555; margin-top: 4px; max-width: 300px; white-space: pre-wrap;"></div>
+            </td>
+        `;
+        elLogitsTableBody.appendChild(row);
+    });
+}
+
+window.generateContinuation = async function(tokenId, btn) {
+    const originalText = btn.textContent;
+    btn.disabled = true;
+    const elGenMax = document.getElementById('gen-max-tokens');
+    const maxNewTokens = elGenMax ? parseInt(elGenMax.value) : 30;
+    
+    btn.textContent = `Generating (${maxNewTokens})...`;
+    const resultDiv = btn.nextElementSibling;
+    resultDiv.textContent = "";
+
+    try {
+        const payload = {
+            prompt: elPrompt.value,
+            max_new_tokens: maxNewTokens,
+            append_token_id: tokenId
+        };
+
+        const response = await fetch(`${API_BASE}/generate`, {
+            method: 'POST',
+            headers: {'Content-Type': 'application/json'},
+            body: JSON.stringify(payload)
+        });
+
+        if (!response.ok) {
+            const err = await response.json();
+            throw new Error(err.detail || "Generation failed");
+        }
+
+        const data = await response.json();
+        resultDiv.textContent = data.generated_text;
+
+    } catch(e) {
+        resultDiv.textContent = "Error: " + e.message;
+        resultDiv.style.color = "red";
+    } finally {
+        btn.disabled = false;
+        btn.textContent = originalText;
+    }
+};
+
+window.selectContrast = function(id, name, logit) {
+    appState.selectedContrastToken = { id, name, logit };
+    
+    // Calculate Diff
+    let diffText = "";
+    if (appState.selectedTargetToken) {
+        // Find top logit (usually the first one in the list if sorted, or stored in appState)
+        // appState.selectedTargetToken comes from computeLogits result[0]
+        const topLogit = appState.selectedTargetToken.logit;
+        const diff = topLogit - logit;
+        diffText = ` | Logit Diff: ${diff.toFixed(4)}`;
+    }
+
+    elContrastTokenDisplay.textContent = `${name} (ID: ${id})${diffText}`;
+    
+    // Visual Updates
+    document.querySelectorAll('.btn-select-ref').forEach(btn => {
+        btn.textContent = "Select for Contrast";
+        btn.classList.remove('selected');
+        btn.style.backgroundColor = ""; // Reset inline style if any (or rely on class)
+    });
+    
+    const activeBtn = document.getElementById(`btn-select-${id}`);
+    if (activeBtn) {
+        activeBtn.textContent = "Selected for Contrast";
+        activeBtn.classList.add('selected');
+    }
+    
+    // Auto Update UI settings
+    elBpMode.value = "logit_diff";
+    updateBpOptions();
+    elBpStrategy.value = "by_ref_token";
+    updateBpOptions();
+    elBpRefId.value = id;
+
+    // Auto Update Input Attribution settings
+    if (elAttrBpMode) {
+        elAttrBpMode.value = "logit_diff";
+        updateAttrBpOptions();
+    }
+    if (elAttrBpStrategy) {
+        elAttrBpStrategy.value = "by_ref_token";
+        updateAttrBpOptions();
+    }
+    if (elAttrBpRefId) {
+        elAttrBpRefId.value = id;
+    }
+};
+
+function updateBpOptions() {
+    const isDiff = elBpMode.value === "logit_diff";
+    document.getElementById('diff-strategy-opts').style.display = isDiff ? 'block' : 'none';
+    
+    if (isDiff) {
+        const strategy = elBpStrategy.value;
+        elOptTopk.classList.toggle('hidden', strategy !== 'by_topk_avg');
+        elOptRef.classList.toggle('hidden', strategy !== 'by_ref_token');
+    }
+}
+
+function updateAttrBpOptions() {
+    const isDiff = elAttrBpMode.value === "logit_diff";
+    if (elAttrDiffStrategyOpts) {
+        elAttrDiffStrategyOpts.style.display = isDiff ? 'block' : 'none';
+    }
+    
+    if (isDiff) {
+        const strategy = elAttrBpStrategy.value;
+        if (elAttrOptRef) {
+             elAttrOptRef.classList.toggle('hidden', strategy !== 'by_ref_token');
+        }
+    }
+}
+
+async function computeInputAttribution() {
+    if (!appState.selectedTargetToken) {
+        alert("Please calculate logits and select a target token first.");
+        return;
+    }
+
+    // Default values if inputs are missing (e.g. reused from circuit config or hardcoded)
+    const bpK = 10; 
+
+    const payload = {
+        target_token_id: appState.selectedTargetToken.token_id,
+        contrast_token_id: appState.selectedContrastToken ? appState.selectedContrastToken.id : null,
+        backprop_config: {
+            mode: elAttrBpMode.value,
+            strategy: elAttrBpStrategy.value,
+            k: bpK,
+            ref_token_id: elAttrBpRefId && elAttrBpRefId.value ? parseInt(elAttrBpRefId.value) : null
+        }
+    };
+
+    elInputAttributionDisplay.innerHTML = "Computing...";
+    elBtnComputeAttr.disabled = true;
+
+    try {
+        const resp = await fetch(`${API_BASE}/compute_input_attribution`, {
+            method: 'POST',
+            headers: {'Content-Type': 'application/json'},
+            body: JSON.stringify(payload)
+        });
+
+        if (!resp.ok) {
+            const err = await resp.json();
+            const errMsg = (typeof err.detail === 'object') ? JSON.stringify(err.detail) : (err.detail || "Request failed");
+            throw new Error(errMsg);
+        }
+
+        const data = await resp.json();
+        // data.relevance is list of floats
+        appState.lastRelevance = data.relevance;
+        renderAttributionMap(data.relevance, appState.tokens);
+
+    } catch (e) {
+        console.error(e);
+        elInputAttributionDisplay.textContent = "Error: " + e.message;
+    } finally {
+        elBtnComputeAttr.disabled = false;
+    }
+}
+if (elAttrHideBos) {
+    elAttrHideBos.addEventListener('change', () => {
+        if (appState.lastRelevance && appState.tokens) {
+            renderAttributionMap(appState.lastRelevance, appState.tokens);
+        }
+    });
+}
+
+
+const SPECIAL_TOKENS_SET = new Set(['<|im_start|>', '<|im_end|>', '<think>', '</think>']);
+
+function renderAttributionMap(relevance, tokens) {
+    elInputAttributionDisplay.innerHTML = "";
+    
+    if (!relevance || relevance.length === 0) {
+        elInputAttributionDisplay.textContent = "No data returned.";
+        return;
+    }
+
+    const hideBos = elAttrHideBos ? elAttrHideBos.checked : false;
+    const hideSpecial = elAttrHideSpecial ? elAttrHideSpecial.checked : false;
+    const reverseSign = elAttrReverseSign ? elAttrReverseSign.checked : false;
+
+    // Helper: is this token skipped?
+    // NOTE: This logic determines both Normalization AND Rendering visibility
+    function isSkipped(idx, tokenData) {
+        if (hideBos && idx === 0) return true;
+        
+        if (hideSpecial) {
+             // tokenData could be string or object
+             let s = (typeof tokenData === 'string') ? tokenData : (tokenData.token_str || tokenData.text);
+             if (s) s = s.trim();
+             if (SPECIAL_TOKENS_SET.has(s)) return true;
+        }
+        return false;
+    }
+    
+    // Determine range for normalization (exclude BOS if hideBos is true)
+    // But we render everything from index 0 now, just styling BOS differently.
+    
+    let maxAbs = 0;
+    // Calculation Loop
+    for (let i = 0; i < relevance.length; i++) {
+        if (isSkipped(i, tokens[i])) continue; // Skip excluded tokens for dynamic range calculation
+        const r = relevance[i];
+        if (Math.abs(r) > maxAbs) maxAbs = Math.abs(r);
+    }
+    if (maxAbs === 0) maxAbs = 1;
+
+    // Helper for color
+    function getColor(val) {
+        if (Math.abs(val) === 0) return 'transparent'; // optimization
+        
+        // Red for positive, Blue for negative
+        const norm = val / maxAbs;
+        
+        // Use HSL for better control? Or RGBA.
+        // We want a clear white text if background is dark? 
+        // Or keep background light.
+        // Let's use slight alpha backgrounds.
+        
+        const alpha = Math.abs(norm); 
+        // Cap alpha to avoid being too dark/unreadable if we don't change text color
+        // But we want it visible.
+        const cappedAlpha = Math.min(alpha, 0.6); 
+
+        // Blend with white background to avoid html2canvas alpha issues (red background bug)
+        // Foreground: Red (255, 0, 0) or Blue (0, 0, 255)
+        // Background: White (255, 255, 255)
+        // Result = FG * alpha + BG * (1 - alpha)
+        
+        if (norm > 0) {
+            // Red case: R=255, G=255(1-a), B=255(1-a)
+            const gb = Math.round(255 * (1 - cappedAlpha));
+            return `rgb(255, ${gb}, ${gb})`; 
+        } else {
+            // Blue case: R=255(1-a), G=255(1-a), B=255
+            const rg = Math.round(255 * (1 - cappedAlpha));
+            return `rgb(${rg}, ${rg}, 255)`;
+        }
+    }
+
+    tokens.forEach((tok, idx) => {
+        // Handle mismatch length if any (shouldn't happen)
+        if (idx >= relevance.length) return;
+
+        let val = relevance[idx];
+        if (reverseSign) val = -val;
+
+        const span = document.createElement("span");
+        span.className = "token-span";
+        
+        // Check if this is a "Hidden" token
+        const isHidden = isSkipped(idx, tok);
+        
+        if (isHidden) {
+            // Apply special style
+            span.style.color = "#aaa";
+            span.style.backgroundColor = "#f0f0f0";
+            // span.style.textDecoration = "line-through"; // Optional? 
+        }
+
+        // tok is a string from backend
+        let tokenText = (typeof tok === 'string') ? tok : (tok.token_str || tok.text);
+        
+        // Use Central Formatter
+        const visualText = formatTokenForDisplay(tokenText, 'visual');
+        
+        // Special structural handling for newlines
+        const isNewline = (tokenText === '\n' || tokenText === '\r\n');
+        const isDoubleNewline = (tokenText === '\n\n' || tokenText === '\r\n\r\n');
+
+        if (isNewline || isDoubleNewline) {
+             span.textContent = visualText;
+             // Ensure it has shape/padding
+             span.style.display = "inline-block";
+             span.style.minWidth = isDoubleNewline ? "24px" : "12px";
+             span.style.textAlign = "center";
+        } else {
+             span.textContent = tokenText; // Keep original for normal text? OR use visual? 
+             // Usually normal text is fine, but if it has internal newlines, visualText handles it.
+             // Let's use visualText if it differs significantly? No, visualText escapes \n.
+             // For attribution map, we want visible \n, but regular text should probably wrap?
+             // Actually, for wrapped text block, literal \n is confusing.
+             // But valid tokens usually don't have internal newlines except specifically newline tokens.
+             // Let's stick to original tokenText for normal tokens, just in case.
+             // But if visualText detected a change (like mixed \n), we might want to show it.
+             if (visualText !== tokenText) span.textContent = visualText;
+             else span.textContent = tokenText;
+        }
+        
+        const normVal = val / maxAbs;
+        if (!isHidden) {
+            span.style.backgroundColor = getColor(val);
+        }
+        span.title = `Token: "${visualText}"\nPos: ${idx}\nRaw Rel: ${val.toFixed(5)}\nNorm Rel: ${normVal.toFixed(3)}`;
+
+        elInputAttributionDisplay.appendChild(span);
+
+        // Preserve line break behavior for newlines -- DISABLED per user request
+        /*
+        if (isNewline) {
+             elInputAttributionDisplay.appendChild(document.createElement("br"));
+        } else if (isDoubleNewline) {
+             elInputAttributionDisplay.appendChild(document.createElement("br"));
+             elInputAttributionDisplay.appendChild(document.createElement("br"));
+        }
+        */
+    });
+}
+
+async function computeCircuit() {
+    // Parse Layers List
+    // Expect: "0, 5, 20" -> [0, 5, 20]
+    const rawVal = elLayersList.value;
+    const layers = rawVal.split(',').map(s => parseInt(s.trim())).filter(n => !isNaN(n));
+    
+    // Sort and Validate
+    layers.sort((a,b) => a - b);
+    
+    if (layers.length < 2) {
+        alert("Please provide at least 2 layers (e.g., '0, 5, 20')");
+        return;
+    }
+    
+    elLayersList.value = layers.join(', ');
+    
+    elBtnCircuit.disabled = true;
+    elBtnCircuit.textContent = "Computing...";
+    
+    // Progress UI
+    const elProgressContainer = document.getElementById("compute-progress-container");
+    const elProgressBar = document.getElementById("compute-progress-bar");
+    const elProgressStatus = document.getElementById("progress-status-text");
+    const elProgressPercent = document.getElementById("progress-percent-text");
+
+    if(elProgressContainer) {
+        elProgressContainer.style.display = "block";
+        elProgressBar.value = 0;
+        elProgressStatus.textContent = "Starting...";
+        elProgressPercent.textContent = "0%";
+    }
+    
+    // Build Backprop Config
+    const bpConfig = {
+        mode: elBpMode.value,
+        strategy: elBpStrategy.value,
+        k: parseInt(elBpK.value) || 10,
+        ref_token_id: parseInt(elBpRefId.value) || 0,
+        contrast_rank: 2 // simplified defaults
+    };
+
+    // Pruning Config
+    const pruningMode = elPruningMode ? elPruningMode.value : "by_per_layer_cum_mass_percentile";
+    const topP = elVisTopP ? parseFloat(elVisTopP.value) : 0.9;
+    const edgeThresh = elValGlobalThresh ? parseFloat(elValGlobalThresh.value) : 0.01;
+    
+    try {
+        const response = await fetch(`${API_BASE}/compute_circuit`, {
+            method: 'POST',
+            headers: {'Content-Type': 'application/json'},
+            body: JSON.stringify({
+                backprop_config: bpConfig,
+                layers: layers,
+                pruning_mode: pruningMode,
+                top_p: topP,
+                edge_threshold: edgeThresh
+            })
+        });
+        
+        if (!response.ok) {
+             // Try to parse error
+             let errMsg = "Server Error";
+             try {
+                const err = await response.json();
+                errMsg = err.detail || errMsg;
+             } catch(e) {}
+             throw new Error(errMsg);
+        }
+        
+        // NDJSON Streaming Reader
+        const reader = response.body.getReader();
+        const decoder = new TextDecoder();
+        let buffer = '';
+        
+        // Clear previous data
+        appState.circuitData = null;
+        
+        while (true) {
+            const { value, done } = await reader.read();
+            if (done) break;
+            
+            const chunk = decoder.decode(value, {stream: true});
+            buffer += chunk;
+            const lines = buffer.split('\n');
+            buffer = lines.pop(); // Keep partial
+            
+            for (const line of lines) {
+                if (!line.trim()) continue;
+                try {
+                    const msg = JSON.parse(line);
+                    if (msg.type === "progress") {
+                        if(elProgressContainer) {
+                            const pct = Math.min(100, Math.max(0, msg.percent || 0));
+                            elProgressBar.value = pct;
+                            elProgressStatus.textContent = msg.msg;
+                            elProgressPercent.textContent = Math.round(pct) + "%";
+                        }
+                    } else if (msg.type === "graph_data") {
+                        // Received FULL graph
+                        appState.circuitData = {
+                            graph: msg.graph,
+                            pruning_details: msg.pruning_details,
+                            layers: layers // Or extract from graph nodes
+                        };
+                        appState.layout = null; // Clear cached layout
+                    } else if (msg.type === "complete") {
+                         if (appState.circuitData) {
+                            requestAnimationFrame(() => updateCircuitLayout()); // directly to layout
+                        }
+                    } else if (msg.type === "error") {
+                        throw new Error(msg.msg);
+                    }
+                } catch (e) {
+                     if (e instanceof SyntaxError) {
+                        console.error("JSON Parse Error on line:", line, e);
+                        continue;
+                     }
+                     throw e; 
+                }
+            }
+        }
+    } catch (e) {
+        alert(e.message);
+    } finally {
+        elBtnCircuit.disabled = false;
+        elBtnCircuit.textContent = "Visualize Connection";
+        if(elProgressContainer) {
+             setTimeout(() => {
+                 elProgressContainer.style.display = "none";
+             }, 2000);
+        }
+    }
+}
+
+function updateStatus(msg, type) {
+    elStatus.textContent = msg;
+    elStatus.className = `status ${type}`;
+}
+
+// --- Visualization Logic ---
+
+// NEW: Separated heavy layout calculation from rendering
+function updateCircuitLayout() {
+    if (!appState.circuitData || !appState.circuitData.graph) return;
+    showLoading("Updating Layout...");
+    
+    // Yield to browser to render spinner
+    setTimeout(() => {
+        try {
+            _updateCircuitLayoutInternal();
+        } catch(e) {
+            console.error(e);
+        } finally {
+             hideLoading();
+        }
+    }, 50);
+}
+
+function _updateCircuitLayoutInternal() {
+    if (!appState.circuitData || !appState.circuitData.graph) return;
+    
+    // Destructure graph early
+    const { graph, layers: requestedLayers } = appState.circuitData;
+    
+    // Define robust accessors
+    // NetworkX alignment: v2 uses 'links', v3 uses 'edges'
+    const nodes = graph.nodes || [];
+    const edges = graph.edges || graph.links || [];
+    
+    // DEBUG: Log Graph Data Structure
+    console.log("Graph Data Received:", graph);
+    if(nodes.length > 0) {
+        console.log("Sample Node:", nodes[0]);
+    }
+    if(edges.length > 0) {
+        console.log("Sample Link/Edge:", edges[0]);
+    } else {
+        console.log("No links/edges found. (Graph might be disconnected or threshold too high)");
+    }
+    
+    // Check tokens
+    if (!appState.tokens || appState.tokens.length === 0) {
+        console.error("appState.tokens is missing! Cannot render graph. Please run Compute Logits first.");
+        alert("Error: Tokens missing. Please run 'Compute Logits' first.");
+        return;
+    }
+
+    // Check toggle listeners
+    const toggle = document.getElementById('hide-bos-node');
+    if (toggle && !toggle.hasAttribute('data-listening')) {
+        toggle.addEventListener('change', updateCircuitLayout); // Need layout update to filter BOs
+        toggle.setAttribute('data-listening', 'true');
+    }
+    const showAllToggle = document.getElementById('show-all-tokens');
+    if (showAllToggle && !showAllToggle.hasAttribute('data-listening')) {
+        showAllToggle.addEventListener('change', updateCircuitLayout); 
+        showAllToggle.setAttribute('data-listening', 'true');
+    }
+    
+    const showNodeValuesToggle = document.getElementById('show-node-values');
+    if (showNodeValuesToggle && !showNodeValuesToggle.hasAttribute('data-listening')) {
+        showNodeValuesToggle.addEventListener('change', drawCircuit); // Visual only
+        showNodeValuesToggle.setAttribute('data-listening', 'true');
+    }
+
+    // const { graph, layers: requestedLayers } = appState.circuitData; // Moved up
+    const tokens = appState.tokens;
+    const layerHeight = parseFloat(elVisLayerSpacing ? elVisLayerSpacing.value : 300);
+    
+    const fullSeqLen = tokens.length;
+    const hideFirstToken = document.getElementById('hide-bos-node') && document.getElementById('hide-bos-node').checked;
+    const showAllTokens = elShowAllTokens ? elShowAllTokens.checked : true;
+    
+    // Data Slicing (if hiding BOS)
+    let displaySeqLen = fullSeqLen;
+    let displayTokens = tokens;
+    let startTokenIdx = 0;
+    
+    if (hideFirstToken && fullSeqLen > 1) {
+        displayTokens = tokens.slice(1);
+        displaySeqLen = fullSeqLen - 1;
+        startTokenIdx = 1;
+    }
+
+    // Process Graph Data
+    // Nodes are { id: [layer, token], layer: L, token: T, relevance: R }
+    // const nodes = graph.nodes; // Moved up
+    // const edges = graph.links; // Moved up
+    
+    // Identify Layers involved in graph
+    // Use requestedLayers if available, or infer from graph
+    let layersList = requestedLayers || [];
+    if (layersList.length === 0) {
+        const layersSet = new Set(nodes.map(n => n.layer));
+        layersList = Array.from(layersSet).sort((a,b)=>a-b);
+    }
+    
+    // Determine active nodes (those in graph)
+    const activeNodeMap = new Map(); // key: "layer,token" -> nodeObj
+    let globalMaxValRel = 0;
+    let globalMaxValEdge = 0;
+
+    nodes.forEach(n => {
+        // Filter BOS if needed
+        if (hideFirstToken && n.token === 0) return; // Assume BOS is 0? Use n.token index check
+        if (hideFirstToken && n.token < startTokenIdx) return;
+        
+        const key = `${n.layer},${n.token}`;
+        activeNodeMap.set(key, n);
+        globalMaxValRel = Math.max(globalMaxValRel, Math.abs(n.relevance || 0));
+    });
+    
+    edges.forEach(e => {
+        // e.source and e.target are arrays [L, T]
+        const srcL = e.source[0];
+        const srcT = e.source[1];
+        const tgtL = e.target[0];
+        const tgtT = e.target[1];
+        
+        if (hideFirstToken && (srcT < startTokenIdx || tgtT < startTokenIdx)) return;
+        
+        globalMaxValEdge = Math.max(globalMaxValEdge, Math.abs(e.weight || 0));
+    });
+    
+    const maxNorm = Math.max(globalMaxValRel, globalMaxValEdge, 1e-9);
+
+    // Layout Geometry Setup
+    const nHops = layersList.length > 1 ? layersList.length - 1 : 1; 
+    const totalHeight = Math.max(500, nHops * layerHeight + 150);
+    const margin = { top: 50, left: 120, right: 50, bottom: 50 };
+    
+    
+    // Generate Layout Nodes
+    const nodesByLayer = {}; 
+    const allNodesFlat = [];
+    const layerTotals = {};
+    
+    // Helper to get active node data
+    const getActiveData = (l, t) => activeNodeMap.get(`${l},${t}`);
+
+    // DYNAMIC LAYOUT LOGIC
+    // 1. Identify "Active Columns": Any token index that participates in an EDGE.
+    // This filters out columns that might have nodes but no connections (pruned).
+    let activeTokenIndices = new Set();
+    
+    // Use edges to find connected tokens
+    edges.forEach(e => {
+        // e.source/target are [layer, token]
+        activeTokenIndices.add(e.source[1]); 
+        activeTokenIndices.add(e.target[1]);
+    });
+    
+    // Also include target node tokens? 
+    // Usually target is connected, but what if we have a top-p that prunes EVERYTHING?
+    // We should probably show the target node column regardless, so the user sees where they started.
+    // If backend prunes everything, we might have edges=[], nodes=[target].
+    // Let's iterate nodes too, but check a flag or just rely on edges.
+    // User Request: "hide the token positions whose nodes in all layers do not have any connections".
+    // This implies strictly edges.
+    // BUT: If the target node has no incoming edges (e.g. threshold too high), it is isolated.
+    // Should we hide the target node? That seems confusing.
+    // Let's prioritize edges, but maybe keep target? 
+    // For now, strict edge adherence based on user prompt.
+    // If empty set (graph empty), we might show nothing or just full seq?
+    if (activeTokenIndices.size === 0 && !showAllTokens && nodes.length > 0) {
+        // Fallback: Show at least the nodes that exist (like target)
+        nodes.forEach(n => activeTokenIndices.add(n.token));
+    }
+
+    // 2. Define "Render Columns"
+    let renderColumns = [];
+    if (showAllTokens) {
+        // All tokens in display range
+        for(let i=0; i<displaySeqLen; i++) renderColumns.push(startTokenIdx + i);
+    } else {
+        // Only active tokens
+        renderColumns = Array.from(activeTokenIndices).sort((a,b) => a - b);
+        // Filter by hideFirstToken if needed (activeNodeMap already filtered? No, activeNodeMap built from nodes list)
+        // Ensure strictly respecting hideFirstToken setting
+        if (hideFirstToken) {
+            renderColumns = renderColumns.filter(t => t >= startTokenIdx);
+        }
+    }
+    
+    // 3. Spacing based on Render Columns count
+    const nodesInRowForSpacing = renderColumns.length; 
+    const minW = nodesInRowForSpacing * (showAllTokens ? 15 : 30) + 100;
+    const width = Math.max(800, minW);
+    const drawingW = width - margin.left - margin.right;
+    const nodeSpacing = drawingW / Math.max(1, nodesInRowForSpacing);
+
+    const yPositions = {};
+    layersList.forEach((lIdx, i) => {
+        yPositions[lIdx] = (totalHeight - margin.bottom) - (i * layerHeight);
+    });
+
+    layersList.forEach((lIdx) => {
+        const y = yPositions[lIdx];
+        const rowNodeMap = {}; 
+        let runningTotal = 0;
+        
+        // Iterate Render Columns
+        renderColumns.forEach((tokenIdx, visualColIdx) => {
+            const activeData = getActiveData(lIdx, tokenIdx);
+            const rel = activeData ? (activeData.relevance || 0) : 0;
+            const normRel = Math.abs(rel) / maxNorm;
+            
+            // accumulate total only if we want layer sum to reflect EVERYTHING or just visible?
+            // Usually layer sum is total relevance. If we hide tokens, do we hide their relevance from sum?
+            // Let's sum only active data found in graph.
+            if (activeData) runningTotal += rel;
+            
+            // X alignment: Based on visual column index
+            const x = margin.left + visualColIdx * nodeSpacing + (nodeSpacing/2);
+            
+            // Token String
+            const rawToken = (tokenIdx < tokens.length) ? tokens[tokenIdx] : "?";
+            const tokenStr = (typeof rawToken === 'string') ? rawToken : (rawToken.token_str || "?");
+
+            const node = {
+                layer: lIdx,
+                index: tokenIdx, // True token index
+                visualIndex: visualColIdx, // For gap logic
+                token: { token_str: tokenStr }, 
+                x: x,
+                y: y,
+                rel: rel,
+                normRel: normRel,
+                isActive: !!activeData
+            };
+            allNodesFlat.push(node);
+            rowNodeMap[tokenIdx] = node;
+        });
+        
+        nodesByLayer[lIdx] = rowNodeMap;
+        layerTotals[lIdx] = runningTotal;
+    });
+
+    // Generate Edges
+    const visibleEdges = [];
+    
+    edges.forEach(e => {
+        const srcL = e.source[0]; 
+        const srcT = e.source[1]; 
+        const tgtL = e.target[0]; 
+        const tgtT = e.target[1];
+        
+        // Lookup in our generated layout nodes
+        // nodesByLayer[layer][tokenIndex]
+        if (!nodesByLayer[srcL] || !nodesByLayer[tgtL]) return;
+        
+        const sNode = nodesByLayer[srcL][srcT];
+        const tNode = nodesByLayer[tgtL][tgtT];
+        
+        if (sNode && tNode) {
+            const val = e.weight;
+            const normVal = Math.abs(val) / maxNorm;
+            
+            visibleEdges.push({
+                source: sNode,
+                target: tNode,
+                val: val,
+                normVal: normVal
+            });
+        }
+    });
+
+    // STORE LAYOUT
+    appState.layout = {
+        width,
+        totalHeight,
+        margin,
+        nodeSpacing,
+        maxNorm,
+        layersList,
+        nodesByLayer,
+        allNodesFlat,
+        visibleEdges,
+        activeIndicesList: [], // Not used in this sparse logic
+        layerTotals,
+        displaySeqLen,
+        displayTokens, // Add this back for gap drawing
+        displayConnections: [] // Legacy compat
+    };
+    
+    drawCircuit();
+}
+
+function drawCircuit() {
+    if (!appState.layout) {
+        // Init layout if data exists
+        if(appState.circuitData) updateCircuitLayout();
+        return;
+    }
+    
+    const L = appState.layout;
+    
+    // Check Spacing Updates first (Geometry Recalc)
+    const currentLayerHeight = parseFloat(elVisLayerSpacing ? elVisLayerSpacing.value : 300);
+    // Use stored layers count for recalculating total height
+    // We infer layersList from layout
+    const nHops = L.displayConnections ? L.displayConnections.length : 1; 
+    let newTotalHeight = Math.max(500, nHops * currentLayerHeight + 150);
+    
+    if (L.layersList && L.layersList.length > 0) {
+        newTotalHeight = Math.max(500, (L.layersList.length - 1) * currentLayerHeight + 200);
+
+        // Pre-calculate Y for each layer (Optimization: O(Layers) instead of O(Nodes*Layers))
+        const layerYMap = {};
+        L.layersList.forEach((lid, idx) => {
+            layerYMap[lid] = (newTotalHeight - L.margin.bottom) - (idx * currentLayerHeight);
+        });
+
+        // Update Node Y Positions in place
+        L.allNodesFlat.forEach(node => {
+             if (layerYMap.hasOwnProperty(node.layer)) {
+                 node.y = layerYMap[node.layer];
+             }
+        });
+        
+        // Edge coordinates update automatically since they reference node objects
+    }
+    L.totalHeight = newTotalHeight;
+
+    const ctx = elCanvas.getContext('2d');
+    
+    // Check canvas dims
+    if (elCanvas.width !== L.width || elCanvas.height !== L.totalHeight) {
+        elCanvas.width = L.width;
+        elCanvas.height = L.totalHeight;
+    }
+    
+    // Vis Settings
+    const strengthScale = parseFloat(elVisStrength.value);
+    const showNodeValues = document.getElementById('show-node-values') ? document.getElementById('show-node-values').checked : true;
+    const showAllTokens = elShowAllTokens ? elShowAllTokens.checked : true;
+    const hideFirstToken = document.getElementById('hide-bos-node') && document.getElementById('hide-bos-node').checked;
+    
+    // Add startTokenIdx logic here
+    let startTokenIdx = 0;
+    if (hideFirstToken && appState.tokens && appState.tokens.length > 1) {
+        startTokenIdx = 1;
+    }
+
+    ctx.clearRect(0, 0, L.width, L.totalHeight);
+    
+    // --- Interaction ---
+    const mouseX = appState.mouseX || -1;
+    const mouseY = appState.mouseY || -1;
+    function checkHit(x, y) { return Math.sqrt((x-mouseX)**2 + (y-mouseY)**2) < 8; }
+    
+    let hoveredNode = null;
+    for (const n of L.allNodesFlat) {
+        if (checkHit(n.x, n.y)) {
+            hoveredNode = n;
+            break;
+        }
+    }
+    const activeNode = hoveredNode || selectedNode;
+    const isHighlightActive = !!activeNode;
+    
+    // --- Highlight Computation (One-Hop) ---
+    const highlightedNodes = new Set();
+    const highlightedEdges = new Set(); // store edges to paint bold
+    
+    if (activeNode) {
+        highlightedNodes.add(activeNode);
+        
+        // Use visibleEdges directly
+        L.visibleEdges.forEach(e => {
+             // 1. Outgoing
+             if (e.source === activeNode) {
+                 highlightedNodes.add(e.target);
+                 highlightedEdges.add(e);
+             }
+             // 2. Incoming
+             if (e.target === activeNode) {
+                 highlightedNodes.add(e.source);
+                 highlightedEdges.add(e);
+             }
+        });
+    }
+
+    // --- Draw Edges ---
+    // Sort edges: Passive first, Active on top
+    // Actually just draw active afterwards.
+    
+    L.visibleEdges.forEach(e => {
+        const isActive = isHighlightActive && highlightedEdges.has(e);
+        const isDim = isHighlightActive && !isActive;
+        
+        // If dim, maybe skip drawing very thin lines to save perf?
+        // Or alpha.
+        
+        let width = e.normVal * strengthScale;
+        if (width < 0.5 && isDim) return; // Culling
+        
+        width = Math.max(width, 0.5);   // Minimal width for visibility
+
+        let alpha = isDim ? 0.1 : 1.0;
+        if (isActive) {
+            alpha = 1.0;
+            width = Math.max(width, 1.0); // Boost active
+        }
+        
+        const color = e.val >= 0 ? `rgba(211, 47, 47, ${alpha})` : `rgba(25, 118, 210, ${alpha})`;
+        
+        ctx.beginPath();
+        ctx.moveTo(e.source.x, e.source.y);
+        ctx.lineTo(e.target.x, e.target.y);
+        ctx.strokeStyle = color;
+        ctx.lineWidth = width;
+        
+        // draw active later? No, simple batch is fine mostly.
+        ctx.stroke();
+
+        if (isActive) { // Use isActive logic for edge label drawing
+            // Draw Edge Label
+            const midX = (e.source.x + e.target.x)/2;
+            const midY = (e.source.y + e.target.y)/2;
+            ctx.save();
+            const dx = e.target.x - e.source.x;
+            const dy = e.target.y - e.source.y;
+            let angle = Math.atan2(dy, dx);
+            if (Math.abs(angle) > Math.PI / 2) angle += Math.PI; // Correct text orientation
+            
+            ctx.translate(midX, midY);
+            ctx.rotate(angle);
+            
+            ctx.font = 'bold 10px Arial';
+            const text = e.val.toFixed(2);
+            const metrics = ctx.measureText(text);
+            const p = 2; // padding
+            
+            ctx.fillStyle = 'rgba(255, 255, 255, 0.9)';
+            ctx.fillRect(-metrics.width/2 - p, -6, metrics.width + 2*p, 12);
+            
+            ctx.fillStyle = e.val >= 0 ? '#b71c1c' : '#0d47a1';
+            ctx.textAlign = 'center';
+            ctx.textBaseline = 'middle';
+            ctx.fillText(text, 0, 0);
+            ctx.restore();
+        }
+    });
+    
+    // --- Draw Nodes ---
+    // Re-use Gap logic variables
+    let prevIdx = -1;
+    let listIdx = -1;
+    const displayedGaps = [];
+    const displayTokens = L.displayTokens;
+    
+    L.allNodesFlat.forEach((node) => {
+        // Logic for tracking gap indices per layer
+        if (node.layer === L.layersList[0]) {
+            listIdx++;
+        } else {
+            // Reset for other layers? Actually gap logic only applied to bottom layer in original code
+        }
+
+        const isSelfActive = activeNode === node;
+        const isConnected = highlightedNodes.has(node) && !isSelfActive;
+        const isInteractive = isHighlightActive;
+        const isDim = isInteractive && !isSelfActive && !isConnected;
+
+        let radius = (node.normRel * strengthScale);
+        radius = Math.max(radius, 2);
+        if (isConnected) radius = Math.max(radius, strengthScale * 0.75); 
+        if (isSelfActive) radius = Math.max(radius, strengthScale * 1.25);
+        
+        let fillStyle = '#e0e0e0';
+        if (node.rel > 0.001) fillStyle = '#ef9a9a'; 
+        else if (node.rel < -0.001) fillStyle = '#90caf9'; 
+        else fillStyle = '#cfd8dc'; 
+
+        if (isHighlightActive && !isDim) {
+             if (node.rel > 0.001) fillStyle = '#d32f2f'; 
+             else if (node.rel < -0.001) fillStyle = '#1976d2'; 
+             else fillStyle = '#455a64'; 
+        }
+
+        if (isDim) {
+            fillStyle = '#f5f5f5';
+            radius = 3; 
+        }
+        
+        ctx.beginPath();
+        ctx.arc(node.x, node.y, radius, 0, 2*Math.PI);
+        ctx.fillStyle = fillStyle;
+        ctx.fill();
+        
+        if (!isDim) {
+             ctx.strokeStyle = (isSelfActive || isConnected) ? '#333' : '#bbb';
+             ctx.lineWidth = 1;
+             ctx.stroke();
+        }
+        
+        // Labels (Layer)
+        // Check if node is the first visible node in its layer OR if it's the very first node of the sequence
+        // We use activeIndicesList logic if available, otherwise just use "first in row" logic
+        const isFirstInLayer = (node.index === 0) || (Object.values(L.nodesByLayer[node.layer])[0] === node);
+
+        if (isFirstInLayer) { 
+            ctx.fillStyle = '#2c3e50';
+            ctx.font = 'bold 12px Arial';
+            ctx.textAlign = 'right';
+            
+            let label = "";
+            if (Array.isArray(node.layer)) {
+                 label = `L${node.layer[0]} (${node.layer[1]})`;
+            } else {
+                 label = node.layer === -1 ? "Embedding" : `Layer ${node.layer}`;
+            }
+            
+            ctx.fillText(label, L.margin.left - 20, node.y + 4);            
+            const total = L.layerTotals[node.layer] || 0;
+            ctx.font = '11px Arial';
+            ctx.fillText(`Sum: ${total.toFixed(2)}`, L.margin.left - 20, node.y + 18);
+        }
+        
+        // Labels (Tokens)
+        const isBottomLayer = (node.layer === L.layersList[0]);
+        let showToken = isBottomLayer; 
+        if (isHighlightActive && (isSelfActive || isConnected)) showToken = true;
+        
+        if (showToken) {
+             ctx.fillStyle = '#000';
+             ctx.font = isSelfActive ? 'bold 11px Arial' : '10px Arial';
+             ctx.textAlign = 'center';
+             ctx.save();
+             
+             if (isBottomLayer && node.layer !== L.layersList[L.layersList.length-1]) {
+                 ctx.translate(node.x, node.y + 15);
+                 ctx.rotate(Math.PI/4);
+             } else {
+                 ctx.translate(node.x, node.y - 12);
+                 ctx.rotate(-Math.PI/4);
+             }
+             
+             let txt = formatTokenForDisplay(node.token.token_str, 'visual');
+             ctx.fillText(txt, 0, 0);
+             ctx.restore();
+             
+             // Gap Marker Logic (Bottom Layer Only)
+             if (isBottomLayer && !showAllTokens) {
+                  // Preceding Gap (Before first visible token)
+                  // If this is the first visible token (listIdx or visualIndex 0), check if there are hidden tokens before it
+                  const trueStartIdx = hideFirstToken ? startTokenIdx : 0;
+                  
+                  if (node.visualIndex === 0 && node.index > trueStartIdx) {
+                      const gx = node.x - L.nodeSpacing/2;
+                      const gy = node.y + 45;
+                      const hidden = [];
+                      for(let k=trueStartIdx; k<node.index; k++) {
+                          if(displayTokens && k < displayTokens.length) {
+                              const tObj = displayTokens[k];
+                              const tStr = (typeof tObj === 'string') ? tObj : tObj.token_str;
+                              hidden.push(formatTokenForDisplay(tStr, 'visual'));
+                          }
+                      }
+                      if (hidden.length > 0) displayedGaps.push({x: gx, y: gy, hidden: hidden});
+                  }
+                  
+                  // Inter-token Gap
+                  // Previous Visible Node Index vs Current Node Index
+                  if (prevIdx !== -1 && (node.index - prevIdx) > 1) {
+                      // We can use visualIndex to backtrack to previous node X
+                      // node.visualIndex is current i. Previous was i-1.
+                      // xPrev = x of column (node.visualIndex - 1)
+                      const xPrev = node.x - L.nodeSpacing;
+                      
+                      const gx = (xPrev + node.x)/2;
+                      const gy = node.y + 45;
+                      const hidden = [];
+                      for(let k=prevIdx+1; k<node.index; k++) {
+                            if(displayTokens && k < displayTokens.length) {
+                                const tObj = displayTokens[k];
+                                const tStr = (typeof tObj === 'string') ? tObj : tObj.token_str;
+                                hidden.push(formatTokenForDisplay(tStr, 'visual'));
+                            }
+                      }
+                      if (hidden.length > 0) displayedGaps.push({x: gx, y: gy, hidden: hidden});
+                  }
+                  prevIdx = node.index;
+             }
+        }
+
+        if (showNodeValues) {
+             const valText = node.rel.toFixed(2);
+             ctx.font = '9px Arial';
+             ctx.fillStyle = node.rel >= 0 ? '#b71c1c' : '#0d47a1';
+             ctx.textAlign = 'center';
+             let textY = node.y - 14; 
+             if (!isBottomLayer) textY = node.y - 16; 
+             if (isSelfActive) ctx.font = 'bold 10px Arial';
+             ctx.fillText(valText, node.x, textY);
+        }
+    });
+    // --- Draw Gaps Loop ---
+    let hoveredGap = null;
+    // const { displayTokens } = L; // Already defined above
+
+    displayedGaps.forEach(gap => {
+        const mx = appState.mouseX;
+        const my = appState.mouseY;
+        const dist = Math.sqrt(Math.pow(gap.x - mx, 2) + Math.pow(gap.y - my, 2));
+        
+        if (dist < 15) hoveredGap = gap;
+        
+        ctx.save();
+        ctx.translate(gap.x, gap.y);
+        ctx.fillStyle = (dist < 15) ? '#e0e0e0' : '#f5f5f5';
+        ctx.strokeStyle = '#bdbdbd';
+        ctx.lineWidth = 1;
+        
+        ctx.beginPath();
+        if(ctx.roundRect) ctx.roundRect(-12, -8, 24, 16, 4);
+        else ctx.rect(-12, -8, 24, 16); 
+        ctx.fill();
+        ctx.stroke();
+        
+        ctx.fillStyle = '#616161';
+        ctx.font = 'bold 12px Arial';
+        ctx.textAlign = 'center';
+        ctx.textBaseline = 'middle';
+        ctx.fillText("...", 0, -2);
+        ctx.restore();
+    });
+    
+    // --- Tooltip Updating ---
+    const elTooltip = document.getElementById('tooltip');
+    if (hoveredGap && elTooltip) {
+        elTooltip.style.display = 'block';
+        elTooltip.style.opacity = '1';
+        
+        // Update content FIRST to measure size
+        const fullText = hoveredGap.hidden ? hoveredGap.hidden.join(" ") : "";
+        elTooltip.textContent = fullText || "[Empty]";
+        elTooltip.style.maxWidth = "300px";
+
+        const rect = elCanvas.getBoundingClientRect();
+        const gapX = rect.left + hoveredGap.x;
+        const gapY = rect.top + hoveredGap.y;
+        
+        // Default: Bottom-Right
+        let finalLeft = gapX + 10;
+        let finalTop = gapY + 10;
+        
+        // Check Bounds
+        const tooltipRect = elTooltip.getBoundingClientRect();
+        
+        // Vertical overflow (flip up)
+        if (finalTop + tooltipRect.height > window.innerHeight) {
+            finalTop = gapY - tooltipRect.height - 10;
+        }
+        
+        // Horizontal overflow (flip left)
+        if (finalLeft + tooltipRect.width > window.innerWidth) {
+            finalLeft = gapX - tooltipRect.width - 10;
+        }
+        
+        elTooltip.style.left = finalLeft + 'px'; 
+        elTooltip.style.top = finalTop + 'px';
+        
+    } else if (elTooltip) {
+        elTooltip.style.display = 'none';
+        elTooltip.style.opacity = '0';
+    }
+    
+    // Mouse Interaction Tracking
+    // Note: Re-binding these every draw call is inefficient but matches original logic structure.
+    // Ideally move these out of drawCircuit.
+    elCanvas.onmousemove = function(e) {
+         const rect = elCanvas.getBoundingClientRect();
+         appState.mouseX = e.clientX - rect.left;
+         appState.mouseY = e.clientY - rect.top;
+         requestAnimationFrame(drawCircuit);
+    };
+    
+    elCanvas.onmouseleave = function() {
+        appState.mouseX = -1;
+        appState.mouseY = -1;
+        requestAnimationFrame(drawCircuit);
+    };
+
+    elCanvas.onclick = function(e) {
+        const rect = elCanvas.getBoundingClientRect();
+        const x = e.clientX - rect.left;
+        const y = e.clientY - rect.top;
+        
+        let hit = null;
+        if(L && L.allNodesFlat) {
+            for (const n of L.allNodesFlat) {
+                 if (Math.sqrt(Math.pow(x - n.x, 2) + Math.pow(y - n.y, 2)) < 10) {
+                     hit = n;
+                     break;
+                 }
+            }
+        }
+        
+        if (hit) {
+            if (activeNode && activeNode.layer === hit.layer && activeNode.index === hit.index && !hoveredNode) {
+                selectedNode = null; 
+            } else {
+                selectedNode = hit;
+            }
+        } else {
+            selectedNode = null;
+        }
+        requestAnimationFrame(drawCircuit);
+    };
+}
+
+// Initialization and Cleanup
+window.addEventListener('DOMContentLoaded', async () => {
+    // Call cleanup endpoint to release GPU memory on page refresh
+    try {
+        await fetch(`${API_BASE}/cleanup`, { method: 'POST' });
+        console.log("Backend memory cleaned up.");
+    } catch (e) {
+        console.warn("Failed to cleanup backend memory:", e);
+    }
+    
+    // Init Datasets
+    await fetchDatasets();
+});
+
+// --- Trace Loading Logic ---
+
+async function fetchDatasets() {
+    try {
+        const resp = await fetch(`${API_BASE}/datasets`);
+        const data = await resp.json();
+        
+        elTraceDataset.innerHTML = '<option value="">-- Select --</option>';
+        data.datasets.forEach(ds => {
+            const opt = document.createElement('option');
+            opt.value = ds;
+            opt.textContent = ds;
+            elTraceDataset.appendChild(opt);
+        });
+    } catch (e) {
+        console.error("Failed to fetch datasets:", e);
+    }
+}
+
+async function fetchTraces(dataset) {
+    elTraceFile.disabled = true;
+    elTraceFile.innerHTML = '<option value="">Loading...</option>';
+    showLoading("Fetching Traces...");
+    
+    try {
+        const resp = await fetch(`${API_BASE}/traces/${dataset}`);
+        if (!resp.ok) throw new Error("Failed");
+        
+        const data = await resp.json();
+        
+        elTraceFile.innerHTML = '<option value="">-- Select Trace --</option>';
+        data.traces.forEach(t => {
+            const opt = document.createElement('option');
+            opt.value = t;
+            opt.textContent = t;
+            elTraceFile.appendChild(opt);
+        });
+        elTraceFile.disabled = false;
+    } catch (e) {
+        console.error("Failed to fetch traces:", e);
+        elTraceFile.innerHTML = '<option value="">Error</option>';
+    } finally {
+        hideLoading();
+    }
+}
+
+async function loadTraceDetails(dataset, traceId) {
+    showLoading("Loading Trace Config...");
+    try {
+        const resp = await fetch(`${API_BASE}/trace_details/${dataset}/${traceId}`);
+        if (!resp.ok) throw new Error("Failed");
+        
+        const data = await resp.json();
+        
+        // Update UI
+        if(data.model_path) elModelPath.value = data.model_path;
+        if(data.prompt) elPrompt.value = data.prompt;
+
+        // Populate "Prompt + Original Completion"
+        if(elPromptOrig) {
+            // Prefer raw_prompt if available (added to backend), else prompt (which might be full concat)
+            const basePrompt = data.raw_prompt !== undefined ? data.raw_prompt : (data.prompt || "");
+            elPromptOrig.value = basePrompt + (data.completion || "");
+        }
+        
+        // Set Config Defaults as requested
+        elQuant.checked = data.quantization; // Default false
+        elModelDtype.value = data.dtype || "bfloat16"; 
+        
+        // Determine "Other" candidates data
+        let otherData = null;
+        let otherLabel = "Other Model Candidates";
+
+        if (data.other_candidates && Object.keys(data.other_candidates).length > 0) {
+            // Pick first likely candidate key
+            const keys = Object.keys(data.other_candidates);
+            const key = keys[0];
+            otherData = data.other_candidates[key];
+            
+            // Heuristic for label
+            if (key === '4b') otherLabel = "Qwen3-4B Candidates";
+            else if (key === '1.7b' || key === '1_7b') otherLabel = "Qwen3-1.7B Candidates";
+            else if (key.toLowerCase().includes('qwen')) otherLabel = key;
+            else otherLabel = key + " Candidates";
+            
+        } else if (data.topk_token_explore_4b && data.topk_token_explore_4b.length) {
+            // Legacy Fallback
+            otherData = data.topk_token_explore_4b;
+            otherLabel = "Qwen3-4B Candidates";
+        }
+
+        // Render Exploration Data if available
+        if (elSectionTraceExplore) {
+            let hasContent = false;
+            
+            if (data.topk_token_explore && data.topk_token_explore.length) {
+                renderExploreTable(elContainerExploreOriginal, data.topk_token_explore);
+                hasContent = true;
+            } else {
+                elContainerExploreOriginal.innerHTML = "<div style='padding:10px;'>No data</div>";
+            }
+            
+            if (otherData && otherData.length) {
+                renderExploreTable(elContainerExplore4b, otherData);
+                if (elHeaderExploreOther) elHeaderExploreOther.textContent = otherLabel;
+                hasContent = true;
+            } else {
+                elContainerExplore4b.innerHTML = "<div style='padding:10px; color:#666;'>No exploration data available.</div>";
+                if (elHeaderExploreOther) elHeaderExploreOther.textContent = otherLabel;
+            }
+
+            if (hasContent) {
+                 elSectionTraceExplore.classList.remove('hidden');
+            } else {
+                 elSectionTraceExplore.classList.add('hidden');
+            }
+        }
+
+        // Flash success?
+        console.log("Trace loaded:", data);
+    } catch (e) {
+        alert("Failed to load trace details: " + e.message);
+    } finally {
+        hideLoading();
+    }
+}
+
+// Init Trace Listeners
+if (elTraceDataset) {
+    elTraceDataset.addEventListener('change', (e) => {
+        const ds = e.target.value;
+        if (ds) {
+            fetchTraces(ds);
+        } else {
+            elTraceFile.innerHTML = '<option value="">-- Select Dataset First --</option>';
+            elTraceFile.disabled = true;
+        }
+    });
+}
+
+if (elTraceFile) {
+    elTraceFile.addEventListener('change', (e) => {
+        const t = e.target.value;
+        const ds = elTraceDataset.value;
+        if (t && ds) {
+            loadTraceDetails(ds, t);
+        }
+    });
+}
+// ---------------------------
+
+function generateLayerPresets() {
+   // Helper update
+}
+
+function setLayerPreset(mode) {
+    const N = appState.n_layers || 28;
+    let layers = [];
+    if (mode === 'all') {
+        layers.push(-1);
+        for(let i=0; i<N; i++) layers.push(i);
+    } else {
+        // Default: 5 parts
+        // 0, N/4, 2N/4, 3N/4, N-1
+        const steps = 5;
+        const stepSize = (N - 1) / steps;
+        const set = new Set();
+        set.add(-1); // Always include Embedding as requested
+        for(let i=0; i<=steps; i++) {
+             set.add(Math.round(i * stepSize));
+        }
+        layers = Array.from(set).sort((a,b)=>a-b);
+    }
+    elLayersList.value = layers.join(', ');
+}
+
+// FORMAT HELPER
+function formatTokenForDisplay(tokenStr, escapeMode='visual') {
+    if (!tokenStr) return tokenStr;
+    
+    // 1. Escaping for Data Attributes or Generic Log usage
+    if (escapeMode === 'data') {
+         return tokenStr.replace(/\\/g, '\\\\')
+            .replace(/'/g, "\\'")
+            .replace(/\n/g, '\\n')
+            .replace(/\r/g, '\\r')
+            .replace(/"/g, '&quot;');
+    }
+    
+    // 2. Visual Display (Visible \n)
+    if (escapeMode === 'visual') {
+        if (tokenStr === '\n') return '\\n';
+        if (tokenStr === '\n\n') return '\\n\\n';
+        if (tokenStr === '\r\n') return '\\r\\n';
+        
+        // Check for mixed content
+        if (tokenStr.trim() === '' && tokenStr.length > 0) {
+            // It is whitespace
+            if (tokenStr === ' ') return '␣'; // Optional
+        }
+        
+        // If containing newlines mixed with text, escape the newlines
+        if (tokenStr.includes('\n')) return tokenStr.replace(/\n/g, '\\n');
+        
+        return tokenStr;
+    }
+    
+    return tokenStr;
+}
+
+// Initialize Preset on Load Model success (or when appState.n_layers is known)
+document.addEventListener('DOMContentLoaded', () => {
+    setLayerPreset('default');
+});
+
+
+function renderExploreTable(container, data) {
+    if (!container) return;
+    
+    if (!data || data.length === 0) {
+        container.innerHTML = '<div style="padding:10px; color: #888;">No exploration data available.</div>';
+        return;
+    }
+    
+    let html = `
+    <table style="width: 100%; border-collapse: collapse; font-size: 0.85em;">
+        <thead>
+            <tr style="background: #f8f8f8; text-align: left; position: sticky; top: 0; z-index: 10;">
+                <th style="padding: 6px; border-bottom: 2px solid #ddd;">Rank</th>
+                <th style="padding: 6px; border-bottom: 2px solid #ddd;">Token</th>
+                <th style="padding: 6px; border-bottom: 2px solid #ddd;">Logit</th>
+                <th style="padding: 6px; border-bottom: 2px solid #ddd;">Res</th>
+                <th style="padding: 6px; border-bottom: 2px solid #ddd;">Completion Start</th>
+            </tr>
+        </thead>
+        <tbody>
+    `;
+    
+    data.forEach(item => {
+        const isCorrect = item.eval_result === true;
+        const correctColor = isCorrect ? '#2e7d32' : '#c62828';
+        const correctBg = isCorrect ? '#e8f5e9' : '#ffebee';
+        const correctIcon = isCorrect ? 'OK' : 'Fail';
+        
+        let tokenStr = item.token_str || "";
+        // Use the existing formatTokenForDisplay helper if available, or simple replacement
+        let displayToken = tokenStr;
+        if (typeof formatTokenForDisplay === 'function') {
+            displayToken = formatTokenForDisplay(tokenStr);
+        } else {
+             displayToken = tokenStr.replace(/Ġ/g, ' ').replace(/Ċ/g, '\n');
+        }
+
+        // Truncate completion
+        let rawCompletion = item.completion || "";
+        let displayCompletion = rawCompletion;
+        
+        const maxLen = 60;
+        if (displayCompletion.length > maxLen) {
+            displayCompletion = displayCompletion.substring(0, maxLen) + "...";
+        }
+        
+        // Escape HTML
+        const escapeHtml = (text) => {
+            return text
+                .replace(/&/g, "&amp;")
+                .replace(/</g, "&lt;")
+                .replace(/>/g, "&gt;")
+                .replace(/"/g, "&quot;")
+                .replace(/'/g, "&#039;");
+        };
+
+        const fullCompletionEscaped = escapeHtml(rawCompletion);
+        const displayCompletionEscaped = escapeHtml(displayCompletion);
+
+        html += `
+        <tr style="border-bottom: 1px solid #eee;">
+            <td style="padding: 4px 6px;">${item.rank}</td>
+            <td style="padding: 4px 6px; font-family: monospace; background: #fafafa;">${displayToken}</td>
+            <td style="padding: 4px 6px;">${(item.logits || 0).toFixed(2)}</td>
+            <td style="padding: 4px 6px;">
+                <span style="font-size: 0.8em; padding: 2px 4px; border-radius: 4px; background: ${correctBg}; color: ${correctColor}; font-weight: bold;">
+                    ${correctIcon}
+                </span>
+            </td>
+            <td style="padding: 4px 6px; color: #555; font-size: 0.9em; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; max-width: 200px;" title="${fullCompletionEscaped}">${displayCompletionEscaped}</td>
+        </tr>
+        `;
+    });
+    
+    html += '</tbody></table>';
+    container.innerHTML = html;
+}
+
+/* MODEL SELECTION LOGIC */
+async function initModelSelectors() {
+    if(!elModelSeries || !elModelSelect) return;
+
+    // Series change handler
+    elModelSeries.addEventListener('change', async () => {
+        const series = elModelSeries.value;
+        elModelSelect.innerHTML = "<option>Loading...</option>";
+        elModelSelect.disabled = true;
+        
+        try {
+            const res = await fetch(`${API_BASE}/list_hf_models?series=${series}`);
+            const data = await res.json();
+            
+            elModelSelect.innerHTML = "";
+            if(data.models && data.models.length > 0) {
+                 // Sort models alphabetically
+                 data.models.sort();
+                 data.models.forEach(modelName => {
+                    const opt = document.createElement("option");
+                    opt.value = modelName;
+                    opt.textContent = modelName;
+                    elModelSelect.appendChild(opt);
+                });
+                // Select first
+                elModelSelect.disabled = false;
+                elModelSelect.dispatchEvent(new Event('change'));
+            } else {
+                 elModelSelect.innerHTML = "<option>No models found</option>";
+            }
+        } catch(e) {
+            console.error("Failed to list models", e);
+            elModelSelect.innerHTML = "<option>Error loading list</option>";
+        }
+    });
+
+    // Model change handler
+    elModelSelect.addEventListener('change', async () => {
+        const modelName = elModelSelect.value;
+        if(!modelName || modelName.includes("Loading")) return;
+        
+        // Update Path Input
+        elModelPath.value = modelName;
+
+        // Update Revisions
+        elModelRevision.innerHTML = '<option>Loading...</option>';
+        elModelRevision.disabled = true;
+        
+        try {
+            const res = await fetch(`${API_BASE}/list_model_revisions?model_id=${encodeURIComponent(modelName)}`);
+            const data = await res.json();
+            
+            elModelRevision.innerHTML = '<option value="">Latest (Default)</option>';
+            
+            const allRevs = [...(data.branches || []), ...(data.tags || [])];
+            
+            if(allRevs.length > 0) {
+                // Sort revisions alphabetically
+                allRevs.sort();
+                allRevs.forEach(rev => {
+                    if (rev === 'main') return; // Skip main as it is usually default
+                    const opt = document.createElement("option");
+                    opt.value = rev;
+                    opt.textContent = rev;
+                    elModelRevision.appendChild(opt);
+                });
+            }
+            elModelRevision.disabled = false;
+        } catch (e) {
+            console.error("Failed to list revisions", e);
+            elModelRevision.innerHTML = '<option value="">Latest (Default)</option>';
+            elModelRevision.disabled = false;
+        }
+    });
+
+    // Trigger initial population
+    elModelSeries.dispatchEvent(new Event('change'));
+}
+console.log("Initializing model selectors...");
+
+function saveCircuitGraph() {
+    if (!elCanvas) return;
+    
+    // Create a temporary link
+    const link = document.createElement('a');
+    
+    // Generate filename with timestamp
+    const date = new Date();
+    const timestamp = date.toISOString().replace(/[:.]/g, '-');
+    link.download = `circuit_graph_${timestamp}.png`;
+    
+    // Convert canvas to blob/dataURL
+    // High quality PNG
+    const dataUrl = elCanvas.toDataURL('image/png', 1.0);
+    
+    link.href = dataUrl;
+    document.body.appendChild(link);
+    link.click();
+    document.body.removeChild(link);
+}
+
+function saveCircuitGraphPDF() {
+    if (!elCanvas) return;
+    
+    try {
+        const { jsPDF } = window.jspdf;
+        
+        // Create temp canvas to flatten background
+        const tempCanvas = document.createElement('canvas');
+        tempCanvas.width = elCanvas.width;
+        tempCanvas.height = elCanvas.height;
+        const tCtx = tempCanvas.getContext('2d');
+        
+        // Fill white
+        tCtx.fillStyle = '#ffffff';
+        tCtx.fillRect(0, 0, tempCanvas.width, tempCanvas.height);
+        
+        // Draw original
+        tCtx.drawImage(elCanvas, 0, 0);
+        
+        // Calculate PDF size
+        // Orientation based on aspect ratio
+        const orientation = (elCanvas.width > elCanvas.height) ? 'l' : 'p';
+        
+        const doc = new jsPDF({
+            orientation: orientation,
+            unit: 'px',
+            format: [elCanvas.width, elCanvas.height] // Custom size matching canvas
+        });
+        
+        const imgData = tempCanvas.toDataURL('image/jpeg', 1.0);
+        
+        doc.addImage(imgData, 'JPEG', 0, 0, elCanvas.width, elCanvas.height);
+        
+        const date = new Date();
+        const timestamp = date.toISOString().replace(/[:.]/g, '-');
+        doc.save(`circuit_graph_${timestamp}.pdf`);
+    } catch (e) {
+        console.error("PDF generation failed:", e);
+        alert("PDF generation failed. Ensure jsPDF is loaded.");
+    }
+}
+
+async function saveAttributionMapPNG() {
+    if (!elInputAttributionDisplay) return;
+    try {
+        const canvas = await html2canvas(elInputAttributionDisplay, {
+            backgroundColor: '#ffffff'
+        });
+        
+        const link = document.createElement('a');
+        const date = new Date();
+        const timestamp = date.toISOString().replace(/[:.]/g, '-');
+        link.download = `attribution_map_${timestamp}.png`;
+        link.href = canvas.toDataURL('image/png');
+        document.body.appendChild(link);
+        link.click();
+        document.body.removeChild(link);
+    } catch (e) {
+        console.error("PNG save failed", e);
+        alert("Failed to save PNG");
+    }
+}
+
+async function saveAttributionMapPDF() {
+    if (!elInputAttributionDisplay) return;
+    try {
+        const { jsPDF } = window.jspdf;
+        // Capture the full DOM node. For long content, we must paginate into multiple PDF pages
+        // instead of trying to create a single ultra-tall page (many viewers/printers will clamp).
+        const canvas = await html2canvas(elInputAttributionDisplay, {
+            backgroundColor: '#ffffff',
+            scale: 2 // Better quality (also increases memory usage)
+        });
+    const imgWidth = canvas.width;
+    const imgHeight = canvas.height;
+        // Create an A4 portrait PDF in px units.
+        // In jsPDF, with unit='px', 1px ~ 1/96 inch.
+        // So A4 is ~ 794 x 1123 px.
+        const doc = new jsPDF({
+            orientation: 'p',
+            unit: 'px',
+            format: 'a4'
+        });
+        const pageWidth = doc.internal.pageSize.getWidth();
+        const pageHeight = doc.internal.pageSize.getHeight();
+        // We scale the capture to fit page width, then slice the source canvas into page-sized chunks.
+        // This is more robust than relying on negative-y offsets (some jsPDF builds/viewers don't clip as expected).
+        const renderWidth = pageWidth;
+        const scale = renderWidth / imgWidth; // PDF px per canvas px
+        // How many source (canvas) pixels fit into one PDF page height after scaling?
+        const pageHeightInCanvasPx = pageHeight / scale;
+        // Basic debug info to help diagnose issues in-browser.
+        const estimatedPages = Math.max(1, Math.ceil(imgHeight / pageHeightInCanvasPx));
+        console.log('[saveAttributionMapPDF] canvas:', { width: imgWidth, height: imgHeight, scale });
+        console.log('[saveAttributionMapPDF] page:', { pageWidth, pageHeight, pageHeightInCanvasPx, estimatedPages });
+        // Re-use a single offscreen canvas for each page slice to save memory.
+        const pageCanvas = document.createElement('canvas');
+        pageCanvas.width = imgWidth;
+        pageCanvas.height = Math.ceil(pageHeightInCanvasPx);
+        const pageCtx = pageCanvas.getContext('2d');
+        let pageIndex = 0;
+        for (let sy = 0; sy < imgHeight; sy += pageHeightInCanvasPx) {
+            const sliceHeight = Math.min(pageHeightInCanvasPx, imgHeight - sy);
+            // Resize pageCanvas height for the last slice to avoid stretching.
+            if (pageCanvas.height !== Math.ceil(sliceHeight)) {
+                pageCanvas.height = Math.ceil(sliceHeight);
+            }
+            // White background
+            pageCtx.fillStyle = '#ffffff';
+            pageCtx.fillRect(0, 0, pageCanvas.width, pageCanvas.height);
+            // Draw slice
+            pageCtx.drawImage(
+                canvas,
+                0,
+                sy,
+                imgWidth,
+                sliceHeight,
+                0,
+                0,
+                imgWidth,
+                sliceHeight
+            );
+            const imgData = pageCanvas.toDataURL('image/jpeg', 0.95);
+            if (pageIndex > 0) doc.addPage();
+            const sliceRenderHeight = sliceHeight * scale;
+            doc.addImage(imgData, 'JPEG', 0, 0, renderWidth, sliceRenderHeight);
+            pageIndex++;
+        }
+        const date = new Date();
+        const timestamp = date.toISOString().replace(/[:.]/g, '-');
+        doc.save(`attribution_map_${timestamp}.pdf`);
+    } catch (e) {
+        console.error("PDF save failed", e);
+        alert("Failed to save PDF. Ensure html2canvas/jspdf loaded.");
+    }
+}
+
+// async function saveAttributionMapPDF() {
+//     if (!elInputAttributionDisplay) return;
+//     try {
+//         const { jsPDF } = window.jspdf;
+//         const canvas = await html2canvas(elInputAttributionDisplay, {
+//              backgroundColor: '#ffffff',
+//              scale: 2 // Better quality
+//         });
+        
+//         const imgData = canvas.toDataURL('image/jpeg', 1.0);
+        
+//         // Calculate PDF dims to match canvas (scaled back to points if needed, or just px)
+//         // Unit 'px' in jsPDF usually corresponds to 1/96 inch, same as canvas pixel?
+//         // Let's match the canvas pixel dimensions exactly for custom format.
+        
+//         const imgWidth = canvas.width; 
+//         const imgHeight = canvas.height;
+        
+//         const orientation = (imgWidth > imgHeight) ? 'l' : 'p';
+        
+//         const doc = new jsPDF({
+//             orientation: orientation,
+//             unit: 'px',
+//             format: [imgWidth, imgHeight]
+//         });
+        
+//         doc.addImage(imgData, 'JPEG', 0, 0, imgWidth, imgHeight);
+        
+//         const date = new Date();
+//         const timestamp = date.toISOString().replace(/[:.]/g, '-');
+//         doc.save(`attribution_map_${timestamp}.pdf`);
+//     } catch (e) {
+//         console.error("PDF save failed", e);
+//         alert("Failed to save PDF. Ensure html2canvas/jspdf loaded.");
+//     }
+// }
+
+
+initModelSelectors();

requirements.txt

@@ -0,0 +1,14 @@
+fastapi>=0.104.0
+uvicorn>=0.24.0
+pydantic>=2.0.0
+torch>=2.0.0
+transformers>=4.40.0
+numpy>=1.24.0
+scipy>=1.10.0
+pandas>=2.0.0
+networkx>=3.0
+huggingface_hub>=0.20.0
+jinja2>=3.1.0
+pyyaml>=6.0
+openai>=1.0.0
+lxt>=0.3.0