README.md

---
tags:
  - modernbert
  - ner
  - medical
  - pytorch
  - dataset-extraction
  - token-classification
  - long-context
license: apache-2.0
---

# ModernBERT-large - Medical Dataset Name Extraction

A **ModernBERT-large** token classification model trained to extract dataset names from medical/scientific research papers. Uses `answerdotai/ModernBERT-large` as the encoder with a linear classification head and cross-entropy loss. Trained with an **8192-token context window** — 16x longer than the SciBERT/CRF baselines (512 tokens).

## Architecture

- **Encoder:** ModernBERT-large (`answerdotai/ModernBERT-large`), 395M params, hidden size 1024, 8192-token context
- **Classification head:** Linear(1024, 3) with CrossEntropyLoss
- **Mixed precision:** bf16 on Blackwell/Ampere GPUs, fp16 + GradScaler fallback on Turing/Volta
- **Gradient checkpointing** enabled to fit ModernBERT-large at 8K context on 16 GB GPUs
- **Differential LR:** 5e-5 encoder / 1e-3 head, warmup 10%, AdamW

## Entity Types

| Entity Type | Description |
|---|---|
| Dataset | Names of datasets, corpus, collections, databases, benchmarks used in scientific research |

## Performance

### Validation Set
| Metric | Precision | Recall | F1 |
|---|---|---|---|
| seqeval (entity) | - | - | 0.9025 |
| Exact Match (chunk) | 0.8689 | 0.9149 | 0.8913 |
| Partial Match (chunk) | 0.9379 | 0.9877 | 0.9622 |

### Test Set (In-Distribution)
| Metric | Precision | Recall | F1 |
|---|---|---|---|
| seqeval (entity) | - | - | 0.9124 |
| Exact Match (chunk) | 0.8815 | 0.9222 | 0.9014 |
| Partial Match (chunk) | 0.9551 | 0.9992 | 0.9766 |

### OOD Set (Out-of-Distribution)
| Metric | Precision | Recall | F1 |
|---|---|---|---|
| seqeval (entity) | - | - | 0.6673 |
| Exact Match (chunk) | 0.6584 | 0.7204 | 0.6880 |
| Partial Match (chunk) | 0.8290 | 0.9071 | 0.8663 |

## Training Details

| Parameter | Value |
|---|---|
| Base model | answerdotai/ModernBERT-large (395M) |
| Max sequence length | 8192 tokens |
| Primary chunk size | 6000 chars (~1800 BPE tokens) |
| Chunk overlap | 500 chars |
| Negative sampling ratio | 0.15 |
| Entity-centered augmentation | window=4000 chars, max 3/doc |
| Training documents | 542 (+ 151 OOD held out) |
| Training examples | 2324 |
| Effective batch size | 4 (batch 1 x accum 4) |
| Epochs | 10 |
| Precision | bf16 |
| GPU used | NVIDIA GeForce RTX 5070 Ti |

## Comparison vs Baselines

All three models (CRF, SciBERT, ModernBERT-large) were trained on the same 542 manually
annotated medical-dataset-mention documents using the same random seed (42), and evaluated
on the same 151-document OOD set. All pipelines apply **negative sampling (ratio 0.15)** and
**entity-centered augmentation** for class balance. ModernBERT uses **6000-char primary
chunks** and **4000-char augmentation windows** to exploit its 8K context window — so chunk
counts differ numerically from the 512-token baselines, but the underlying documents, seed,
and balancing techniques are identical. Comparison is meaningful at the document / entity F1
level (see `eval_metrics.json`).

## Usage

```python
from inference import load_model, predict

model, tokenizer, id2label, config = load_model(".", device="cuda")
text = "We evaluated our model on the MIMIC-III dataset and the PhysioNet challenge corpus."
entities = predict(text, model, tokenizer, id2label, device="cuda")
for ent in entities:
    print(ent)
```

## Dependencies

```
torch>=2.1
transformers>=4.48   # required for native ModernBERT support
seqeval
```

## Limitations

- Recognizes only dataset / corpus / database / benchmark names, not other biomedical entities
- Although the encoder supports 8192 tokens, inputs are chunked at 6000 chars during inference to match training; cross-chunk entities are deduplicated by span
- Trained on manually annotated medical research abstracts and full-text sections; generalization to other scientific domains is not guaranteed
- Long-document inference is memory-heavy on small GPUs — use a quantized build or CPU-offload for the largest inputs