labor-arbitration-system/nlp-service/training/prepare_training_data.py

"""
Auto-labeling pipeline: runs the rule-based extractor on all cleaned cases,
converts extracted fields to BIO token-level labels plus classification labels,
and outputs a training dataset ready for model training.

Run from project root:
  python -m nlp-service.training.prepare_training_data
"""

from __future__ import annotations

import json
import re
import sys
from pathlib import Path
from typing import Any

# Ensure backend is on path
BACKEND = Path(__file__).resolve().parent.parent.parent / "backend"
sys.path.insert(0, str(BACKEND))

from app.extractor import (
    RuleBasedLaborExtractor,
    _clean_text,
    parse_amount,
)

from bio_schema import (
    ENTITY_TYPES,
    ENTITY_TO_FIELD,
    LABEL2ID,
    NUM_LABELS,
)

RAW_CORPUS = BACKEND / "data" / "raw_corpus.json"
OUTPUT = BACKEND / "data" / "training_dataset.json"

CASE_CAUSE_CLASSES = [
    "劳动关系纠纷类",
    "工伤保险待遇纠纷",
    "追索劳动报酬",
    "经济补偿金纠纷",
    "赔偿金纠纷",
    "生育保险待遇纠纷",
]

CONTRACT_TYPE_CLASSES = [
    "无固定期限劳动合同",
    "固定期限劳动合同",
    "未订立书面劳动合同",
    "未知",
]

EMPLOYMENT_TYPE_CLASSES = [
    "劳务派遣",
    "非全日制用工",
    "全日制用工",
    "劳务关系",
    "未知",
]


def _find_span_position(text: str, value: str) -> tuple[int, int] | None:
    """Find the character position of a value in the text. Returns (start, end)."""
    if not value or not text:
        return None
    idx = text.find(str(value))
    if idx >= 0:
        return idx, idx + len(str(value))
    return None


def _find_spans_for_list(text: str, values: list[str]) -> list[tuple[int, int, str]]:
    """Find character positions for a list of values. Returns [(start, end, value), ...]."""
    spans = []
    for val in values:
        pos = _find_span_position(text, val)
        if pos:
            spans.append((pos[0], pos[1], val))
    return spans


def _char_spans_to_token_bio(
    text: str,
    char_spans: list[tuple[int, int, str, str]],  # (start, end, value, entity_type)
) -> tuple[list[str], list[int]]:
    """
    Convert character-level entity spans to token-level BIO labels.
    Uses a simple character-level tokenization approach (since Chinese
    has no natural word boundaries, we split by whitespace and then by
    individual characters within each segment).
    """
    # Tokenize: keep characters together, split on whitespace
    # For Chinese text, tokens will be individual Chinese characters
    # with multi-character sequences kept together when they form words
    # between whitespace boundaries
    tokens: list[str] = []
    token_char_starts: list[int] = []
    token_char_ends: list[int] = []

    for m in re.finditer(r"[^\s]+|\s+", text):
        word = m.group(0)
        if word.strip():
            # For Chinese, split each character as a token
            for i, ch in enumerate(word):
                tokens.append(ch)
                token_char_starts.append(m.start() + i)
                token_char_ends.append(m.start() + i + 1)
        else:
            # Keep whitespace as single token
            tokens.append(word)
            token_char_starts.append(m.start())
            token_char_ends.append(m.end())

    # Initialize all labels as O
    labels = [LABEL2ID["O"]] * len(tokens)

    # Mark BIO labels based on character spans
    for char_start, char_end, _value, entity_type in char_spans:
        b_key = f"B-{entity_type}"
        i_key = f"I-{entity_type}"
        if b_key not in LABEL2ID:
            continue

        first_token_idx = None
        for ti, (ts, te) in enumerate(zip(token_char_starts, token_char_ends)):
            if ts >= char_start and te <= char_end:
                if first_token_idx is None:
                    first_token_idx = ti
                    labels[ti] = LABEL2ID[b_key]
                else:
                    labels[ti] = LABEL2ID[i_key]

    return tokens, labels


