Upload scripts/tree_utils.py with huggingface_hub

scripts/tree_utils.py

@@ -0,0 +1,74 @@
+"""Structural parsing shared by data generation, training, and evaluation."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+from dataclasses import dataclass
+
+REF_RE = re.compile(r"\[ref=([^\]]+)\]")
+STAGEHAND_REF_RE = re.compile(r"^\s*\[([^\]]+)\]")
+WORD_RE = re.compile(r"[a-z0-9]{2,}")
+
+
+@dataclass(frozen=True)
+class Window:
+    id: str
+    text: str
+    refs: tuple[str, ...]
+    start_line: int
+    end_line: int
+
+
+def references(text: str, source_format: str = "betterwright") -> set[str]:
+    if source_format == "stagehand":
+        return {m.group(1) for line in text.splitlines() if (m := STAGEHAND_REF_RE.match(line))}
+    # Auto-detect mixed corpora: Stagehand's outline is `[id] role: name`,
+    # while BetterWright/Playwright use `[ref=id]`. The patterns cannot collide.
+    result = set(REF_RE.findall(text))
+    result.update(
+        m.group(1) for line in text.splitlines() if (m := STAGEHAND_REF_RE.match(line))
+    )
+    return result
+
+
+def words(text: str) -> set[str]:
+    return set(WORD_RE.findall(text.casefold()))
+
+
+def structural_windows(text: str, *, max_chars: int = 3600, overlap_lines: int = 4) -> list[Window]:
+    """Create overlapping, indentation-aware windows without cutting a line."""
+    lines = [line.rstrip() for line in text.splitlines() if line.strip()]
+    if not lines:
+        return []
+    windows: list[Window] = []
+    start = 0
+    while start < len(lines):
+        size = 0
+        end = start
+        while end < len(lines) and (size + len(lines[end]) + 1 <= max_chars or end == start):
+            size += len(lines[end]) + 1
+            end += 1
+        # Prefer a structural boundary near the limit.
+        if end < len(lines):
+            floor = max(start + 1, end - 12)
+            candidates = [i for i in range(floor, end) if len(lines[i]) - len(lines[i].lstrip()) <= 2]
+            if candidates:
+                end = candidates[-1]
+        body = "\n".join(lines[start:end])
+        refs = tuple(sorted(references(body)))
+        digest = hashlib.sha1(f"{start}\0{end}\0{body}".encode()).hexdigest()[:16]
+        windows.append(Window(digest, body, refs, start, end))
+        if end >= len(lines):
+            break
+        start = max(start + 1, end - overlap_lines)
+    return windows
+
+
+def lexical_score(query: str, text: str) -> float:
+    q = words(query)
+    if not q:
+        return 0.0
+    t = words(text)
+    overlap = len(q & t)
+    return overlap / (len(q) ** 0.5 * max(1, len(t)) ** 0.25)