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

"""
BIO tagging schema for labor arbitration case element extraction.
15 entity types, 31 labels total (15 x B/I + O).
"""

from __future__ import annotations

from typing import Any

# Entity types with their BIO prefixes
ENTITY_TYPES = [
    "APPLICANT_NAME",       # 申请人姓名
    "RESPONDENT_NAME",      # 被申请人名称
    "ENTRY_DATE",           # 入职日期
    "LEAVE_DATE",           # 离职日期
    "FILING_DATE",          # 立案日期
    "MONTH_SALARY",         # 月工资标准
    "CLAIM_AMOUNT",         # 主张金额
    "WORKER_POSITION",      # 劳动者岗位
    "ARBITRATION_ORG",      # 仲裁机构
    "LAW_REF",              # 法律条款引用
    "CASE_NUMBER",          # 案号
    "EVIDENCE",             # 证据材料
    "TERMINATION_REASON",   # 解除原因
    "OVERTIME_DESC",        # 加班描述
    "WORK_DURATION",        # 工作年限
]

# Build label list: O + B-ENTITY + I-ENTITY for each type
LABELS = ["O"]
LABEL2ID = {"O": 0}
for entity in ENTITY_TYPES:
    b_label = f"B-{entity}"
    i_label = f"I-{entity}"
    LABELS.append(b_label)
    LABELS.append(i_label)
    LABEL2ID[b_label] = len(LABEL2ID)
    LABEL2ID[i_label] = len(LABEL2ID)

ID2LABEL = {v: k for k, v in LABEL2ID.items()}
NUM_LABELS = len(LABELS)


def get_entity_spans_from_bio(
    tokens: list[str],
    labels: list[int],
) -> list[dict[str, Any]]:
    """Convert BIO tag sequence back to entity spans."""
    spans = []
    current_entity = None
    current_tokens = []
    current_start = -1

    for i, (token, label_id) in enumerate(zip(tokens, labels)):
        label = ID2LABEL.get(label_id, "O")
        if label.startswith("B-"):
            if current_entity is not None:
                spans.append({
                    "entity": current_entity,
                    "text": "".join(current_tokens),
                    "start": current_start,
                    "end": i,
                })
            current_entity = label[2:]
            current_tokens = [token]
            current_start = i
        elif label.startswith("I-") and current_entity == label[2:]:
            current_tokens.append(token)
        else:
            if current_entity is not None:
                spans.append({
                    "entity": current_entity,
                    "text": "".join(current_tokens),
                    "start": current_start,
                    "end": i,
                })
            current_entity = None
            current_tokens = []
            current_start = -1

    if current_entity is not None:
        spans.append({
            "entity": current_entity,
            "text": "".join(current_tokens),
            "start": current_start,
            "end": len(tokens),
        })

    return spans


def spans_to_bio(
    tokens: list[str],
    spans: list[dict[str, Any]],
) -> list[int]:
    """Convert entity spans to BIO tag sequence."""
    labels = ["O"] * len(tokens)
    for span in spans:
        entity = span["entity"]
        start = span["start"]
        end = span["end"]
        b_key = f"B-{entity}"
        i_key = f"I-{entity}"
        if b_key in LABEL2ID:
            labels[start] = b_key
            for i in range(start + 1, min(end, len(tokens))):
                labels[i] = i_key
    return [LABEL2ID[l] for l in labels]


# Mapping from entity types to flat schema field keys (for converting
# extracted spans into the element dictionary expected by the backend)
ENTITY_TO_FIELD = {
    "APPLICANT_NAME": "applicant_name",
    "RESPONDENT_NAME": "respondent_name",
    "ENTRY_DATE": "entry_date",
    "LEAVE_DATE": "leave_date",
    "FILING_DATE": "filing_date",
    "MONTH_SALARY": "month_salary",
    "CLAIM_AMOUNT": "claims",  # merges into claims.amount_total
    "WORKER_POSITION": "worker_position",
    "ARBITRATION_ORG": "arbitration_org",
    "LAW_REF": "law_refs",     # list field
    "CASE_NUMBER": "case_number",
    "EVIDENCE": "evidence_materials",  # list field
    "TERMINATION_REASON": "termination_reason",
    "OVERTIME_DESC": "overtime_desc",
    "WORK_DURATION": "work_duration_text",
}