app.py

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):
    _nan_warning_logged = False

    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)
            if not np.isfinite(f):
                if not NumpyEncoder._nan_warning_logged:
                    logger.warning(f"NumpyEncoder: Converting non-finite value ({f}) to 0.0. "
                                   "This may indicate numerical instability in LRP computation.")
                    NumpyEncoder._nan_warning_logged = True
                return 0.0
            return f
        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
    target_token_id: Optional[int] = None  # Target token ID for backprop (top-1 if None)

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:
        # Prompt comes properly decoded from JSON - no unescaping needed.
        # unescape_string would corrupt special tokens like <|im_start|>
        # and mishandle non-ASCII characters via unicode_escape.
        prompt = 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=False)

        # 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:
        # Prompt/completion come properly decoded from JSON - no unescaping needed.
        prompt = request.prompt
        completion = request.completion
        ground_truth = 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:
        # Invalidate circuit cache when a new forward pass is run
        CACHED_CONNECTION_DATA["config_hash"] = None
        CACHED_CONNECTION_DATA["data"] = None

        # Prompt comes properly decoded from JSON - no unescaping needed.
        prompt = 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")

# Data directory for trace files
DATA_DIR = os.path.join(PROJECT_ROOT, "data")

@app.get("/api/datasets")
async def get_datasets():
    """Scan data/ directory for available datasets (subdirectories)."""
    datasets = []
    if os.path.isdir(DATA_DIR):
        for name in sorted(os.listdir(DATA_DIR)):
            full_path = os.path.join(DATA_DIR, name)
            if os.path.isdir(full_path):
                datasets.append(name)
    return {"datasets": datasets}

@app.get("/api/traces/{dataset}")
async def get_traces(dataset: str):
    """List trace files (JSON) in a dataset directory."""
    dataset_dir = os.path.join(DATA_DIR, dataset)
    if not os.path.isdir(dataset_dir):
        raise HTTPException(status_code=404, detail=f"Dataset '{dataset}' not found")

    traces = []
    for name in sorted(os.listdir(dataset_dir)):
        if name.endswith(".json"):
            traces.append(name)
    return {"traces": traces}

@app.get("/api/trace_details/{dataset}/{trace_file}")
async def get_trace_details(dataset: str, trace_file: str):
    """Load and return trace file details."""
    file_path = os.path.join(DATA_DIR, dataset, trace_file)
    if not os.path.isfile(file_path):
        raise HTTPException(status_code=404, detail=f"Trace file '{trace_file}' not found in dataset '{dataset}'")

    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            trace_data = json.load(f)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to parse trace file: {str(e)}")

    # Extract fields for the frontend
    metadata = trace_data.get("metadata", {})
    result = {
        "model_path": metadata.get("model", ""),
        "dtype": str(metadata.get("dtype", "float16")).replace("torch.", ""),
        "quantization": False,
        "prompt": trace_data.get("prompt", ""),
        "raw_prompt": trace_data.get("prompt", ""),
        "completion": trace_data.get("completion", ""),
        "ground_truth": trace_data.get("ground_truth", ""),
        "eval_result": trace_data.get("eval_result", None),
        "topk_token_explore": trace_data.get("topk_token_explore", []),
    }

    # Check for other model candidates (e.g., 4b variants)
    other_candidates = {}
    for key in trace_data:
        if key.startswith("topk_token_explore_") and key != "topk_token_explore":
            suffix = key.replace("topk_token_explore_", "")
            other_candidates[suffix] = trace_data[key]
    if other_candidates:
        result["other_candidates"] = other_candidates

    return result

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)