PII extraction experiments with GLiNER (v1) + GLiNER2 (v2)

This is a small, config-driven experiment harness to compare four training setups ("instances") for long-span NER / PII extraction.

It is designed around the dataset:

• arthrod/amd-pii-generator-dataprivateDataset (Hugging Face)

What you get

• scripts/prepare_data.py
  Downloads the HF dataset and produces canonical JSONL plus derived datasets for:

  • GLiNER v1 (tokenized_text + ner spans)
  • GLiNER2 (text + entities dict + optional classification tasks)

• scripts/train_gliner2.py
  Trains GLiNER2 using the official gliner2.training.* APIs (supports LoRA via config).

• scripts/train_gliner1.py
  Trains GLiNER v1 using gliner.training.Trainer and DataCollator.

• scripts/evaluate.py
  Runs inference + writes JSONL predictions + computes a big pile of metrics (strict + overlap + fuzzy), logs everything to W&B, and writes metrics_summary.json.

Quickstart

0) Install deps

pip install -r requirements.txt

1) Prepare datasets

python scripts/prepare_data.py --config configs/instance4_gliner2_lora_multitask.yaml

This writes:

• data/canonical/{train,val,test}.jsonl
• data/gliner1/{instance_name}/{train,val,test}.jsonl
• data/gliner2/{instance_name}/{train,val,test}.jsonl

2) Train

GLiNER2 (Instance 4 example):

accelerate launch scripts/train_gliner2.py --config configs/instance4_gliner2_lora_multitask.yaml

GLiNER v1 (Instance 1 example):

accelerate launch scripts/train_gliner1.py --config configs/instance1_gliner1_short.yaml

3) Evaluate

python scripts/evaluate.py --config configs/instance4_gliner2_lora_multitask.yaml --checkpoint runs/instance4_gliner2_lora_multitask/model

Outputs are saved under runs/<run_name>/.

Notes for ROCm / AMD GPUs

• Prefer bf16 if supported (MI200/MI300 class). Otherwise fall back to fp16 or full fp32.
• accelerate is the simplest launcher for multi-GPU on both CUDA and ROCm.
• If you hit instability, set:
  • train.precision: fp32
  • train.gradient_checkpointing: true
  • smaller batch size + higher grad accumulation.

Modal fallback

If your container has any GPU/driver weirdness, take the exact same scripts and run them on Modal. All configs are self-contained; you typically only need to change:

• runtime.output_root
• wandb.* settings
• dataset cache location via HF_HOME