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