Implement v1.1 model quality improvements

Address v1 weaknesses: Tier 1 recall 0.722 (target 0.98) and Tier 2
recall 0.934 (target 0.95) caused by 323x entity frequency imbalance.

Changes:
- Tier-weighted CRF loss: sequences with critical PII get higher loss
  weight (3x for Tier 1, 2x for Tier 2, 1.5x for Tier 3)
- Oversampling: duplicate Tier 1 entity examples x3 in training set
  (applied after train/val/test split to avoid data leakage)
- Halved backbone LR (2e-5 -> 1e-5) to prevent late-epoch regression
  observed in v1 at epochs 6-7
- Cosine LR schedule (from linear) for smoother decay
- 15 epochs (from 10) with 500-step warmup
- Entity frequency audit script for data analysis

Preflight passed on RTX 3090: BF16 loss 1166.0 -> 1150.6 (decreasing).

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: Sid Mohan

pii-ner-v1/configs/rtx3090-v1.1.yaml

@@ -0,0 +1,46 @@
+model:
+  backbone: microsoft/deberta-v3-xsmall
+  char_embed_dim: 50
+  char_vocab_size: 256
+  char_cnn_filters: [50, 50, 50]
+  char_cnn_widths: [3, 4, 5]
+  max_char_len: 20
+  dropout: 0.1
+
+data:
+  max_seq_len: 256
+  val_ratio: 0.1
+  test_ratio: 0.1
+  seed: 42
+  # Oversample examples containing Tier 1 entities (SSN, Credit Card, etc.)
+  oversample_tiers: [1]
+  oversample_factor: 3
+
+training:
+  epochs: 15
+  batch_size: 8
+  gradient_accumulation_steps: 4  # effective batch = 32
+  lr_backbone: 1.0e-5             # halved from v1 (2e-5) to prevent late-epoch regression
+  lr_head: 1.0e-3
+  lr_scheduler_type: cosine        # smoother decay than linear
+  warmup_steps: 500                # ~1 epoch warmup at 4237 steps/epoch
+  weight_decay: 0.01
+  eval_strategy: epoch
+  save_strategy: epoch
+  metric_for_best_model: overall_f1
+  logging_steps: 50
+  dataloader_num_workers: 4
+  save_total_limit: 3
+  output_dir: output-v1.1
+  run_name: pii-ner-v1.1-rtx3090
+
+  # Tier-weighted CRF loss: sequences with critical PII get higher loss weight
+  tier_weights:
+    1: 3.0    # SSN, Credit Card, Bank Account, Passport, DL, Tax ID
+    2: 2.0    # Person, Email, Phone, DOB, Address, IP
+    3: 1.5    # Username, Date, Location, Org, URL, etc.
+    4: 1.0    # Domain-specific types
+
+wandb:
+  enabled: true
+  project: datafog-pii-ner

pii-ner-v1/notebooks/full_training.ipynb

