Upload scripts/inference_server.py with huggingface_hub

scripts/inference_server.py

@@ -0,0 +1,254 @@
+#!/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()