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

"""
Data augmentation for labor arbitration training data.
Strategies:
  1. Synonym replacement (legal domain vocabulary)
  2. Entity masking and variation
  3. Qwen synthetic case generation via Ollama
"""

from __future__ import annotations

import json
import random
import re
import sys
from pathlib import Path
from typing import Any
from copy import deepcopy

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

TRAINING_DATASET = BACKEND / "data" / "training_dataset.json"
OUTPUT = BACKEND / "data" / "augmented_dataset.json"

# Legal domain synonym groups
SYNONYM_MAP = {
    "工资": ["薪资", "劳动报酬", "薪酬", "工钱"],
    "加班": ["超时工作", "延时工作", "延长工作时间"],
    "辞退": ["解雇", "开除", "解除劳动合同"],
    "经济补偿": ["经济赔偿金", "补偿款", "经济补助"],
    "仲裁": ["劳动仲裁", "争议仲裁", "调解仲裁"],
    "申请人": ["申请方", "申请人方", "申诉人"],
    "被申请人": ["被申请方", "被诉人", "答辩方"],
    "劳动合同": ["劳动协议", "用工合同", "聘用合同"],
    "拖欠": ["欠付", "未支付", "拖欠支付", "未付"],
    "社保": ["社会保险", "五险", "社会保险费"],
    "工伤": ["因工受伤", "工作伤害", "职业伤害"],
    "劳动者": ["员工", "职工", "工人", "雇员"],
    "用人单位": ["雇主", "公司", "企业", "用工单位"],
    "月薪": ["月工资", "月收入", "月均工资"],
    "入职": ["到岗", "入职工", "开始工作", "参加工作"],
    "离职": ["离开公司", "辞职", "解除劳动关系"],
}

# Company name templates
COMPANY_TEMPLATES = [
    "XX市{name}有限公司",
    "XX省{name}{industry}有限公司",
    "{name}科技（XX）有限公司",
    "XX市{name}{industry}厂",
    "XX新区{name}实业有限公司",
    "{name}（中国）有限公司XX分公司",
]

# Person name pool
SURNAMES = ["张", "李", "王", "刘", "陈", "杨", "赵", "黄", "周", "吴", "徐", "孙", "马", "朱", "胡"]
GIVEN_NAMES = ["明", "华", "强", "伟", "芳", "丽", "敏", "静", "勇", "军", "磊", "洋", "涛", "鑫", "宇"]

# Case cause templates for Qwen generation
CAUSE_PROMPT_TEMPLATES = {
    "劳动关系纠纷类": "确认劳动关系的存在，涉及未签订书面劳动合同，请求二倍工资差额",
    "工伤保险待遇纠纷": "在工作中受伤，涉及工伤认定、劳动能力鉴定及一次性伤残补助金等工伤保险待遇",
    "追索劳动报酬": "用人单位拖欠工资、加班费、高温津贴等劳动报酬，要求支付拖欠款项",
    "经济补偿金纠纷": "因用人单位未依法缴纳社保或未及时足额支付工资，劳动者提出解除劳动合同并要求经济补偿金",
    "赔偿金纠纷": "用人单位违法解除或终止劳动合同，要求支付违法解除劳动合同赔偿金（2N）",
    "生育保险待遇纠纷": "女职工因生育未能享受生育津贴和生育医疗费用，要求用人单位支付生育保险待遇",
}


def synonym_replace(text: str, probability: float = 0.3) -> str:
    """Randomly replace legal terms with synonyms."""
    result = list(text)
    for match in re.finditer(r"[一-鿿]{2,6}", text):
        word = match.group(0)
        if word in SYNONYM_MAP and random.random() < probability:
            replacement = random.choice(SYNONYM_MAP[word])
            result[match.start():match.end()] = list(replacement)
    return "".join(result)


def mask_entity_variations(text: str) -> list[str]:
    """
    Generate 2 entity-masked variants of the same text.
    Replaces names/dates/amounts with different but realistic values.
    """
    variants = []
    for _ in range(2):
        variant = text

        # Replace company names
        variant = re.sub(
            r"XX[市省]?\S{0,8}(?:有限|责任|实业|科技|贸易|工程|装饰)?(?:公司|厂|企业|集团)",
            lambda m: random.choice(COMPANY_TEMPLATES).format(
                name=random.choice(["恒达", "鑫源", "瑞丰", "永盛", "华泰", "新锐", "博远"]),
                industry=random.choice(["电子", "机械", "纺织", "餐饮", "物流", "建筑", "软件"]),
            ),
            variant,
        )

        # Replace XX placeholders with realistic names
        while "XX" in variant:
            name = random.choice(SURNAMES) + random.choice(GIVEN_NAMES)
            variant = variant.replace("XX", name, 1)

        # Replace dates within reasonable range
        def date_shifter(m: re.Match) -> str:
            year = int(m.group(1))
            month = int(m.group(2))
            day = int(m.group(3))
            year += random.choice([-1, 0, 1])
            month = max(1, min(12, month + random.choice([-1, 0, 1])))
            day = max(1, min(28, day + random.choice([-2, -1, 0, 1, 2])))
            return f"{year}年{month}月{day}日"

        variant = re.sub(
            r"([12][0-9]{3})[年./-]([01]?[0-9])[月./-]([0-3]?[0-9])[日]?",
            date_shifter,
            variant,
        )

        variants.append(variant)

    return variants