@@ -0,0 +1,156 @@
+#!/usr/bin/env python3
+"""Audit entity type frequencies across all training datasets.
+
+Counts entity occurrences per type and per dataset to identify
+class imbalance issues.
+"""
+
+import sys
+from collections import Counter
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
+
+from datasets import load_dataset
+from datafog_pii_ner.data.label_schema import (
+    ALL_ENTITY_TYPES,
+    TIER_1,
+    TIER_2,
+    TIER_3,
+    TIER_4,
+    TIER_MAP,
+    map_label,
+)
+from datafog_pii_ner.data.dataset import DATASET_CONFIGS, _spans_to_bio, _detect_columns, _normalize_tags
+
+import re
+_WORD_RE = re.compile(r"\S+")
+
+
+def count_entities_token_format(raw, dataset_name):
+    """Count entities in token-format dataset (AI4Privacy)."""
+    token_col, tag_col = _detect_columns(raw)
+
+    tag_names = None
+    try:
+        tag_names = raw.features[tag_col].feature.names
+    except Exception:
+        pass
+
+    entity_counts = Counter()
+    examples_with_entity = Counter()
+
+    for i, example in enumerate(raw):
+        tags = _normalize_tags(example[tag_col], dataset_name, tag_names)
+        types_in_example = set()
+        for tag in tags:
+            if tag.startswith("B-"):
+                entity_type = tag[2:]
+                entity_counts[entity_type] += 1
+                types_in_example.add(entity_type)
+        for t in types_in_example:
+            examples_with_entity[t] += 1
+
+        if (i + 1) % 50000 == 0:
+            print(f"  Processed {i+1} examples...")
+
+    return entity_counts, examples_with_entity, len(raw)
+
+
+def count_entities_span_format(raw, dataset_name, text_col, spans_col):
+    """Count entities in span-format dataset (Nemotron, Gretel)."""
+    entity_counts = Counter()
+    examples_with_entity = Counter()
+
+    for i, example in enumerate(raw):
+        _, bio_tags = _spans_to_bio(example[text_col], example[spans_col], dataset_name)
+        types_in_example = set()
+        for tag in bio_tags:
+            if tag.startswith("B-"):
+                entity_type = tag[2:]
+                entity_counts[entity_type] += 1
+                types_in_example.add(entity_type)
+        for t in types_in_example:
+            examples_with_entity[t] += 1
+
+        if (i + 1) % 50000 == 0:
+            print(f"  Processed {i+1} examples...")
+
+    return entity_counts, examples_with_entity, len(raw)
+
+
+def main():
+    total_counts = Counter()
+    total_examples = Counter()
+    dataset_stats = {}
+
+    for name, config in DATASET_CONFIGS.items():
+        print(f"\n{'='*60}")
+        print(f"Loading {name} from {config['path']}...")
+        print(f"{'='*60}")
+
+        raw = load_dataset(config["path"], split=config["split"])
+
+        # Filter to English
+        lang_col = config.get("language_col", "language")
+        lang_val = config.get("language_value", "en")
+        if lang_col in raw.column_names:
+            raw = raw.filter(lambda x: x[lang_col] == lang_val, desc="Filtering English")
+            print(f"  Filtered to {len(raw)} English examples")
+
+        if config.get("format") == "span":
+            counts, ex_counts, n = count_entities_span_format(
+                raw, name, config["text_col"], config["spans_col"]
+            )
+        else:
+            counts, ex_counts, n = count_entities_token_format(raw, name)
+
+        dataset_stats[name] = {
+            "total_examples": n,
+            "entity_counts": counts,
+            "examples_with_entity": ex_counts,
+        }
+
+        total_counts += counts
+        total_examples += ex_counts
+
+        print(f"\n{name}: {n} examples")
+        for etype in sorted(counts.keys(), key=lambda x: counts[x], reverse=True):
+            tier = TIER_MAP.get(etype, "?")
+            print(f"  T{tier} {etype:30s} {counts[etype]:>8,} entities  in {ex_counts[etype]:>8,} examples")
+
+    # Summary
+    print(f"\n{'='*60}")
+    print("COMBINED ENTITY FREQUENCY AUDIT")
+    print(f"{'='*60}")
+
+    total_entities = sum(total_counts.values())
+    print(f"\nTotal entities: {total_entities:,}")
+
+    for tier_name, tier_list in [("TIER 1 (target >=0.98)", TIER_1),
+                                  ("TIER 2 (target >=0.95)", TIER_2),
+                                  ("TIER 3 (target >=0.90)", TIER_3),
+                                  ("TIER 4 (target >=0.85)", TIER_4)]:
+        print(f"\n--- {tier_name} ---")
+        for etype in tier_list:
+            count = total_counts.get(etype, 0)
+            ex_count = total_examples.get(etype, 0)
+            pct = 100.0 * count / total_entities if total_entities > 0 else 0
+            print(f"  {etype:30s} {count:>8,} ({pct:5.2f}%)  in {ex_count:>8,} examples")
+
+    # Imbalance ratio
+    most_common = total_counts.most_common(1)[0] if total_counts else ("?", 0)
+    least_common = total_counts.most_common()[-1] if total_counts else ("?", 0)
+    if least_common[1] > 0:
+        ratio = most_common[1] / least_common[1]
+        print(f"\nImbalance ratio: {most_common[0]}({most_common[1]:,}) / "
+              f"{least_common[0]}({least_common[1]:,}) = {ratio:.0f}x")
+
+    # Types with zero occurrences
+    zero_types = [t for t in ALL_ENTITY_TYPES if total_counts.get(t, 0) == 0]
+    if zero_types:
+        print(f"\nZERO OCCURRENCE TYPES: {zero_types}")
+
+
+if __name__ == "__main__":
+    main()

