finetune-ner

name	finetune-ner
description	Fine-tune Thai NER transformer models (PhayaThaiBERT, WangchanBERTa, or new checkpoints). Covers full pipeline — data prep, config generation, hyperparameter tuning, training, evaluation, comparison, and deployment. Use when the user wants to train, retrain, tune, or experiment with NER models.
allowed-tools	Read, Edit, Write, Bash(python3 *), Bash(source *), Bash(cat *), Bash(ls *), Bash(diff *), Bash(wc *), Bash(git *), Grep, Glob
argument-hint	[model-name-or-checkpoint] [optional: experiment-description]

Fine-Tune Thai NER Model

You are an expert NER fine-tuning assistant for this Thai job-post entity extraction project. Follow the workflow below to guide the user through a complete fine-tuning run.

Context

• Project: Thai Job NLP NER — extracts 7 entity types from Thai job posts
• Best model: PhayaThaiBERT
• Key finding (Sprint 6): Synthetic F1=0.956 but real-world F1=0.143 — massive domain gap. Real-world data is essential.
• Hardware: Apple Silicon MPS (FP32 only, no bf16)
• Datasets:
  • Synthetic: 1,253 posts in data/raw/synthetic*.json → data/processed/ (1,001/126/126 split)
  • Real-world: 54 posts in data/raw/prod-secret/real_export.jsonl → data/processed/prod-secret/ (37/8/8 split)
  • Mixed: combine both sources for training
• Labels: O, B/I-{HARD_SKILL, PERSON, LOCATION, COMPENSATION, EMPLOYMENT_TERMS, CONTACT, DEMOGRAPHIC}
• Training entry point: python3 -m src.training.train_ner --config <yaml> --output results
• CRITICAL: Always evaluate on BOTH synthetic and real-world test sets. Real-world F1 is the metric that matters.

Workflow

Step 1: Clarify the Experiment

Ask the user (if not already clear from $ARGUMENTS):

1. Which base model? PhayaThaiBERT (clicknext/phayathaibert), WangchanBERTa (airesearch/wangchanberta-base-att-spm-uncased), or a new HuggingFace checkpoint?
2. What's the goal? Beat current F1=0.956? Improve a weak entity (DEMOGRAPHIC F1=0.915)? Test a new technique? Speed up training?
3. Any specific technique to try? (LoRA, CRF layer, progressive unfreezing, new optimizer, data augmentation, etc.)

Step 2: Generate or Modify Config

Create a new config YAML in configs/ based on the experiment. Use the reference doc for recommended hyperparameters.

Config template (adjust per model):

model:
  checkpoint: "<checkpoint>"
  num_labels: 15
  max_length: 512

labels:
  - "O"
  - "B-HARD_SKILL"
  - "I-HARD_SKILL"
  - "B-PERSON"
  - "I-PERSON"
  - "B-LOCATION"
  - "I-LOCATION"
  - "B-COMPENSATION"
  - "I-COMPENSATION"
  - "B-EMPLOYMENT_TERMS"
  - "I-EMPLOYMENT_TERMS"
  - "B-CONTACT"
  - "I-CONTACT"
  - "B-DEMOGRAPHIC"
  - "I-DEMOGRAPHIC"

training:
  learning_rate: 3.0e-5
  batch_size: <2 for large vocab, 8 for 25K vocab>
  gradient_accumulation_steps: <8 for batch=2, 2 for batch=8>
  epochs: 15
  warmup_ratio: 0.1
  weight_decay: 0.01
  fp16: false
  bf16: false
  eval_strategy: "epoch"
  save_strategy: "epoch"
  metric_for_best_model: "f1"
  mps_high_watermark_ratio: "<0.0 for large vocab, 1.5 for small>"
  report_to: "none"
  early_stopping_patience: 3
  gradient_checkpointing: <true for large vocab>
  freeze_embeddings: <true for large vocab>

