""" 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)