pii-ner-v1/scripts/entity_audit.py

@@ -37,7 +37,7 @@
 
 from datafog_pii_ner.data.collator import PiiDataCollator
 from datafog_pii_ner.data.dataset import load_pii_datasets
-from datafog_pii_ner.data.label_schema import NUM_LABELS
+from datafog_pii_ner.data.label_schema import NUM_LABELS, build_label_weights
 from datafog_pii_ner.model.pii_model import PiiNerConfig, PiiNerModel
 from datafog_pii_ner.training.metrics import compute_metrics
 from datafog_pii_ner.training.train import PiiTrainer
@@ -241,13 +241,16 @@ def train(config: dict, precision: dict, logger: logging.Logger, status: StatusT
     # --- Data ---
     logger.info("Loading datasets (may take a few minutes on first run)...")
     status.update(stage="loading_data")
+    data_cfg = config["data"]
     datasets = load_pii_datasets(
         tokenizer=tokenizer,
-        max_seq_len=config["data"]["max_seq_len"],
+        max_seq_len=data_cfg["max_seq_len"],
         max_char_len=config["model"]["max_char_len"],
-        val_ratio=config["data"]["val_ratio"],
-        test_ratio=config["data"]["test_ratio"],
-        seed=config["data"].get("seed", 42),
+        val_ratio=data_cfg["val_ratio"],
+        test_ratio=data_cfg["test_ratio"],
+        seed=data_cfg.get("seed", 42),
+        oversample_tiers=data_cfg.get("oversample_tiers"),
+        oversample_factor=data_cfg.get("oversample_factor", 2),
     )
     logger.info(f"Train: {len(datasets['train']):,}  Val: {len(datasets['validation']):,}  "
                 f"Test: {len(datasets['test']):,}")
@@ -271,6 +274,17 @@ def train(config: dict, precision: dict, logger: logging.Logger, status: StatusT
     trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
     logger.info(f"Parameters: {param_count:,} total, {trainable:,} trainable")
 
+    # --- Tier-weighted loss ---
+    tier_weights_cfg = config.get("training", {}).get("tier_weights")
+    if tier_weights_cfg:
+        # Convert string keys from YAML to ints
+        tw = {int(k): float(v) for k, v in tier_weights_cfg.items()}
+        weights = build_label_weights(tw)
+        model.crf_head.set_label_weights(torch.tensor(weights, dtype=torch.float32))
+        logger.info(f"Tier-weighted CRF loss: {tw}")
+    else:
+        logger.info("CRF loss: uniform (no tier weighting)")
+
     # --- Collator ---
     collator = PiiDataCollator(tokenizer=tokenizer, max_char_len=config["model"]["max_char_len"])
 
@@ -285,14 +299,22 @@ def train(config: dict, precision: dict, logger: logging.Logger, status: StatusT
     wandb_cfg = config.get("wandb", {})
     report_to = "wandb" if wandb_cfg.get("enabled", False) else "none"
 
+    # Warmup: prefer warmup_steps over warmup_ratio (ratio is deprecated in HF 5.2+)
+    warmup_kwargs = {}
+    if "warmup_steps" in tcfg:
+        warmup_kwargs["warmup_steps"] = tcfg["warmup_steps"]
+    elif "warmup_ratio" in tcfg:
+        warmup_kwargs["warmup_ratio"] = tcfg["warmup_ratio"]
+
     training_args = TrainingArguments(
         output_dir=output_dir,
         num_train_epochs=tcfg["epochs"],
         per_device_train_batch_size=tcfg["batch_size"],
         per_device_eval_batch_size=tcfg["batch_size"],
         gradient_accumulation_steps=tcfg.get("gradient_accumulation_steps", 1),
         learning_rate=tcfg["lr_backbone"],
-        warmup_ratio=tcfg.get("warmup_ratio", 0.1),
+        lr_scheduler_type=tcfg.get("lr_scheduler_type", "linear"),
+        **warmup_kwargs,
         weight_decay=tcfg.get("weight_decay", 0.01),
         fp16=precision["fp16"],
         bf16=precision["bf16"],

pii-ner-v1/scripts/run_training.py

@@ -17,7 +17,7 @@
 from transformers import PreTrainedTokenizerFast
 
 from .char_vocab import text_to_char_ids
-from .label_schema import LABEL_TO_ID, map_label
+from .label_schema import LABEL_TO_ID, TIER_1, TIER_2, map_label
 
 logger = logging.getLogger(__name__)
 
@@ -333,6 +333,62 @@ def preprocess_token(example):
     return processed
 
 
+def _build_oversample_label_ids(tiers: list[int]) -> set[int]:
+    """Build set of label IDs that belong to the specified tiers."""
+    tier_types = []
+    if 1 in tiers:
+        tier_types.extend(TIER_1)
+    if 2 in tiers:
+        tier_types.extend(TIER_2)
+
+    label_ids = set()
+    for etype in tier_types:
+        b_key = f"B-{etype}"
+        i_key = f"I-{etype}"
+        if b_key in LABEL_TO_ID:
+            label_ids.add(LABEL_TO_ID[b_key])
+        if i_key in LABEL_TO_ID:
+            label_ids.add(LABEL_TO_ID[i_key])
+    return label_ids
+
+
+def _oversample_rare_entities(
+    dataset: Dataset,
+    oversample_tiers: list[int],
+    oversample_factor: int,
+) -> Dataset:
+    """Oversample examples containing rare entity types.
+
+    Identifies examples that contain entities from the specified tiers,
+    then duplicates them `oversample_factor` times.
+    """
+    target_label_ids = _build_oversample_label_ids(oversample_tiers)
+    if not target_label_ids:
+        return dataset
+
+    # Find indices of examples containing target entities
+    rare_indices = []
+    for i in range(len(dataset)):
+        labels = dataset[i]["labels"]
+        if any(label_id in target_label_ids for label_id in labels):
+            rare_indices.append(i)
+
+    if not rare_indices:
+        logger.warning("No examples found containing target tier entities for oversampling")
+        return dataset
+
+    logger.info(
+        f"Oversampling: {len(rare_indices)} examples with tier {oversample_tiers} entities "
+        f"x{oversample_factor} ({len(rare_indices) * oversample_factor} additional examples)"
+    )
+
+    # Create duplicated dataset
+    oversampled_indices = rare_indices * oversample_factor
+    oversampled = dataset.select(oversampled_indices)
+
+    return concatenate_datasets([dataset, oversampled])
+
+
 def load_pii_datasets(
     tokenizer: PreTrainedTokenizerFast,
     max_seq_len: int = 256,
@@ -341,8 +397,16 @@ def load_pii_datasets(
     val_ratio: float = 0.1,
     test_ratio: float = 0.1,
     seed: int = 42,
+    oversample_tiers: list[int] | None = None,
+    oversample_factor: int = 2,
 ) -> DatasetDict:
-    """Load all PII datasets, combine, and split into train/val/test."""
+    """Load all PII datasets, combine, and split into train/val/test.
+
+    Args:
+        oversample_tiers: List of tier numbers to oversample (e.g., [1, 2]).
+            None disables oversampling.
+        oversample_factor: Number of times to duplicate rare-entity examples.
+    """
     all_datasets = []
     for name in DATASET_CONFIGS:
         try:
@@ -361,16 +425,23 @@ def load_pii_datasets(
     combined = concatenate_datasets(all_datasets)
     logger.info(f"Combined: {len(combined)} examples")
 
-    # Split: train / val / test
+    # Split: train / val / test (BEFORE oversampling to avoid data leakage)
     test_size = test_ratio
     val_size = val_ratio / (1 - test_ratio)  # Adjust for two-stage split
 
     split1 = combined.train_test_split(test_size=test_size, seed=seed)
     split2 = split1["train"].train_test_split(test_size=val_size, seed=seed)
 
+    train_ds = split2["train"]
+
+    # Oversample rare entities in training set only
+    if oversample_tiers:
+        train_ds = _oversample_rare_entities(train_ds, oversample_tiers, oversample_factor)
+        logger.info(f"Train after oversampling: {len(train_ds)} examples")
+
     return DatasetDict(
         {
-            "train": split2["train"],
+            "train": train_ds,
             "validation": split2["test"],
             "test": split1["test"],
         }