scripts/inference_server.py

#!/usr/bin/env python3
"""Persistent JSONL sidecar for the local BetterWright 2.0 alpha."""

from __future__ import annotations

import argparse
import bisect
import json
import sys
import time
from pathlib import Path

import torch
from transformers import AutoModel, AutoTokenizer

from tree_utils import references, structural_windows

MUST_KEEP = ("[active]", "[selected]", "[checked]", "dialog", "alert", "error", "textbox")


def ancestors(lines: list[str], index: int) -> set[int]:
    keep = set()
    indent = len(lines[index]) - len(lines[index].lstrip())
    for i in range(index - 1, -1, -1):
        candidate = len(lines[i]) - len(lines[i].lstrip())
        if candidate < indent:
            keep.add(i)
            indent = candidate
            if indent == 0:
                break
    return keep


def render_indices(tree: str, indices: set[int]) -> str:
    lines = [line.rstrip() for line in tree.splitlines() if line.strip()]
    indices = {index for index in indices if 0 <= index < len(lines)}
    for i in list(indices):
        indices.update(ancestors(lines, i))
    out = []
    previous = -1
    for i in sorted(indices):
        if i > previous + 1:
            out.append("- text: … irrelevant subtree pruned …")
        out.append(lines[i])
        previous = i
    return "\n".join(out)


def coarse_ref_indices(
    tree: str,
    windows,
    scores: list[float],
    max_ranked_windows: int,
    ref_context_lines: int,
) -> set[int]:
    """Mirror the validated coarse-ref-context policy exactly."""
    lines = [line.rstrip() for line in tree.splitlines() if line.strip()]
    ranked_windows = sorted(
        range(len(windows)), key=lambda index: scores[index], reverse=True
    )[:max_ranked_windows]
    candidate_lines = set()
    for window_index in ranked_windows:
        window = windows[window_index]
        candidate_lines.update(
            range(window.start_line, min(window.end_line, len(lines)))
        )
    chosen = {
        index
        for index in candidate_lines
        if references(lines[index])
        or any(term in lines[index].casefold() for term in MUST_KEEP)
    }
    for index in list(chosen):
        base_indent = len(lines[index]) - len(lines[index].lstrip())
        kept = 0
        for child in range(index + 1, len(lines)):
            child_indent = len(lines[child]) - len(lines[child].lstrip())
            if child_indent <= base_indent:
                break
            if child in candidate_lines:
                chosen.add(child)
                kept += 1
            if kept >= ref_context_lines:
                break
    return chosen


