#!/usr/bin/env python3 """ Training script for Chinese RoBERTa Multi-Task Labor Arbitration Extractor. Usage: cd nlp-service/training PYTHONPATH="..;../../backend" python train_bert.py \ --dataset ../../backend/data/augmented_dataset.json \ --output ../models/chinese_roberta_labor_extractor \ --epochs 10 --batch_size 8 --lr 2e-5 """ from __future__ import annotations import argparse import json import os import random import sys from pathlib import Path from typing import Any # NOTE: transformers before torch on Windows from transformers import ( AutoTokenizer, get_linear_schedule_with_warmup, ) import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from tqdm import tqdm # Add parent dirs for imports sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from app.services.bert_multi_task_model import ChineseRobertaMultiTask from training.bio_schema import LABEL2ID, NUM_LABELS class LaborCaseDataset(Dataset): """PyTorch Dataset for labor arbitration case element extraction.""" def __init__( self, cases: list[dict], tokenizer: AutoTokenizer, max_length: int = 512, for_training: bool = True, ): self.cases = cases self.tokenizer = tokenizer self.max_length = max_length self.for_training = for_training def __len__(self) -> int: return len(self.cases) def __getitem__(self, idx: int) -> dict[str, Any]: case = self.cases[idx] text = case["text"] cls_labels = case.get("classification_labels", {}) num_labels = case.get("numeric_labels", {}) # Tokenize encoding = self.tokenizer( text, truncation=True, padding="max_length", max_length=self.max_length, return_tensors="pt", return_offsets_mapping=True, ) input_ids = encoding["input_ids"].squeeze(0) # [seq_len] attention_mask = encoding["attention_mask"].squeeze(0) # [seq_len] offset_mapping = encoding["offset_mapping"].squeeze(0) # [seq_len, 2] result = { "input_ids": input_ids, "attention_mask": attention_mask, "case_id": case.get("case_id", 0), } if self.for_training: # BIO labels: align character-level labels to subword tokens bio_labels = case.get("bio_labels", []) tokens = case.get("tokens", []) if bio_labels and tokens: token_labels = self._align_bio_to_tokens( bio_labels, tokens, encoding, offset_mapping ) else: token_labels = [-100] * self.max_length result["ner_labels"] = torch.tensor(token_labels, dtype=torch.long) # Classification labels result["cause_label"] = torch.tensor( cls_labels.get("case_cause", 0), dtype=torch.long ) result["contract_label"] = torch.tensor( cls_labels.get("contract_type", 3), dtype=torch.long ) result["employment_label"] = torch.tensor( cls_labels.get("employment_type", 4), dtype=torch.long ) # Boolean labels (9 fields) bool_field_keys = [ "bool_lr1_contract_signed", "bool_lr1_open_ended_contract", "bool_lr1_double_wage_no_contract", "bool_lr1_si_joined", "bool_wi_si_joined", "bool_wi_recognize_ec", "bool_dm_contract_exists", "bool_dm_contract_continue", "bool_mi_contract_continue", ] bool_vals = [cls_labels.get(k, 2) for k in bool_field_keys] result["bool_labels"] = torch.tensor(bool_vals, dtype=torch.long) # Numeric labels: use first numeric field (month_salary) as primary target reg_val = num_labels.get("month_salary", None) result["reg_value"] = torch.tensor( float(reg_val) if reg_val is not None else -1.0, dtype=torch.float, ) result["reg_mask"] = torch.tensor( 1.0 if reg_val is not None else 0.0, dtype=torch.float ) return result def _align_bio_to_tokens( self, char_bio_labels: list[int], char_tokens: list[str], encoding: Any, offset_mapping: torch.Tensor, ) -> list[int]: """ Align character-level BIO labels to subword token labels. For each subword token, find the corresponding character position and use the BIO label from that character. """ # Build a mapping from char position → BIO label char_to_label: dict[int, int] = {} char_pos = 0 for token, label in zip(char_tokens, char_bio_labels): for _ in token: char_to_label[char_pos] = label char_pos += 1 # Map subword tokens to BIO labels based on their char offset token_labels = [] for i, (start, end) in enumerate(offset_mapping.tolist()): if i >= self.max_length: break if start == 0 and end == 0: # Special tokens ([CLS], [SEP], [PAD]) token_labels.append(-100) elif start >= char_pos: # Token is beyond our character annotations token_labels.append(-100) else: # Use label from the first character of this token label = char_to_label.get(start, 0) # 0 = O token_labels.append(label) # Pad to max_length while len(token_labels) < self.max_length: token_labels.append(-100) return token_labels[:self.max_length] def collate_batch(batch: list[dict]) -> dict[str, torch.Tensor]: """Custom collate function for multi-task training.""" return { "input_ids": torch.stack([b["input_ids"] for b in batch]), "attention_mask": torch.stack([b["attention_mask"] for b in batch]), "ner_labels": torch.stack([b["ner_labels"] for b in batch]), "cause_label": torch.stack([b["cause_label"] for b in batch]), "contract_label": torch.stack([b["contract_label"] for b in batch]), "employment_label": torch.stack([b["employment_label"] for b in batch]), "bool_labels": torch.stack([b["bool_labels"] for b in batch]), "reg_values": torch.stack([b["reg_value"] for b in batch]), "reg_mask": torch.stack([b["reg_mask"] for b in batch]), } def train_epoch( model: ChineseRobertaMultiTask, dataloader: DataLoader, optimizer: torch.optim.Optimizer, scheduler: Any, device: torch.device, epoch: int, ) -> float: """Train one epoch. Returns average loss.""" model.train() total_loss = 0.0 num_batches = 0 pbar = tqdm(dataloader, desc=f"Epoch {epoch}") for batch in pbar: batch = {k: v.to(device) for k, v in batch.items()} optimizer.zero_grad() outputs = model(**batch) loss = outputs["loss"] if loss is not None: loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) optimizer.step() scheduler.step() total_loss += loss.item() num_batches += 1 pbar.set_postfix({"loss": f"{loss.item():.4f}"}) return total_loss / max(num_batches, 1) @torch.no_grad() def evaluate( model: ChineseRobertaMultiTask, dataloader: DataLoader, device: torch.device, ) -> dict[str, float]: """Evaluate on dev set. Returns per-task loss.""" model.eval() total_loss = 0.0 num_batches = 0 cause_correct = 0 cause_total = 0 for batch in tqdm(dataloader, desc="Eval"): batch = {k: v.to(device) for k, v in batch.items()} outputs = model(**batch) loss = outputs["loss"] if loss is not None: total_loss += loss.item() num_batches += 1 # Cause accuracy cause_preds = torch.argmax(outputs["cause_logits"], dim=-1) cause_correct += (cause_preds == batch["cause_label"]).sum().item() cause_total += batch["cause_label"].size(0) return { "loss": total_loss / max(num_batches, 1), "cause_accuracy": cause_correct / max(cause_total, 1), } def create_data_splits(dataset_path: str, train_ratio: float = 0.7, dev_ratio: float = 0.15): """Split dataset into train/dev/test.""" with open(dataset_path, encoding="utf-8") as f: data = json.load(f) cases = data["dataset"] # Separate original vs augmented originals = [c for c in cases if "source" not in c] augmenteds = [c for c in cases if "source" in c] random.shuffle(originals) random.shuffle(augmenteds) n_orig = len(originals) n_train_orig = max(1, int(n_orig * train_ratio)) n_dev_orig = max(1, int(n_orig * dev_ratio)) train_cases = originals[:n_train_orig] + augmenteds dev_cases = originals[n_train_orig:n_train_orig + n_dev_orig] test_cases = originals[n_train_orig + n_dev_orig:] print(f"Split: train={len(train_cases)} ({n_train_orig} orig + {len(augmenteds)} aug), " f"dev={len(dev_cases)}, test={len(test_cases)}") return train_cases, dev_cases, test_cases def main(): parser = argparse.ArgumentParser(description="Train BERT multi-task extractor") parser.add_argument("--dataset", required=True, help="Path to augmented_dataset.json") parser.add_argument("--output", required=True, help="Output directory for trained model") parser.add_argument("--epochs", type=int, default=10, help="Number of training epochs") parser.add_argument("--batch_size", type=int, default=8, help="Training batch size") parser.add_argument("--lr", type=float, default=2e-5, help="Learning rate") parser.add_argument("--max_length", type=int, default=512, help="Max sequence length") parser.add_argument("--model_name", default="hfl/chinese-roberta-wwm-ext", help="Pretrained model name or path") parser.add_argument("--patience", type=int, default=3, help="Early stopping patience") parser.add_argument("--gradient_accumulation", type=int, default=2, help="Gradient accumulation steps") parser.add_argument("--warmup_ratio", type=float, default=0.1, help="Warmup ratio") parser.add_argument("--no_cuda", action="store_true", help="Disable CUDA") args = parser.parse_args() # Device device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") print(f"Using device: {device}") # Load and split data train_cases, dev_cases, test_cases = create_data_splits(args.dataset) # Tokenizer tokenizer = AutoTokenizer.from_pretrained(args.model_name) print(f"Loaded tokenizer: {args.model_name}") # Datasets train_dataset = LaborCaseDataset(train_cases, tokenizer, args.max_length, for_training=True) dev_dataset = LaborCaseDataset(dev_cases, tokenizer, args.max_length, for_training=True) test_dataset = LaborCaseDataset(test_cases, tokenizer, args.max_length, for_training=True) train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_batch, num_workers=0, ) dev_loader = DataLoader( dev_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_batch, num_workers=0, ) # Model model = ChineseRobertaMultiTask(model_name=args.model_name) model.to(device) print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}") # Optimizer no_decay = ["bias", "LayerNorm.weight"] optimizer_grouped = [ { "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], "weight_decay": 0.01, }, { "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)], "weight_decay": 0.0, }, ] optimizer = torch.optim.AdamW(optimizer_grouped, lr=args.lr) # Scheduler total_steps = len(train_loader) * args.epochs // args.gradient_accumulation warmup_steps = int(total_steps * args.warmup_ratio) scheduler = get_linear_schedule_with_warmup( optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, ) # Training loop best_dev_loss = float("inf") patience_counter = 0 for epoch in range(1, args.epochs + 1): train_loss = train_epoch(model, train_loader, optimizer, scheduler, device, epoch) dev_metrics = evaluate(model, dev_loader, device) print(f"Epoch {epoch:>2}: train_loss={train_loss:.4f}, " f"dev_loss={dev_metrics['loss']:.4f}, " f"dev_cause_acc={dev_metrics['cause_accuracy']:.3f}") # Early stopping if dev_metrics["loss"] < best_dev_loss: best_dev_loss = dev_metrics["loss"] patience_counter = 0 os.makedirs(args.output, exist_ok=True) model.save_pretrained(args.output) tokenizer.save_pretrained(args.output) print(f" -> Saved best model to {args.output}") else: patience_counter += 1 if patience_counter >= args.patience: print(f"Early stopping at epoch {epoch}") break print(f"\nTraining complete. Best dev loss: {best_dev_loss:.4f}") print(f"Model saved to: {args.output}") if __name__ == "__main__": main()