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

#!/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()