class Server:
    def __init__(self, model_id: str, config_path: Path):
        self.config = json.loads(config_path.read_text())
        device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(
            model_id, trust_remote_code=True, torch_dtype=torch.bfloat16
        ).to(device).eval()

    @torch.inference_mode()
    def prune(self, query: str, tree: str, max_chars: int | None = None) -> dict:
        started = time.perf_counter()
        if not query.strip() or len(tree) < 1800:
            return {"tree": tree, "fallback": True, "reason": "missing-query-or-small-tree", "savings": 0}
        validated_max_chars = int(self.config.get("validated_max_chars", 10_000))
        if max_chars is not None and max_chars < validated_max_chars:
            return {
                "tree": tree,
                "fallback": True,
                "reason": "requested-budget-below-validated-budget",
                "savings": 0,
            }
        window_chars = int(self.config.get("window_chars", 3600))
        strategy = self.config.get("strategy", "token-lines")
        windows = structural_windows(tree, max_chars=window_chars, overlap_lines=4)
        scores = []
        confidences = []
        total_lines = max((window.end_line for window in windows), default=0)
        line_scores = [float("-inf")] * total_lines
        for start in range(0, len(windows), 16):
            part = windows[start : start + 16]
            prefixes = [
                f"[BETTERWRIGHT_TASK]\n{query}\n[ACCESSIBILITY_SUBTREE]\n"
                for _ in part
            ]
            texts = [prefix + window.text for prefix, window in zip(prefixes, part)]
            tokenizer_args = {
                "padding": True,
                "truncation": True,
                "max_length": 2048,
                "return_tensors": "pt",
            }
            if strategy != "coarse-ref-context":
                tokenizer_args["return_offsets_mapping"] = True
            batch = self.tokenizer(texts, **tokenizer_args)
            offsets = batch.pop("offset_mapping", None)
            output = self.model(**batch.to(self.device))
            scores.extend(torch.sigmoid(output.logits).float().cpu().tolist())
            confidences.extend(torch.sigmoid(output.uncertainty_logits).float().cpu().tolist())
            if strategy == "coarse-ref-context":
                continue
            token_scores = torch.sigmoid(output.token_logits).float().cpu()
            for batch_index, window in enumerate(part):
                prefix_chars = len(prefixes[batch_index])
                line_starts = []
                offset = 0
                for line in window.text.splitlines():
                    line_starts.append(offset)
                    offset += len(line) + 1
                for token_index, (token_start, token_end) in enumerate(
                    offsets[batch_index].tolist()
                ):
                    if token_end <= token_start or token_end <= prefix_chars:
                        continue
                    relative = max(token_start, prefix_chars) - prefix_chars
                    if relative >= len(window.text):
                        continue
                    local_line = bisect.bisect_right(line_starts, relative) - 1
                    if local_line < 0:
                        continue
                    global_line = window.start_line + local_line
                    if global_line < len(line_scores):
                        line_scores[global_line] = max(
                            line_scores[global_line],
                            float(token_scores[batch_index, token_index]),
                        )
        if not scores:
            return {"tree": tree, "fallback": True, "reason": "no-windows", "savings": 0}
        threshold = self.config["relevance_threshold"]
        if strategy == "coarse-ref-context":
            chosen = coarse_ref_indices(
                tree,
                windows,
                scores,
                int(self.config["max_ranked_windows"]),
                int(self.config["ref_context_lines"]),
            )
        else:
            lines = [line.rstrip() for line in tree.splitlines() if line.strip()]
            chosen = {
                index
                for index, line in enumerate(lines)
                if any(term in line.casefold() for term in MUST_KEEP)
            }
            ranked = sorted(
                (
                    index
                    for index, score in enumerate(line_scores)
                    if score != float("-inf")
                ),
                key=lambda index: line_scores[index],
                reverse=True,
            )
            chosen.update(ranked[: int(self.config["max_ranked_lines"])])
        peak_i = max(range(len(scores)), key=lambda i: scores[i])
        if (
            not chosen
            or scores[peak_i] < threshold
            or confidences[peak_i] < self.config["confidence_threshold"]
        ):
            return {"tree": tree, "fallback": True, "reason": "low-confidence", "savings": 0}
        pruned = render_indices(tree, chosen)
        savings = 1 - len(pruned) / max(1, len(tree))
        if len(pruned) > validated_max_chars:
            return {
                "tree": tree,
                "fallback": True,
                "reason": "candidate-over-limit",
                "savings": 0,
                "candidate_chars": len(pruned),
                "candidate_savings": savings,
            }
        if savings < 0.08:
            return {
                "tree": tree,
                "fallback": True,
                "reason": "insufficient-benefit",
                "savings": 0,
                "candidate_chars": len(pruned),
                "candidate_savings": savings,
            }
        return {
            "tree": pruned, "fallback": False, "reason": "model", "savings": savings,
            "latency_ms": round((time.perf_counter() - started) * 1000, 2),
            "peak_score": scores[peak_i], "peak_confidence": confidences[peak_i],
        }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="ProCreations/betterwright-encoder-350m")
    parser.add_argument("--config", type=Path, required=True)
    args = parser.parse_args()
    server = Server(args.model, args.config)
    print(json.dumps({"ready": True, "model": args.model}), flush=True)
    for line in sys.stdin:
        request = {}
        try:
            request = json.loads(line)
            raw_max_chars = request.get("max_chars")
            max_chars = int(raw_max_chars) if raw_max_chars is not None else None
            response = {
                "id": request.get("id"),
                **server.prune(
                    str(request.get("query", "")),
                    str(request.get("tree", "")),
                    max_chars=max_chars,
                ),
            }
        except Exception as error:
            response = {"id": request.get("id"), "fallback": True, "error": str(error)}
        print(json.dumps(response, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()