build-small-hackathon/trace-field-notes

1"""Local small-model assistance for Trace Field Notes on Hugging Face ZeroGPU.2 3The analysis models run on the Space GPU through ``transformers``. Heavy imports4(``torch``, ``transformers``) are loaded lazily inside the generator so that the5deterministic analyzer, the test suite, and local development keep working6without GPU dependencies installed. If a model cannot be loaded or its output is7not valid JSON, :func:`analyzer.analyze_trace_file` falls back to the8deterministic codebook and records the reason in the model notes.9"""10 11from __future__ import annotations12 13import json14import re15import time16from collections.abc import Mapping17from dataclasses import dataclass18from typing import Any, Callable19 20from profiling import get_logger21from schemas import (22    APPRAISALS,23    DETOUR_TYPES,24    DIFFICULTY_TYPES,25    OUTCOME_CLAIMS,26    RECOVERY_PATTERNS,27    RESOLUTION_MODES,28)29 30logger = get_logger()31 32 33PRIMARY_MODEL_ID = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"34QUICK_MODEL_ID = "openbmb/MiniCPM5-1B"35MODEL_MAX_NEW_TOKENS = 819236 37MODEL_CHOICES = {38    "minicpm": {39        "label": "MiniCPM5 1B — quick analysis",40        "model_id": QUICK_MODEL_ID,41    },42    "nemotron": {43        "label": "NVIDIA Nemotron 3 Nano 30B-A3B — deeper analysis",44        "model_id": PRIMARY_MODEL_ID,45    },46    "deterministic": {47        "label": "Rule-based — instant, no model",48        "model_id": None,49    },50}51 52# (messages, *, model_id, max_new_tokens) -> raw model text.53GenerateFn = Callable[..., str]54 55_MODEL_CACHE: dict[str, Any] = {}56 57 58@dataclass(slots=True)59class ModelAnalysisResult:60    model_id: str61    analysis: dict[str, Any]62    note: str63 64 65def model_id_for_engine(engine: str) -> str | None:66    choice = MODEL_CHOICES.get(engine)67    if not choice:68        return None69    model_id = choice["model_id"]70    return str(model_id) if model_id else None71 72 73def resolve_device(device: str | None = None) -> str:74    """Pick the compute device: explicit override, else cuda -> mps -> cpu."""75 76    if device:77        return device78    import torch79 80    if torch.cuda.is_available():81        return "cuda"82    mps = getattr(torch.backends, "mps", None)83    if mps is not None and mps.is_available():84        return "mps"85    return "cpu"86 87 88def run_model_analysis(89    *,90    engine: str,91    numbered_narrative: str,92    agent_type: str = "unknown",93    codebook_hint: str = "",94    generate: GenerateFn | None = None,95    device: str | None = None,96) -> ModelAnalysisResult:97    """Run the selected model as the primary analyst and return a field report.98 99    The model identifies and classifies the difficulty episodes and writes the100    session verdict directly from the visible narrative; the deterministic codebook101    is only a fallback (used by the caller if this raises). ``device`` forces the102    compute device for the default local generator; an injected ``generate`` is103    used as-is.104    """105 106    model_id = model_id_for_engine(engine)107    if not model_id:108        raise ValueError(f"No model is configured for analysis engine {engine!r}.")109 110    prompt = build_analysis_prompt(111        numbered_narrative, agent_type=agent_type, codebook_hint=codebook_hint112    )113    messages = [114        {115            "role": "system",116            "content": (117                "You are an expert analyst of coding-agent session traces. "118                "Judge only the visible narrative; never invent hidden reasoning. "119                "Return one JSON object and nothing else."120            ),121        },122        {"role": "user", "content": prompt},123    ]124 125    started = time.perf_counter()126    if generate is not None:127        content = generate(messages, model_id=model_id, max_new_tokens=MODEL_MAX_NEW_TOKENS)128        device_label = "injected"129    else:130        device_label = resolve_device(device)131        content = _local_generator(132            messages,133            model_id=model_id,134            max_new_tokens=MODEL_MAX_NEW_TOKENS,135            device=device_label,136        )137    logger.info(138        "model analysis: %s on %s in %.2fs (%d chars in)",139        model_id,140        device_label,141        time.perf_counter() - started,142        len(numbered_narrative),143    )144    analysis = parse_analysis_json(content)145    return ModelAnalysisResult(146        model_id=model_id,147        analysis=analysis,148        note=f"Analysis produced by {model_id}.",149    )150 151 152def _local_generator(153    messages: list[dict[str, str]],154    *,155    model_id: str,156    max_new_tokens: int,157    device: str | None = None,158) -> str:159    """Generate text with a locally loaded model on the chosen device.160 161    Imported lazily: ``torch`` only needs to exist on the GPU Space (or a local162    machine running the model), never for the deterministic path, tests, or163    light local development.164    """165 166    import torch167 168    tokenizer, model = _load_model(model_id, device=device)169    chat_inputs = tokenizer.apply_chat_template(170        messages,171        add_generation_prompt=True,172        return_tensors="pt",173        **_chat_template_kwargs(model_id),174    )175    generation_inputs, prompt_token_count = _prepare_generation_inputs(176        chat_inputs,177        device=model.device,178    )179    with torch.no_grad():180        generated = model.generate(181            **generation_inputs,182            max_new_tokens=max_new_tokens,183            do_sample=False,184        )185    completion = generated[0][prompt_token_count:]186    return tokenizer.decode(completion, skip_special_tokens=True)187 188 189def _prepare_generation_inputs(chat_inputs: Any, *, device: Any) -> tuple[dict[str, Any], int]:190    """Move tokenizer output to device and return kwargs plus prompt length.191 192    ``apply_chat_template`` may return either a tensor-like object or a193    ``BatchEncoding``/mapping depending on the tokenizer. ``generate`` accepts194    tensor input through the ``inputs=`` keyword and mapping input through195    expanded kwargs such as ``input_ids`` and ``attention_mask``.196    """197 198    moved = _move_to_device(chat_inputs, device)199    if isinstance(moved, Mapping):200        generation_inputs = {201            key: _move_to_device(value, device)202            for key, value in moved.items()203        }204        input_ids = generation_inputs.get("input_ids")205        if input_ids is None or not hasattr(input_ids, "shape"):206            raise ValueError("Tokenizer output did not include tensor-shaped input_ids.")207        return generation_inputs, int(input_ids.shape[-1])208 209    if not hasattr(moved, "shape"):210        raise ValueError("Tokenizer output was neither a tensor nor a mapping.")211    return {"inputs": moved}, int(moved.shape[-1])212 213 214def _move_to_device(value: Any, device: Any) -> Any:215    if hasattr(value, "to"):216        return value.to(device)217    return value218 219 220def _chat_template_kwargs(model_id: str) -> dict[str, Any]:221    """Model-specific chat-template controls."""222 223    if model_id.startswith("openbmb/"):224        # MiniCPM5 supports hybrid reasoning; the quick engine keeps thinking225        # off for fast, reliably parseable JSON memos.226        return {"enable_thinking": False}227    return {}228 229 230def _load_model(model_id: str, device: str | None = None) -> Any:231    """Lazily load and cache a (tokenizer, model) pair on the chosen device.232 233    The cache keeps weights resident across requests so only the first call per234    (model, device) pays the load cost. ZeroGPU exposes CUDA inside the235    ``@spaces.GPU`` context; CPU/MPS support lets the app run off-Space (e.g. for236    users without GPU quota, or local development).237    """238 239    import torch240 241    resolved = resolve_device(device)242    cache_key = f"{model_id}@{resolved}"243    cached = _MODEL_CACHE.get(cache_key)244    if cached is not None:245        return cached246 247    from transformers import AutoModelForCausalLM, AutoTokenizer248 249    started = time.perf_counter()250    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)251    if resolved == "cuda":252        # The ZeroGPU Space path: load straight onto the GPU in bfloat16.253        model = AutoModelForCausalLM.from_pretrained(254            model_id,255            dtype=torch.bfloat16,256            device_map="cuda",257            trust_remote_code=True,258        )259    else:260        # CPU / Apple MPS: fp16 on MPS, fp32 on CPU for numerical stability.261        dtype = torch.float16 if resolved == "mps" else torch.float32262        model = AutoModelForCausalLM.from_pretrained(263            model_id,264            dtype=dtype,265            trust_remote_code=True,266        ).to(resolved)267    model.eval()268    logger.info("loaded %s on %s in %.1fs", model_id, resolved, time.perf_counter() - started)269    _MODEL_CACHE[cache_key] = (tokenizer, model)270    return tokenizer, model271 272 273def _vocab_block(name: str, vocab: dict[str, str]) -> str:274    return f"{name}:\n" + "\n".join(f"- {key}: {meaning}" for key, meaning in vocab.items())275 276 277def build_analysis_prompt(278    numbered_narrative: str, *, agent_type: str = "unknown", codebook_hint: str = ""279) -> str:280    narrative = numbered_narrative[:16000]281    vocab = "\n\n".join(282        [283            _vocab_block("difficulty_type", DIFFICULTY_TYPES),284            _vocab_block("appraisal", APPRAISALS),285            _vocab_block("detour_type", DETOUR_TYPES),286            _vocab_block("resolution_mode", RESOLUTION_MODES),287            _vocab_block("recovery_pattern", RECOVERY_PATTERNS),288            _vocab_block("outcome_claim", OUTCOME_CLAIMS),289        ]290    )291    return f"""Read the agent's visible narrative and produce a structured field report as JSON.292 293Identify the real DIFFICULTY EPISODES — moments where the agent hit a snag, reassessed,294detoured, recovered, or claimed completion. Ignore instructions, skill files, prompts,295or boilerplate the agent merely read or quoted; those are NOT difficulties. Merge296duplicates. Prefer 1-8 substantive episodes; if there is genuinely no difficulty,297return an empty episodes list.298 299Return ONE JSON object (first character {{ and last character }}), no prose, EXACTLY:300{{301  "verdict": {{302    "tone": one of ["stable","iterative","detour","partial","risk","unknown"],303    "headline": "<= 12 words, plain language",304    "detail": "2-4 sentences a developer can act on",305    "honesty": one of ["candid","mixed","overclaimed"]306  }},307  "overall_patterns": {{308    "difficulty_style": "1 sentence", "detour_style": "1 sentence",309    "recovery_style": "1 sentence", "risk_or_caveat": "1 sentence"310  }},311  "episodes": [312    {{313      "start_index": <a message index shown below>,314      "end_index": <a message index shown below>,315      "title": "<= 10 words",316      "initial_intention": "1 sentence", "reported_difficulty": "1-2 sentences",317      "difficulty_type": "<one key below>", "appraisal": "<one key below>",318      "strategy_before": "1 sentence", "strategy_after": "1 sentence",319      "detour_type": "<one key below>", "resolution_mode": "<one key below>",320      "recovery_pattern": "<one key below>", "outcome_claim": "<one key below>",321      "productive_detour": one of ["yes","no","mixed","unknown"],322      "evidence_quotes": ["short verbatim quote", "up to 3"],323      "analyst_memo": "1-3 sentences of real insight, NOT a restatement of the codes"324    }}325  ]326}}327 328Controlled vocabulary (use these keys exactly):329{vocab}330 331Guidance:332- Every field must contain real content drawn from the trace. NEVER output a333  placeholder such as "<= 10 words", "1 sentence", or "<one key below>" literally.334- difficulty_type, appraisal, detour_type, resolution_mode, recovery_pattern, and335  outcome_claim must each be EXACTLY one key from the vocabulary above (lowercase,336  with underscores). If unsure, use "unknown".337- Be accurate, not generous. If the agent ended unresolved or overclaimed, say so in tone/honesty.338- honesty = "overclaimed" when a success claim outruns the visible evidence.339- start_index / end_index must be message indices that appear below.340- Quote the agent's own words; keep the original language of the quote.341- Do not include secrets or long tool dumps.342 343Agent type: {agent_type}344Rule-based pre-scan candidate spans (hints only — keep, drop, merge, or add freely): {codebook_hint or "(none)"}345 346Numbered visible messages:347{narrative}348"""349 350 351def parse_analysis_json(content: str) -> dict[str, Any]:352    """Validate the structural shape of the model's field report (codes coerced later)."""353 354    parsed = _loads_lenient(content)355    episodes = parsed.get("episodes")356    if not isinstance(episodes, list):357        raise ValueError("Model response did not include an 'episodes' list.")358    parsed["episodes"] = [episode for episode in episodes if isinstance(episode, dict)]359    if not isinstance(parsed.get("overall_patterns"), dict):360        parsed["overall_patterns"] = {}361    if not isinstance(parsed.get("verdict"), dict):362        parsed["verdict"] = {}363    return parsed364 365 366def _loads_lenient(content: str) -> dict[str, Any]:367    """Parse JSON from a model that may wrap it in prose or code fences."""368 369    if not isinstance(content, str) or not content.strip():370        raise ValueError("Model response content was empty.")371 372    text = content.strip()373    fence = re.match(r"^```[a-zA-Z0-9]*\s*(.*?)\s*```$", text, re.DOTALL)374    if fence:375        text = fence.group(1).strip()376 377    try:378        parsed: Any = json.loads(text)379    except json.JSONDecodeError:380        candidates = list(_json_object_candidates(text))381        if not candidates:382            raise ValueError("Model response was not valid JSON.")383        parsed = candidates[-1]384 385    if not isinstance(parsed, dict):386        raise ValueError("Model response was not a JSON object.")387    return parsed388 389 390def _json_object_candidates(text: str) -> list[dict[str, Any]]:391    decoder = json.JSONDecoder()392    candidates: list[dict[str, Any]] = []393    cursor = 0394    while True:395        start = text.find("{", cursor)396        if start == -1:397            return candidates398        try:399            parsed, consumed = decoder.raw_decode(text[start:])400        except json.JSONDecodeError:401            cursor = start + 1402            continue403        if isinstance(parsed, dict):404            candidates.append(parsed)405        cursor = start + max(consumed, 1)406