code/prepare_data.py

#!/usr/bin/env python3
import argparse
import json
import math
from collections import defaultdict
from pathlib import Path


LABEL_MAP = {"反复折返": "折返"}
TIME_FIELDS = [
    "elapsed_seconds_in_current_behavior",
    "estimated_remaining_seconds",
    "full_remaining_seconds",
    "expected_end_time",
]
OUTPUT_FIELDS = [
    "current_behavior",
    "is_transition",
    "elapsed_seconds_in_current_behavior",
    "estimated_remaining_seconds",
    "full_remaining_seconds",
    "expected_end_time",
    "next_possible_behavior",
    "stage_index",
    "total_stages",
    "sequence_so_far",
]


def normalize_label(value):
    return LABEL_MAP.get(value, value)


def normalize_tree(value):
    if isinstance(value, dict):
        out = {}
        for key, item in value.items():
            if key in {"current_behavior", "next_possible_behavior", "label"} and isinstance(item, str):
                out[key] = normalize_label(item)
            elif key in {"full_sequence_order", "label_space"} and isinstance(item, list):
                out[key] = [normalize_label(x) if isinstance(x, str) else x for x in item]
            else:
                out[key] = normalize_tree(item)
        return out
    if isinstance(value, list):
        return [normalize_tree(item) for item in value]
    if isinstance(value, str):
        return normalize_label(value)
    return value


def read_chat_jsonl(path):
    with path.open(encoding="utf-8") as f:
        for line_no, line in enumerate(f, 1):
            if not line.strip():
                continue
            try:
                obj = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_no} is not valid JSONL: {exc}") from exc
            yield obj


def get_assistant_json(example):
    content = example["messages"][-1]["content"]
    return json.loads(content) if isinstance(content, str) else content


def set_assistant_json(example, data):
    example["messages"][-1]["content"] = json.dumps(data, ensure_ascii=False, separators=(",", ":"))


def clean_example(example):
    example = normalize_tree(example)
    for message in example.get("messages", []):
        content = message.get("content")
        if isinstance(content, str):
            try:
                parsed = json.loads(content)
            except Exception:
                continue
            message["content"] = json.dumps(normalize_tree(parsed), ensure_ascii=False, separators=(",", ":"))
    assistant = get_assistant_json(example)
    assistant = normalize_tree(assistant)
    set_assistant_json(example, assistant)
    return example


def write_jsonl(path, rows):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")


def percentile(values, q):
    if not values:
        return 0.0
    values = sorted(values)
    idx = (len(values) - 1) * q
    lo = math.floor(idx)
    hi = math.ceil(idx)
    if lo == hi:
        return float(values[lo])
    return float(values[lo] * (hi - idx) + values[hi] * (idx - lo))


def build_thresholds(clean_train):
    by_label = defaultdict(list)
    for ex in clean_train:
        assistant = get_assistant_json(ex)
        label = assistant.get("current_behavior")
        elapsed = assistant.get("elapsed_seconds_in_current_behavior")
        if isinstance(elapsed, (int, float)):
            by_label[label].append(float(elapsed))
    thresholds = {label: max(5.0, percentile(vals, 0.95)) for label, vals in by_label.items()}
    thresholds["__default__"] = max(30.0, percentile([v for vals in by_label.values() for v in vals], 0.95))
    return thresholds


def used_areas(sequence):
    mapping = {
        "门": {"进入", "门锁", "靠近门", "离开"},
        "马桶": {"靠近马桶", "马桶盖", "马桶垫纸", "坐下", "坐用马桶", "站用马桶", "卷筒厕纸", "起身", "冲水"},
        "洗手池": {"靠近洗手池", "洗手", "刷牙"},
        "垃圾桶": {"垃圾桶"},
    }
    labels = {item.get("label") for item in sequence or []}
    return [area for area, area_labels in mapping.items() if labels & area_labels]


def qa_target(assistant, thresholds):
    sequence = assistant.get("sequence_so_far") or []
    current = assistant.get("current_behavior")
    full_remaining = assistant.get("full_remaining_seconds")
    elapsed = assistant.get("elapsed_seconds_in_current_behavior")
    occupied = bool(sequence) and current != "离开" and (not isinstance(full_remaining, (int, float)) or full_remaining > 0)
    time_to_free_minutes = round(max(0.0, float(full_remaining or 0.0)) / 60.0, 2)
    threshold = thresholds.get(current, thresholds.get("__default__", 120.0))
    frequent_switching = len(sequence) >= 12
    long_current = isinstance(elapsed, (int, float)) and elapsed > threshold
    return {
        "occupied": occupied,
        "time_to_free_minutes": time_to_free_minutes,
        "used_areas": used_areas(sequence),
        "is_abnormal": bool(long_current or frequent_switching),
    }


def make_qa_example(source_example, thresholds):
    assistant = get_assistant_json(source_example)
    compact = {key: assistant.get(key) for key in OUTPUT_FIELDS}
    user_payload = {
        "task": "qa_from_behavior_json",
        "model_output_json": compact,
        "questions": [
            "卫生间是否被占用？",
            "预计多长时间之后卫生间会空出？",
            "卫生间内哪些区域被使用过？",
            "卫生间是否存在异常？",
        ],
    }
    return {
        "messages": [
            {
                "role": "system",
                "content": "你是卫生间状态问答模型。请只根据输入的行为分析 JSON 回答，并只输出 JSON：occupied(bool), time_to_free_minutes(number), used_areas(array), is_abnormal(bool)。",
            },
            {"role": "user", "content": json.dumps(user_payload, ensure_ascii=False, separators=(",", ":"))},
            {"role": "assistant", "content": json.dumps(qa_target(assistant, thresholds), ensure_ascii=False, separators=(",", ":"))},
        ]
    }


def dataset_stats(rows):
    labels = defaultdict(int)
    sample_ids = set()
    for ex in rows:
        assistant = get_assistant_json(ex)
        labels[assistant.get("current_behavior")] += 1
        try:
            sample_ids.add(json.loads(ex["messages"][1]["content"]).get("sample_id"))
        except Exception:
            pass
    return {"num_examples": len(rows), "num_sample_ids": len(sample_ids), "label_counts": dict(sorted(labels.items()))}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--train", default="train.jsonl")
    parser.add_argument("--val", default="val.jsonl")
    parser.add_argument("--out-dir", default="data/processed")
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    train = [clean_example(ex) for ex in read_chat_jsonl(Path(args.train))]
    val = [clean_example(ex) for ex in read_chat_jsonl(Path(args.val))]
    thresholds = build_thresholds(train)
    train_qa = [make_qa_example(ex, thresholds) for ex in train]
    val_qa = [make_qa_example(ex, thresholds) for ex in val]

    write_jsonl(out_dir / "train_struct.jsonl", train)
    write_jsonl(out_dir / "val_struct.jsonl", val)
    write_jsonl(out_dir / "train_qa.jsonl", train_qa)
    write_jsonl(out_dir / "val_qa.jsonl", val_qa)
    write_jsonl(out_dir / "train_mixed.jsonl", train + train_qa)
    write_jsonl(out_dir / "val_mixed.jsonl", val + val_qa)

    summary = {
        "normalization": LABEL_MAP,
        "train_struct": dataset_stats(train),
        "val_struct": dataset_stats(val),
        "train_qa": {"num_examples": len(train_qa)},
        "val_qa": {"num_examples": len(val_qa)},
        "abnormal_elapsed_thresholds_p95": thresholds,
        "qa_schema": ["occupied", "time_to_free_minutes", "used_areas", "is_abnormal"],
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()