data:
  raw_dir: "data/raw"
  processed_dir: "data/processed"
  test_size: 0.1
  val_size: 0.1
  seed: 42

alignment:
  fuzzy_threshold: 85.0
  tcc_engine: "tcc"

Critical rules for config generation:

• Vocab > 100K tokens → MUST set batch_size: 2, gradient_accumulation_steps: 8, freeze_embeddings: true, gradient_checkpointing: true, mps_high_watermark_ratio: "0.0"
• Vocab < 50K tokens → can use batch_size: 8, gradient_accumulation_steps: 2, no freeze needed
• NEVER set fp16: true or bf16: true on MPS — causes NaN loss
• Always keep warmup_ratio: 0.1 — removing it drops F1 by ~4.3 points
• Effective batch size should be 16 (batch_size × gradient_accumulation_steps)

Step 3: Prepare Data (if needed)

If the user has new data or wants to rebuild the dataset:

source .venv/bin/activate
# Build IOB2-aligned dataset
python3 -c "
from src.alignment.iob2_formatter import run_pipeline
run_pipeline(checkpoint='<checkpoint>')
"

Verify dataset:

python3 -c "
from datasets import DatasetDict
ds = DatasetDict.load_from_disk('data/processed')
print(f'train={len(ds[\"train\"])}, val={len(ds[\"validation\"])}, test={len(ds[\"test\"])}')
"

Step 4: Train

source .venv/bin/activate
python3 -m src.training.train_ner --config configs/<config_file>.yaml --output results

Monitor for:

• OOM errors → reduce batch_size, enable freeze_embeddings and gradient_checkpointing
• NaN loss → check fp16/bf16 are false, reduce learning_rate
• Flat F1 → increase learning_rate, check warmup_ratio is set
• Early stopping triggered too soon → increase patience to 5
• Training too slow → this is expected on MPS with PhayaThaiBERT (~1.5 it/s, ~10 min total)

Step 5: Evaluate

After training completes, run per-entity evaluation:

source .venv/bin/activate
python3 -m src.evaluation.per_entity_report \
  --model-dir results/<model_name>/final \
  --dataset-dir data/processed \
  --output-dir results/<model_name>

Step 6: Compare with Baseline

python3 scripts/compare_models.py results/<baseline_dir> results/<new_dir> \
  --names Baseline NewExperiment

Baselines to compare against:

• PhayaThaiBERT: results/phayathaibert (F1=0.956)
• WangchanBERTa v3: results_v3 (F1=0.897)

Step 7: Report Results

Present results as a markdown table:

Metric	Baseline	Experiment	Delta
Overall F1	...	...	...
Per-entity F1	...	...	...

Highlight:

• Any entity with F1 improvement > 0.01
• Any entity with F1 regression > 0.005
• Training time comparison
• Memory usage differences

Step 8: Next Steps

Based on results, suggest:

• If F1 improved → save config, consider deploying
• If F1 regressed → diagnose (check learning curves, confusion matrix)
• If marginal change → suggest a different technique from the reference doc

Important Constraints

1. DATA PROTECTION: NEVER display contents of data/raw/ner_export.json or any **/prod-secret/ files
2. MPS ONLY: This project trains on Apple Silicon. Never suggest CUDA-specific features
3. FP32 ONLY: Mixed precision does not work reliably on MPS
4. Existing pipeline: Prefer modifying configs over rewriting train_ner.py. If code changes are needed (e.g., adding CRF layer, LoRA), make them backward-compatible
5. Reproducibility: Always use seed: 42 for data splits

Advanced Techniques Reference

See reference.md for detailed guidance on:

• LoRA / PEFT for token classification
• CRF decoder layer
• APOLLO optimizer (memory-efficient AdamW alternative)
• Schedule-Free optimizers
• Progressive unfreezing & discriminative learning rates
• Focal loss for class imbalance
• Data augmentation strategies
• Ensemble methods
• R-Drop regularization