#!/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()