def _extract_spans_from_elements(text: str, elements: dict[str, Any]) -> list[tuple[int, int, str, str]]:
    """
    Convert extracted elements dictionary into character-level entity spans.
    Returns list of (char_start, char_end, value_string, entity_type).
    """
    char_spans: list[tuple[int, int, str, str]] = []

    # Direct string fields
    str_fields = [
        ("applicant_name", "APPLICANT_NAME"),
        ("respondent_name", "RESPONDENT_NAME"),
        ("entry_date", "ENTRY_DATE"),
        ("leave_date", "LEAVE_DATE"),
        ("filing_date", "FILING_DATE"),
        ("arbitration_org", "ARBITRATION_ORG"),
        ("worker_position", "WORKER_POSITION"),
        ("case_number", "CASE_NUMBER"),
        ("termination_reason", "TERMINATION_REASON"),
        ("overtime_desc", "OVERTIME_DESC"),
        ("work_duration_text", "WORK_DURATION"),
    ]

    for field, entity_type in str_fields:
        val = elements.get(field)
        if val and isinstance(val, str) and val.strip():
            pos = _find_span_position(text, val.strip())
            if pos:
                char_spans.append((pos[0], pos[1], val.strip(), entity_type))

    # Numeric fields (need to convert to string for span matching)
    num_fields = [
        ("month_salary", "MONTH_SALARY"),
    ]
    for field, entity_type in num_fields:
        val = elements.get(field)
        if val is not None and val != "":
            val_str = str(int(val)) if isinstance(val, float) and val == int(val) else str(val)
            pos = _find_span_position(text, val_str)
            if pos:
                char_spans.append((pos[0], pos[1], val_str, entity_type))

    # List fields (law_refs, evidence_materials)
    law_refs = elements.get("law_refs") or []
    for law in law_refs:
        pos = _find_span_position(text, law)
        if pos:
            char_spans.append((pos[0], pos[1], law, "LAW_REF"))

    evidence = elements.get("evidence_materials") or []
    for ev in evidence:
        pos = _find_span_position(text, ev)
        if pos:
            char_spans.append((pos[0], pos[1], ev, "EVIDENCE"))

    # Claims amounts
    claims = elements.get("claims") or {}
    if isinstance(claims, dict):
        amount = claims.get("amount_total")
        if amount is not None:
            amount_str = str(int(amount)) if isinstance(amount, float) else str(amount)
            pos = _find_span_position(text, amount_str)
            if pos:
                char_spans.append((pos[0], pos[1], amount_str, "CLAIM_AMOUNT"))

    return char_spans


def _extract_classification_labels(elements: dict[str, Any]) -> dict[str, int]:
    """Convert element fields to classification label indices."""
    labels = {}

    # Primary cause type (6 classes)
    cause = elements.get("tmpl_primary_cause") or elements.get("primary_cause_type") or "劳动关系纠纷类"
    if cause in CASE_CAUSE_CLASSES:
        labels["case_cause"] = CASE_CAUSE_CLASSES.index(cause)
    else:
        labels["case_cause"] = 0

    # Contract type (4 classes)
    ct = elements.get("contract_type") or "未知"
    if ct in CONTRACT_TYPE_CLASSES:
        labels["contract_type"] = CONTRACT_TYPE_CLASSES.index(ct)
    else:
        labels["contract_type"] = 3

    # Employment type (5 classes)
    et = elements.get("employment_type") or "未知"
    if et in EMPLOYMENT_TYPE_CLASSES:
        labels["employment_type"] = EMPLOYMENT_TYPE_CLASSES.index(et)
    else:
        labels["employment_type"] = 4

    # Binary/ternary fields
    binary_fields = {
        "lr1_contract_signed": {"是": 0, "否": 1, None: 2},
        "lr1_open_ended_contract": {"是": 0, "否": 1, None: 2},
        "lr1_double_wage_no_contract": {"是": 0, "否": 1, None: 2},
        "lr1_si_joined": {"是": 0, "否": 1, None: 2},
        "wi_si_joined": {"是": 0, "否": 1, None: 2},
        "wi_recognize_ec": {"是": 0, "否": 1, None: 2},
        "dm_contract_exists": {"是": 0, "否": 1, None: 2},
        "dm_contract_continue": {"是": 0, "否": 1, None: 2},
        "mi_contract_continue": {"是": 0, "否": 1, None: 2},
    }
    for field, mapping in binary_fields.items():
        raw = elements.get(field)
        labels[f"bool_{field}"] = mapping.get(raw, 2)

    return labels


