refactor: move eval/ and data_adapters/ to tools/

tools/data_adapters/__init__.py

@@ -0,0 +1 @@
+"""Data adapters: convert external datasets to the project's traces.jsonl format."""

tools/data_adapters/__main__.py

@@ -0,0 +1,13 @@
+"""CLI dispatch: `python -m ai_mouse.data_adapters.balabit ...`
+
+Note: This file makes `python -m ai_mouse.data_adapters` invokable but for
+clarity prefer the explicit form `python -m ai_mouse.data_adapters.balabit`.
+"""
+from __future__ import annotations
+
+import sys
+
+from tools.data_adapters.balabit import main
+
+if __name__ == "__main__":
+    sys.exit(main())

tools/data_adapters/balabit.py

@@ -0,0 +1,326 @@
+"""Adapter for the Balabit Mouse Dynamics Challenge dataset.
+
+Source: https://github.com/balabit/Mouse-Dynamics-Challenge
+
+Each session is a CSV file with columns:
+    record timestamp, client timestamp, button, state, x, y
+
+Where:
+    state ∈ {Move, Pressed, Released, Drag, Scroll}
+    button ∈ {NoButton, Left, Right, Wheel}
+
+We extract "click-anchored" trajectory segments: each Pressed event
+defines a target, and the W ms of Move events preceding it form one
+training trace.
+"""
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MouseEvent:
+    """A single mouse event from a Balabit CSV row."""
+
+    t_ms: int        # client timestamp in milliseconds (relative to session start)
+    button: str      # "NoButton", "Left", "Right", "Wheel"
+    state: str       # "Move", "Pressed", "Released", "Drag", "Scroll"
+    x: int
+    y: int
+
+
+@dataclass
+class Segment:
+    """A click-anchored trajectory segment ready to be written to JSONL."""
+
+    events: list[MouseEvent]   # only Move events, sorted by t_ms ascending
+    click_x: int               # the Pressed event's x coordinate
+    click_y: int               # the Pressed event's y coordinate
+    click_t_ms: int            # the Pressed event's timestamp
+    session_id: str            # e.g. "user7_session_42"
+
+
+def parse_session_csv(path: Path) -> list[MouseEvent]:
+    """Parse a Balabit session CSV file into MouseEvent objects.
+
+    Malformed rows are logged and skipped (not raised).
+    Client timestamps (seconds, float) are converted to int milliseconds.
+
+    Args:
+        path: Path to a Balabit session CSV file.
+
+    Returns:
+        List of MouseEvent in original order. Empty list if file is empty.
+    """
+    import csv as csv_module
+
+    events: list[MouseEvent] = []
+    with path.open("r", encoding="utf-8", newline="") as f:
+        reader = csv_module.DictReader(f)
+        for row_idx, row in enumerate(reader, 2):  # 1-based, header is line 1
+            try:
+                client_ts = float(row["client timestamp"])
+                t_ms = int(round(client_ts * 1000))
+                button = row["button"].strip()
+                state = row["state"].strip()
+                x = int(row["x"])
+                y = int(row["y"])
+            except (KeyError, ValueError, TypeError) as exc:
+                logger.debug("Skipping malformed row %d in %s: %s", row_idx, path.name, exc)
+                continue
+            events.append(MouseEvent(t_ms=t_ms, button=button, state=state, x=x, y=y))
+    return events
+
+
+def segment_by_clicks(
+    events: list[MouseEvent],
+    window_ms: int,
+    session_id: str,
+) -> list[Segment]:
+    """Extract click-anchored segments from a session.
+
+    For each Left-button Pressed event, collect all Move events within
+    [click_t - window_ms, click_t) into one segment.
+
+    Args:
+        events: Full session events (any state, any order is OK but typically sorted).
+        window_ms: How far back to look before each click.
+        session_id: String tag attached to every segment for debugging.
+
+    Returns:
+        List of Segment, one per Left Pressed event that has at least one preceding Move.
+    """
+    segments: list[Segment] = []
+    for ev in events:
+        if ev.state != "Pressed" or ev.button != "Left":
+            continue
+        click_t = ev.t_ms
+        window_start = click_t - window_ms
+        moves = [
+            m for m in events
+            if m.state == "Move" and window_start <= m.t_ms < click_t
+        ]
+        if not moves:
+            continue
+        moves.sort(key=lambda m: m.t_ms)
+        segments.append(Segment(
+            events=moves,
+            click_x=ev.x,
+            click_y=ev.y,
+            click_t_ms=click_t,
+            session_id=session_id,
+        ))
+    return segments
+
+
+def filter_segments(
+    segments: list[Segment],
+    min_events: int,
+    min_dist: int,
+    max_span_ms: int,
+    max_gap_ms: int,
+    coord_max: int = 5000,
+) -> list[Segment]:
+    """Drop segments that fail any quality check.
+
+    A segment is dropped if any of these are true:
+    - len(events) < min_events
+    - Euclidean dist(events[0], (click_x, click_y)) < min_dist
+    - events[-1].t_ms - events[0].t_ms > max_span_ms
+    - any adjacent Move pair has dt > max_gap_ms (sampling drop-out)
+    - any coord (start/end/click) outside [0, coord_max]
+    - total arc length < min_dist (high-frequency jitter only)
+
+    Args:
+        segments: Candidate segments.
+        min_events: Minimum number of Move events.
+        min_dist: Minimum start→click pixel distance AND minimum total arc length.
+        max_span_ms: Maximum time span of the segment (events[-1] - events[0]).
+        max_gap_ms: Maximum allowed gap between adjacent Move events.
+        coord_max: Maximum allowed pixel coordinate value (5000 catches multi-monitor anomalies).
+
+    Returns:
+        Filtered list, original order preserved.
+    """
+    import math
+
+    keep: list[Segment] = []
+    for seg in segments:
+        if len(seg.events) < min_events:
+            continue
+
+        sx, sy = seg.events[0].x, seg.events[0].y
+        ex, ey = seg.click_x, seg.click_y
+
+        # Coord range check
+        if any(c < 0 or c > coord_max for c in (sx, sy, ex, ey)):
+            continue
+
+        # Endpoint distance
+        dist = math.hypot(ex - sx, ey - sy)
+        if dist < min_dist:
+            continue
+
+        # Time span
+        span = seg.events[-1].t_ms - seg.events[0].t_ms
+        if span > max_span_ms:
+            continue
+
+        # Gap check + total arc length
+        total_arc = 0.0
+        bad_gap = False
+        for i in range(1, len(seg.events)):
+            dt = seg.events[i].t_ms - seg.events[i - 1].t_ms
+            if dt > max_gap_ms:
+                bad_gap = True
+                break
+            dx = seg.events[i].x - seg.events[i - 1].x
+            dy = seg.events[i].y - seg.events[i - 1].y
+            total_arc += math.hypot(dx, dy)
+        if bad_gap:
+            continue
+        if total_arc < min_dist:
+            continue
+
+        keep.append(seg)
+
+    return keep
+
+
+def process_session(
+    csv_path: Path,
+    output_jsonl: Path,
+    config,
+) -> int:
+    """Convert one Balabit session CSV to JSONL traces, append to output.
+
+    Args:
+        csv_path: Path to a Balabit session CSV.
+        output_jsonl: Output JSONL file (will be appended to).
+        config: BalabitAdapterConfig with window_ms / min_dist / etc.
+
+    Returns:
+        Number of valid segments written.
+    """
+    import math
+
+    session_id = csv_path.stem  # e.g. "session_42"
+    events = parse_session_csv(csv_path)
+    if not events:
+        return 0
+
+    raw_segments = segment_by_clicks(
+        events, window_ms=config.window_ms, session_id=session_id
+    )
+    valid_segments = filter_segments(
+        raw_segments,
+        min_events=config.min_events,
+        min_dist=config.min_dist,
+        max_span_ms=config.max_span_ms,
+        max_gap_ms=config.max_gap_ms,
+    )
+
+    if not valid_segments:
+        return 0
+
+    output_jsonl.parent.mkdir(parents=True, exist_ok=True)
+    n_written = 0
+    with output_jsonl.open("a", encoding="utf-8") as f:
+        for seg in valid_segments:
+            sx, sy = seg.events[0].x, seg.events[0].y
+            ex, ey = seg.click_x, seg.click_y
+            dist = math.hypot(ex - sx, ey - sy)
+            angle = math.degrees(math.atan2(ey - sy, ex - sx))
+            t0 = seg.events[0].t_ms
+            record = {
+                "meta": {
+                    "start": [sx, sy],
+                    "end": [ex, ey],
+                    "dist": int(round(dist)),
+                    "angle": round(angle, 1),
+                    "source": "balabit",
+                    "session_id": seg.session_id,
+                },
+                "events": [
+                    {"type": "move", "x": e.x, "y": e.y, "t": e.t_ms - t0}
+                    for e in seg.events
+                ],
+            }
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+            n_written += 1
+    return n_written
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point: convert a directory of Balabit sessions to one JSONL file."""
+    import argparse
+
+    from tools.config import BalabitAdapterConfig
+
+    parser = argparse.ArgumentParser(description="Convert Balabit dataset to traces.jsonl format")
+    parser.add_argument(
+        "--input", type=Path, required=True,
+        help="Directory containing Balabit session CSV files (recursive)",
+    )
+    parser.add_argument(
+        "--output", type=Path, default=Path("data/pretrain_traces.jsonl"),
+        help="Output JSONL path (default: data/pretrain_traces.jsonl)",
+    )
+    parser.add_argument("--window-ms", type=int, default=1200)
+    parser.add_argument("--min-dist", type=int, default=50)
+    parser.add_argument("--min-events", type=int, default=5)
+    parser.add_argument("--max-span-ms", type=int, default=5000)
+    parser.add_argument("--max-gap-ms", type=int, default=200)
+    parser.add_argument(
+        "--overwrite", action="store_true",
+        help="Truncate output file before writing (default: append)",
+    )
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    if not args.input.is_dir():
+        logger.error("Input is not a directory: %s", args.input)
+        return 2
+
+    config = BalabitAdapterConfig(
+        window_ms=args.window_ms,
+        min_dist=args.min_dist,
+        min_events=args.min_events,
+        max_span_ms=args.max_span_ms,
+        max_gap_ms=args.max_gap_ms,
+    )
+
+    if args.overwrite and args.output.exists():
+        args.output.unlink()
+
+    # Recursively walk the input directory looking for session files.
+    # Balabit files have no extension; we accept any regular file.
+    csv_files = sorted([p for p in args.input.rglob("*") if p.is_file() and not p.name.startswith(".")])
+    if not csv_files:
+        logger.error("No session files found under %s", args.input)
+        return 2
+
+    total = 0
+    for i, csv_path in enumerate(csv_files, 1):
+        try:
+            n = process_session(csv_path, args.output, config)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("Skipping %s due to error: %s", csv_path.name, exc)
+            continue
+        total += n
+        if i % 10 == 0 or i == len(csv_files):
+            logger.info("Processed %d/%d sessions, %d segments so far", i, len(csv_files), total)
+
+    logger.info("Done. Wrote %d segments to %s", total, args.output)
+    return 0
+
+
+if __name__ == "__main__":
+    import sys
+    sys.exit(main())