Blame for code/qev_multitask_prepare.py

#!/usr/bin/env python3
"""Build an isolated QEV2 DLP multitask dataset.

The source DLP corpus has document-level findings.  This converter keeps that
ground truth and adds a small, deterministic bundle of type questions to every
record: is_dlp plus up to three type decisions.  Japanese vulnerability rows
are converted from their existing explicit annotations and are never used to
overwrite the official corpus.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import random
from collections import Counter
from pathlib import Path

TYPE_MAP = {
    "PERSON_NAME": "is_person_name",
    "GOVERNMENT_ID": "is_government_id",
    "CREDIT_CARD_NUMBER": "is_credit_card",
    "FINANCIAL_ACCOUNT_NUMBER": "is_financial_account",
    "PASSWORD": "is_password",
    "MEDICAL_RECORD_NUMBER": "is_medical_record",
    "EMAIL_ADDRESS": "is_email",
    "PHONE_NUMBER": "is_phone",
    "API_KEY": "is_api_key",
    "AUTH_TOKEN": "is_auth_token",
    "DEVICE_ID": "is_device_id",
    "MAC_ADDRESS": "is_mac_address",
    "DATE_OF_BIRTH": "is_birth_date",
    "IP_ADDRESS": "is_ip_address",
    "DRIVERS_LICENSE_NUMBER": "is_drivers_license",
    "STREET_ADDRESS": "is_address",
    "BLOOD_TYPE": "is_blood_type",
    "PASSPORT": "is_passport",
    "AGE": "is_age",
    "VEHICLE_IDENTIFICATION_NUMBER": "is_vehicle_id",
    "IBAN_CODE": "is_iban",
}

TYPE_ORDER = list(dict.fromkeys(TYPE_MAP.values()))
JP_KEY_MAP = {
    "phone": "is_phone",
    "date_of_birth": "is_birth_date",
    "drivers_license": "is_drivers_license",
    "passport": "is_passport",
    "government_id": "is_government_id",
    "medical": "is_medical_record",
    "biometric": "is_biometric",
    "email": "is_email",
    "address": "is_address",
    "credit_card": "is_credit_card",
    "password": "is_password",
    "api_key": "is_api_key",
    "auth_token": "is_auth_token",
}

LABEL_TEXT = {
    "en": {
        "is_dlp": ("Does this source text contain one or more actual privacy-sensitive values?", "No actual privacy-sensitive value is present.", "At least one actual privacy-sensitive value is present."),
        "is_email": "an actual email address",
        "is_phone": "an actual telephone number",
        "is_birth_date": "an actual date of birth",
        "is_drivers_license": "an actual driver's license number",
        "is_passport": "an actual passport or travel document number",
        "is_government_id": "an actual government-issued identifier",
        "is_medical_record": "an actual medical record or patient identifier",
        "is_person_name": "an actual person's name",
        "is_address": "an actual street address",
    },
    "jp": {
        "is_dlp": ("このテキストに実際のプライバシー・機密情報の値が1つ以上含まれていますか？", "実際の値は含まれていない。", "実際の値が含まれている。"),
        "is_email": "実際のメールアドレス",
        "is_phone": "実際の電話番号",
        "is_birth_date": "実際の生年月日（西暦・和暦・全角数字を含む）",
        "is_drivers_license": "実際の運転免許証番号・免許番号",
        "is_passport": "実際のパスポート番号・旅券番号",
        "is_government_id": "実際の国民ID・個人番号など政府発行ID",
        "is_medical_record": "実際の患者ID・診療録番号など医療識別子",
        "is_person_name": "実際の個人名",
        "is_address": "実際の住所",
    },
    "ko": {
        "is_dlp": ("이 텍스트에 실제 개인정보·민감정보 값이 하나 이상 포함되어 있습니까?", "실제 값이 포함되어 있지 않습니다.", "실제 값이 하나 이상 포함되어 있습니다."),
        "is_email": "실제 이메일 주소",
        "is_phone": "실제 전화번호",
        "is_birth_date": "실제 생년월일",
        "is_drivers_license": "실제 운전면허번호",
        "is_passport": "실제 여권번호",
        "is_government_id": "실제 정부 발행 식별번호",
        "is_medical_record": "실제 환자·진료 기록 식별자",
        "is_person_name": "실제 사람 이름",
        "is_address": "실제 주소",
    },
}


def lang_of(row: dict, default="en") -> str:
    value = row.get("language") or row.get("_meta", {}).get("language") or default
    value = {"ja": "jp", "japanese": "jp", "en-us": "en", "ko-kr": "ko"}.get(str(value).lower(), str(value).lower())
    return value if value in {"en", "jp", "ko"} else default


def question(qid: str, lang: str, value: bool, *, is_dlp=False) -> dict:
    if is_dlp:
        instruction, no_text, yes_text = LABEL_TEXT[lang]["is_dlp"]
    else:
        phrase = LABEL_TEXT[lang].get(qid, qid)
        if lang == "jp":
            instruction = f"このテキストに{phrase}が含まれていますか？全角・半角、区切り記号、表形式、改行を含む表記を判定します。"
        elif lang == "ko":
            instruction = f"이 텍스트에 {phrase}가 포함되어 있습니까? 전각·반각, 구분기호, 표 형식, 줄바꿈 표기를 포함해 판정합니다."
        else:
            instruction = f"Does this source text contain {phrase}? Include full-width/half-width, separators, tables, and line-wrapped forms."
        no_text = "実際の値は含まれていない。" if lang == "jp" else ("실제 값이 포함되어 있지 않습니다." if lang == "ko" else "No actual value is present.")
        yes_text = "実際の値が含まれている。" if lang == "jp" else ("실제 값이 포함되어 있습니다." if lang == "ko" else "An actual value is present.")
    return {
        "type": "noul",
        "instructions": instruction,
        "criteria": {"false": no_text, "true": yes_text},
        "label": "true" if value else "false",
    }


def stable_choice(seed: str, candidates: list[str], n: int) -> list[str]:
    rng = random.Random(int(hashlib.sha256(seed.encode()).hexdigest()[:16], 16))
    shuffled = list(candidates)
    rng.shuffle(shuffled)
    return shuffled[:n]


def official_record(row: dict, split: str, ordinal: int) -> dict:
    lang = lang_of(row)
    positive = {TYPE_MAP[x.get("info_type")] for x in row.get("findings", []) if x.get("info_type") in TYPE_MAP}
    # Three type slots per record keeps the demo-compatible four-question limit.
    chosen = list(sorted(positive))[:3]
    if len(chosen) < 3:
        pool = [x for x in TYPE_ORDER if x not in positive]
        chosen += stable_choice(row.get("id", str(ordinal)), pool, 3 - len(chosen))
    qs = {"is_dlp": question("is_dlp", lang, bool(row.get("findings")), is_dlp=True)}
    for qid in chosen:
        qs[qid] = question(qid, lang, qid in positive)
    return {
        "id": f"mt-{split}-{row.get('id', ordinal)}",
        "state": row.get("text", ""),
        "questions": qs,
        "_meta": {"origin": "official-dlp-findings", "source_id": row.get("id"), "source_split": split,
                  "language": lang, "info_types": sorted(x for x in positive), "binary_label": int(bool(row.get("findings")))},
    }


def ja_record(row: dict, split: str, ordinal: int) -> list[dict]:
    lang = "jp"
    state = row.get("state", "")
    src_q = row.get("questions", {})
    labels = {JP_KEY_MAP[k]: (q.get("label") == "true") for k, q in src_q.items() if k in JP_KEY_MAP}
    pos = {k for k, v in labels.items() if v}
    # retain explicitly annotated types; fill with negatives to create bundles.
    keys = list(labels)
    bundles = [keys[i:i + 3] for i in range(0, len(keys), 3)] or [[]]
    out = []
    for bi, bundle in enumerate(bundles):
        qs = {"is_dlp": question("is_dlp", lang, bool(row.get("_meta", {}).get("info_types") or row.get("questions", {}).get("is_dlp", {}).get("label") == "true"), is_dlp=True)}
        for qid in bundle:
            qs[qid] = question(qid, lang, labels[qid])
        out.append({"id": f"mt-ja-{split}-{row.get('id', ordinal)}-b{bi}", "state": state, "questions": qs,
                    "_meta": {"origin": "ja-vulnerability-annotations", "source_id": row.get("id"), "source_split": split,
                              "language": "jp", "info_types": sorted(pos), "binary_label": int(bool(pos))}})
    return out


def general_record(row: dict) -> dict:
    # Keep replay records unchanged except for a provenance marker.
    row = dict(row)
    meta = dict(row.get("_meta", {})); meta["origin"] = "general-replay"; row["_meta"] = meta
    # Make the multitask contract explicit for replay too.  The replay corpus
    # is the non-DLP general set, so this additional target is always false.
    if isinstance(row.get("questions"), dict) and "is_dlp" not in row["questions"]:
        lang = lang_of(row, "en")
        row["questions"] = {"is_dlp": question("is_dlp", lang, False, is_dlp=True), **row["questions"]}
    return row


def read_jsonl(path: Path):
    with path.open(encoding="utf-8") as f:
        for line in f:
            if line.strip():
                yield json.loads(line)


def write_jsonl(path: 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 build(args):
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    source_root, dlp_root, ja_root = map(Path, (args.source_root, args.dlp_root, args.ja_root))
    split_map = {"train": "train", "validation": "development", "test": "historical"}
    stats = Counter(); generated = {}
    for source_split, target_split in split_map.items():
        rows = []
        for i, row in enumerate(read_jsonl(source_root / f"{source_split}.jsonl")):
            rows.append(official_record(row, source_split, i)); stats[f"official_{target_split}"] += 1
        if target_split == "train":
            # Existing official QEV2 replay is already a carefully selected 30% general mix.
            for row in read_jsonl(dlp_root / "train-r0-replay30.jsonl"):
                if row.get("_meta", {}).get("origin") == "general-replay" or row.get("_meta", {}).get("group_id"):
                    rows.append(general_record(row)); stats["replay"] += 1
            ja_path = ja_root / "train-r4-focused.jsonl"
        elif target_split == "development":
            ja_path = ja_root / "development-r3.jsonl"
        else:
            ja_path = ja_root / "historical-test-multitask.jsonl"
        # JA source is explicitly annotated; include only Japanese records to avoid duplicating EN/KO.
        if ja_path.exists():
            for i, row in enumerate(read_jsonl(ja_path)):
                raw_lang = row.get("language") or row.get("_meta", {}).get("language")
                if str(raw_lang).lower() in {"jp", "ja", "japanese"} and row.get("_meta", {}).get("origin") != "official-dlp-findings":
                    converted = ja_record(row, target_split, i); rows.extend(converted); stats[f"ja_{target_split}"] += len(converted)
        write_jsonl(out / f"{target_split}-multitask.jsonl", rows)
        generated[target_split] = len(rows)
    manifest = {"schema": "qev2-train-multitask-v1", "question_limit": 4, "type_map": TYPE_MAP,
                "logical_split": {"train": "train-multitask.jsonl", "development": "development-multitask.jsonl", "historical": "historical-multitask.jsonl"},
                "official_source": str(source_root), "replay_source": str(dlp_root / "train-r0-replay30.jsonl"),
                "ja_source": str(ja_root), "counts": generated, "stats": dict(stats),
                "label_rule": "is_dlp=true iff source findings/info_types is non-empty; type labels come only from findings or explicit JA annotations"}
    (out / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(manifest, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--source-root", required=True)
    ap.add_argument("--dlp-root", required=True)
    ap.add_argument("--ja-root", required=True)
    build(ap.parse_args())