def generate_synthetic_cases_with_qwen(
    num_per_cause: int = 5,
    ollama_host: str = "http://localhost:11434",
    ollama_model: str = "qwen2.5:3b",
) -> list[dict]:
    """
    Use Ollama Qwen to generate synthetic labor arbitration case texts.
    Generates num_per_cause cases for each of the 6 cause types.
    """
    try:
        import ollama
    except ImportError:
        print("  [SKIP] ollama package not installed, skipping synthetic generation")
        return []

    client = ollama.Client(host=ollama_host)
    synthetic_cases = []

    for cause_type, description in CAUSE_PROMPT_TEMPLATES.items():
        for i in range(num_per_cause):
            prompt = f"""请生成一个虚构的中国劳动仲裁案件申请书文本，要求如下：

案由类型：{cause_type}
案件特征：{description}

请按以下格式输出完整的仲裁申请书：

申请人：[虚构姓名]，[性别]，[出生年份]年出生，住址[虚构地址]。
被申请人：[虚构公司名称]，住所[虚构地址]。

请求事项：
1、[仲裁请求1]
2、[仲裁请求2]
3、[仲裁请求3，可选]

事实与理由：
[包含入职时间、工作岗位、工资标准、争议发生经过的事实描述，不少于200字]

注意：所有信息必须虚构，不能出现真实人名地名。金额使用合理范围（月薪3000-15000元，总请求金额在5000-500000元之间）。日期使用2018-2024年之间的日期。"""

            try:
                response = client.chat(
                    model=ollama_model,
                    messages=[
                        {"role": "system", "content": "你是一个法律文书生成助手，只输出仲裁申请书正文，不输出其他内容。"},
                        {"role": "user", "content": prompt},
                    ],
                    options={"temperature": 0.8, "num_predict": 2048},
                )
                generated_text = response["message"]["content"].strip()
                # Remove markdown code fences if present
                generated_text = re.sub(r"^```[^\n]*\n?", "", generated_text)
                generated_text = re.sub(r"\n```$", "", generated_text)

                if len(generated_text) >= 100:
                    synthetic_cases.append({
                        "text": generated_text,
                        "cause_type": cause_type,
                        "source": "qwen_synthetic",
                    })
                    print(f"  Generated: {cause_type} #{i+1} ({len(generated_text)} chars)")
            except Exception as e:
                print(f"  [ERROR] Failed to generate {cause_type} #{i+1}: {e}")
                continue

    print(f"  Total synthetic cases generated: {len(synthetic_cases)}")
    return synthetic_cases


def augment_dataset(
    include_synthetic: bool = True,
    synthetic_per_cause: int = 5,
) -> dict:
    """Main augmentation function."""
    if not TRAINING_DATASET.exists():
        sys.exit(f"Training dataset not found: {TRAINING_DATASET}")

    data = json.loads(TRAINING_DATASET.read_text(encoding="utf-8"))
    dataset = data["dataset"]

    augmented = []
    orig_ids: set[int] = set()

    for case in dataset:
        orig_ids.add(case["case_id"])
        augmented.append(case)  # Keep original

        text = case["text"]

        # Strategy 1: Synonym replacement (3x per case)
        for _ in range(3):
            variant_text = synonym_replace(text, probability=0.25)
            if variant_text != text and len(variant_text) >= 100:
                augmented.append({
                    **{k: v for k, v in case.items() if k not in ("text", "tokens", "bio_labels", "bio_label_names", "case_id")},
                    "case_id": case["case_id"] * 1000 + len(augmented),
                    "text": variant_text,
                    "tokens": [],  # will be re-tokenized during training
                    "bio_labels": [],
                    "bio_label_names": [],
                    "source": "synonym_aug",
                    "orig_case_id": case["case_id"],
                })

        # Strategy 2: Entity-masked variations (2x per case)
        for variant_text in mask_entity_variations(text):
            if variant_text != text and len(variant_text) >= 100:
                augmented.append({
                    **{k: v for k, v in case.items() if k not in ("text", "tokens", "bio_labels", "bio_label_names", "case_id")},
                    "case_id": case["case_id"] * 1000 + len(augmented),
                    "text": variant_text,
                    "tokens": [],
                    "bio_labels": [],
                    "bio_label_names": [],
                    "source": "entity_mask_aug",
                    "orig_case_id": case["case_id"],
                })

    print(f"  Original: {len(orig_ids)} cases")
    print(f"  After augmentation: {len(augmented)} cases")

    # Strategy 3: Qwen synthetic cases
    if include_synthetic:
        synthetic = generate_synthetic_cases_with_qwen(
            num_per_cause=synthetic_per_cause,
        )
        for i, syn in enumerate(synthetic):
            max_id = max(orig_ids) if orig_ids else 0
            synthetic_id = (max_id + 1) * 1000 + i
            augmented.append({
                "case_id": synthetic_id,
                "file": f"synthetic_{i}.txt",
                "text": syn["text"],
                "tokens": [],
                "bio_labels": [],
                "bio_label_names": [],
                "classification_labels": {
                    "case_cause": [
                        "劳动关系纠纷类", "工伤保险待遇纠纷",
                        "追索劳动报酬", "经济补偿金纠纷",
                        "赔偿金纠纷", "生育保险待遇纠纷",
                    ].index(syn["cause_type"]),
                },
                "numeric_labels": {},
                "rule_elements": {},
                "source": "qwen_synthetic",
                "cause_type": syn["cause_type"],
            })

    # Save augmented dataset
    output = {
        **{k: v for k, v in data.items() if k != "dataset"},
        "augmented": True,
        "total_augmented_cases": len(augmented),
        "dataset": augmented,
    }
    OUTPUT.write_text(json.dumps(output, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nAugmented dataset saved: {OUTPUT}")
    print(f"Total: {len(augmented)} cases")
    return output


if __name__ == "__main__":
    augment_dataset(include_synthetic=True, synthetic_per_cause=5)