labor-arbitration-system/backend/tools/prepare_dataset.py

#!/usr/bin/env python3
"""
Data cleaning and preparation for labor arbitration case corpus.
Reads uploaded case files, filters non-labor cases, sanitizes privacy markers,
and outputs a cleaned JSON corpus.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

UPLOADS_DIR = Path(__file__).resolve().parent.parent / "uploads"
OUTPUT_PATH = Path(__file__).resolve().parent.parent / "data" / "raw_corpus.json"

# Keywords that indicate a labor arbitration case (must have at least one)
LABOR_KEYWORDS = [
    "申请人", "被申请人", "劳动关系", "工资", "加班",
    "劳动合同", "辞退", "解除劳动", "经济补偿", "工伤",
    "仲裁请求", "请求事项", "劳动争议", "社保", "仲裁委员会",
]

# Keywords that suggest a non-labor case (civil loan, contract dispute, etc.)
NON_LABOR_KEYWORDS = [
    "借款", "欠条", "民间借贷", "买卖合同", "租赁合同",
    "房屋买卖", "知识产权", "侵犯商标",
]


def _clean_privacy(text: str) -> str:
    """Replace privacy-sensitive patterns with placeholders."""
    # Base64-like identity hashes (long alphanumeric/+//= strings)
    text = re.sub(r"[A-Za-z0-9+/]{30,}={0,3}", "[ID_HASH]", text)

    # Chinese ID card numbers (18 digits, possibly with X at end)
    text = re.sub(r"\b[1-9]\d{5}(?:19|20)\d{2}(?:0[1-9]|1[0-2])(?:0[1-9]|[12]\d|3[01])\d{3}[\dXx]\b", "[ID_NUM]", text)

    # Phone numbers
    text = re.sub(r"1[3-9]\d{9}", "[PHONE]", text)

    # Generic XX placeholder (keep as is - already anonymized)
    # text = re.sub(r"X{2,}", "[REDACTED]", text)

    return text


def _is_labor_case(text: str) -> bool:
    """Check if the text appears to be a labor arbitration case."""
    labor_score = sum(1 for kw in LABOR_KEYWORDS if kw in text)
    non_labor_score = sum(1 for kw in NON_LABOR_KEYWORDS if kw in text)
    return labor_score >= 2 and non_labor_score == 0


def _normalize_text(text: str) -> str:
    """Normalize whitespace and encoding."""
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    # Normalize multiple newlines
    text = re.sub(r"\n{3,}", "\n\n", text)
    # Strip trailing/leading whitespace
    text = text.strip()
    return text


def prepare_corpus() -> list[dict]:
    """Main data preparation function."""
    if not UPLOADS_DIR.exists():
        raise FileNotFoundError(f"Uploads directory not found: {UPLOADS_DIR}")

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)

    corpus = []
    skipped = []
    files = sorted(UPLOADS_DIR.glob("*.txt"))

    for idx, filepath in enumerate(files):
        text = filepath.read_text(encoding="utf-8")
        text = _normalize_text(text)

        if not text.strip():
            skipped.append({"file": filepath.name, "reason": "empty"})
            continue

        if not _is_labor_case(text):
            skipped.append({"file": filepath.name, "reason": "non-labor"})
            continue

        text = _clean_privacy(text)

        # Extract case ID from filename (e.g., "10_xxx.txt" -> 10, or "364-014-2022-0001.txt")
        try:
            case_id = int(filepath.stem.split("_")[0])
        except ValueError:
            case_id = idx + 1000  # fallback for numeric filenames without underscore prefix

        corpus.append({
            "case_id": case_id,
            "file": filepath.name,
            "text": text,
        })

    # Save corpus
    output = {
        "total_files": len(files),
        "valid_cases": len(corpus),
        "skipped": skipped,
        "cases": corpus,
    }
    OUTPUT_PATH.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"Processed {len(files)} files:")
    print(f"  Valid labor cases: {len(corpus)}")
    print(f"  Skipped: {len(skipped)}")
    for s in skipped:
        print(f"    - {s['file']}: {s['reason']}")
    print(f"Output saved to: {OUTPUT_PATH}")

    return corpus


if __name__ == "__main__":
    prepare_corpus()