def _extract_numeric_labels(elements: dict[str, Any]) -> dict[str, float | None]:
    """Extract numeric field values for regression targets."""
    amount_keys = [
        "month_salary",
        "lr1_pay_amount", "lr1_si_benefit_amount",
        "wi_benefit_amount_total", "wi_benefit_disability",
        "wi_benefit_prosthetic", "wi_benefit_medical_allowance",
        "wi_benefit_travel", "wi_benefit_rehab", "wi_benefit_nursing",
        "wi_benefit_meal", "wi_si_benefit_amt",
        "sr_claim_amount", "sr_claim_deducted_pay",
        "sr_claim_overtime_pay", "sr_claim_living_allowance",
        "sr_high_temp_allowance", "sr_annual_leave_pay",
        "sr_overtime_amount",
        "ec_avg_salary_12m", "ec_claim_amount",
        "ec_double_wage_part", "ec_illegal_term_part",
        "ec_illegal_probation_part", "ec_extra_compensation_part",
        "ec_notice_pay", "ec_additional_damages",
        "dm_claim_amount", "dm_illegal_dismissal_damages",
        "mi_claim_amount", "mi_maternity_medical",
        "mi_maternity_allowance_salary", "mi_additional_damages",
        "mi_travel_accommodation",
    ]

    nums: dict[str, float | None] = {}
    for key in amount_keys:
        val = elements.get(key)
        if val is not None and val != "":
            try:
                nums[key] = float(val)
            except (ValueError, TypeError):
                nums[key] = None
        else:
            nums[key] = None

    return nums


def generate_training_dataset() -> dict:
    """Main function: load corpus, auto-label, output training dataset."""
    if not RAW_CORPUS.exists():
        sys.exit(f"Raw corpus not found. Run prepare_dataset.py first.\n  Missing: {RAW_CORPUS}")

    corpus = json.loads(RAW_CORPUS.read_text(encoding="utf-8"))
    cases = corpus["cases"]

    if not cases:
        sys.exit("No valid cases in corpus.")

    extractor = RuleBasedLaborExtractor()
    dataset = []

    for case in cases:
        text = case["text"]
        cleaned = _clean_text(text)

        # Run rule-based extraction
        try:
            elements = extractor.extract(cleaned)
        except Exception as e:
            print(f"  ERROR extracting case {case['case_id']}: {e}")
            elements = {}

        # Generate BIO spans
        char_spans = _extract_spans_from_elements(cleaned, elements)
        tokens, bio_labels = _char_spans_to_token_bio(cleaned, char_spans)

        # Generate classification labels
        cls_labels = _extract_classification_labels(elements)

        # Generate numeric labels
        num_labels = _extract_numeric_labels(elements)

        dataset.append({
            "case_id": case["case_id"],
            "file": case["file"],
            "text": cleaned,
            "tokens": tokens,
            "bio_labels": bio_labels,
            "bio_label_names": [
                {v: k for k, v in LABEL2ID.items()}.get(l, "O")
                for l in bio_labels
            ],
            "classification_labels": cls_labels,
            "numeric_labels": num_labels,
            "rule_elements": elements,
        })

        # Count entities found
        n_entities = sum(1 for l in bio_labels if l != LABEL2ID["O"])
        print(f"  Case {case['case_id']:>3}: {len(tokens):>5} tokens, "
              f"{len(char_spans):>3} spans, {n_entities:>4} labeled entities")

    # Save dataset
    output = {
        "version": "1.0",
        "total_cases": len(dataset),
        "label_schema": {
            "num_bio_labels": NUM_LABELS,
            "bio_label_to_id": LABEL2ID,
            "case_cause_classes": CASE_CAUSE_CLASSES,
            "contract_type_classes": CONTRACT_TYPE_CLASSES,
            "employment_type_classes": EMPLOYMENT_TYPE_CLASSES,
        },
        "dataset": dataset,
    }

    OUTPUT.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nTraining dataset saved: {OUTPUT}")
    print(f"Total cases: {len(dataset)}")
    return output


if __name__ == "__main__":
    generate_training_dataset()