feat(adapter): implement process_session and CLI
This commit is contained in:
13
ai_mouse/data_adapters/__main__.py
Normal file
13
ai_mouse/data_adapters/__main__.py
Normal file
@@ -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 ai_mouse.data_adapters.balabit import main
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -197,5 +197,130 @@ def process_session(
|
||||
output_jsonl: Path,
|
||||
config,
|
||||
) -> int:
|
||||
"""Stub — implemented in Task 6."""
|
||||
raise NotImplementedError
|
||||
"""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 ai_mouse.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())
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
"""Tests for Balabit Mouse Dynamics Challenge data adapter."""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@@ -261,3 +262,75 @@ class TestFilterSegments:
|
||||
result = filter_segments([seg], min_events=5, min_dist=50,
|
||||
max_span_ms=5000, max_gap_ms=200)
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestProcessSession:
|
||||
def test_writes_jsonl_in_expected_format(self, tmp_path):
|
||||
from ai_mouse.config import BalabitAdapterConfig
|
||||
from ai_mouse.data_adapters.balabit import process_session
|
||||
|
||||
# Construct a Balabit-format CSV with one valid segment
|
||||
csv_path = tmp_path / "user_session_42"
|
||||
rows = []
|
||||
# 10 Move events going from (100,100) to (500,200) over 500ms
|
||||
for i in range(10):
|
||||
t = i * 0.05
|
||||
x = 100 + i * 40
|
||||
y = 100 + i * 10
|
||||
rows.append(f"0,{t:.3f},NoButton,Move,{x},{y}")
|
||||
rows.append(f"0,0.550,Left,Pressed,510,210")
|
||||
_write_csv(csv_path, rows)
|
||||
|
||||
out = tmp_path / "out.jsonl"
|
||||
config = BalabitAdapterConfig(
|
||||
window_ms=1200, min_dist=50, min_events=5,
|
||||
max_span_ms=5000, max_gap_ms=200,
|
||||
)
|
||||
n = process_session(csv_path, out, config)
|
||||
assert n == 1
|
||||
assert out.exists()
|
||||
|
||||
line = out.read_text(encoding="utf-8").strip()
|
||||
record = json.loads(line)
|
||||
assert record["meta"]["start"] == [100, 100]
|
||||
assert record["meta"]["end"] == [510, 210]
|
||||
assert record["meta"]["dist"] > 0
|
||||
assert record["meta"]["source"] == "balabit"
|
||||
assert record["meta"]["session_id"] == "user_session_42"
|
||||
# Events: only move events, no down/up
|
||||
assert all(e["type"] == "move" for e in record["events"])
|
||||
# Timestamps relative (start at 0)
|
||||
assert record["events"][0]["t"] == 0
|
||||
|
||||
def test_returns_zero_for_session_with_no_valid_segments(self, tmp_path):
|
||||
from ai_mouse.config import BalabitAdapterConfig
|
||||
from ai_mouse.data_adapters.balabit import process_session
|
||||
|
||||
csv_path = tmp_path / "empty_session"
|
||||
_write_csv(csv_path, ["0,0.000,NoButton,Move,100,100"]) # no clicks
|
||||
|
||||
out = tmp_path / "out.jsonl"
|
||||
config = BalabitAdapterConfig()
|
||||
n = process_session(csv_path, out, config)
|
||||
assert n == 0
|
||||
# File may not exist or may be empty — both acceptable
|
||||
if out.exists():
|
||||
assert out.read_text() == ""
|
||||
|
||||
def test_appends_to_existing_jsonl(self, tmp_path):
|
||||
from ai_mouse.config import BalabitAdapterConfig
|
||||
from ai_mouse.data_adapters.balabit import process_session
|
||||
|
||||
out = tmp_path / "out.jsonl"
|
||||
out.write_text('{"meta":{"start":[0,0],"end":[1,1]},"events":[]}\n', encoding="utf-8")
|
||||
|
||||
csv_path = tmp_path / "session_x"
|
||||
rows = [f"0,{i*0.05:.3f},NoButton,Move,{100+i*40},100" for i in range(10)]
|
||||
rows.append("0,0.550,Left,Pressed,510,100")
|
||||
_write_csv(csv_path, rows)
|
||||
|
||||
config = BalabitAdapterConfig()
|
||||
process_session(csv_path, out, config)
|
||||
|
||||
lines = out.read_text(encoding="utf-8").strip().split("\n")
|
||||
assert len(lines) == 2 # original + appended
|
||||
|
||||
Reference in New Issue
Block a user