PALAUNER

Reference implementation for PALAUNER, a reinforcement-learning policy that selects which augmented sentences to keep when training a few-shot biomedical NER model. A DQN agent is trained over a pool of augmented candidates and incrementally adds the most useful ones to the training set, instead of using the whole augmented pool indiscriminately.

This repository accompanies the paper "Palauner: policy-based active learning to augment named entity recognition datasets" (Data Mining and Knowledge Discovery, 2026, doi:10.1007/s10618-026-01239-2) and is provided for reproducibility. See Citation below.

Repository layout

Path	Description
main.py	Entry point: few-shot sampling, augmentation, DQN training loop, evaluation.
palauner.py	The DQN agent and the CNN state encoders (sentence_cnn, prob_cnn, DQN_PAL).
augmentation.py	Augmentation backends (see below).
utils.py	Few-shot sampling, prediction normalization, state construction.
scripts/	Batch drivers used to run the experiment matrix and aggregate results.
data/	The five NER datasets in a common JSON format.

Datasets

Five biomedical NER corpora are included under data/, each with train.json / valid.json / test.json: BC5CDR, NCBI Disease, BC2GM, Linnaeus, CRAFT.

Augmentation backends

Selected with --data_augmentation_method:

Value	Method
random_mention_replacement	Replace entity mentions with same-type mentions drawn from the training set.
cosiner	COSINER-style similarity-driven mention replacement.
melm	Masked Entity Language Modeling (requires the MELM code, see below).
gpt	OpenAI generative augmentation (requires --openai_api_key).
llama3	Local Llama-3-8B-Instruct generative augmentation.
qwen25	Local Qwen-2.5-7B-Instruct generative augmentation.

The llama3 / qwen25 backends download weights from the HuggingFace Hub; set HF_HOME if you need the cache somewhere other than ~/.cache/huggingface.

The melm backend is not self-contained: third-party code is not redistributed here. MELMAugmenter shells out into a MELM/ directory at the repository root and invokes its 01_train.sh / 02_generate.sh. To use it, place the reference implementation of MELM: Data Augmentation with Masked Entity Language Modeling for Low-Resource NER (Zhou et al., ACL 2022) in MELM/. All other backends work out of the box.

Installation

pip install -r requirements.txt

requirements.txt pins the exact versions used for the reported results, including a CUDA 11.8 build of PyTorch 2.1.0 for Linux / Python 3.10. On other platforms, install a matching torch build for your CUDA version instead.

Quick start

A single run on BC5CDR with 100-shot training data:

python main.py \
    --data_path 'data/BC5CDR' \
    --model_name "dmis-lab/biobert-base-cased-v1.2" \
    --episodes 5 \
    --few_shot_size 100 \
    --evaluate \
    --evaluate_episode \
    --verbose

--evaluate reports three numbers on the test set: the few-shot-only baseline, training on the full augmented pool, and training on the PALAUNER-selected subset. --evaluate_episode additionally reports PALAUNER after every episode.

Useful flags: --seed, --augmentation_k (candidates generated per sentence), --eval_mode {valid,test}, --results_output_dir, --training_output_dir. Run python main.py --help for the full list.

Reproducing the paper experiments

The LLM-augmentation experiments are run in two stages, so the LLM is loaded once per model rather than once per configuration.

1. Pre-generate the augmented pools (writes cache/pools/<llm>/<dataset>/k{k}_seed{seed}/{train,valid,pool}.json):

python scripts/generate_pools.py --llms llama3 qwen25

2. Run the training matrix over datasets x few-shot sizes x seeds, with one baseline and one PALAUNER run per cell:

bash scripts/run_matrix.sh              # full matrix
DRY_RUN=1 bash scripts/run_matrix.sh    # print the commands only
LLMS="llama3" bash scripts/run_matrix.sh

The driver dispatches jobs across the GPUs in GPUS (default 1 2) and is configurable through environment variables: CACHE_ROOT, RESULTS_ROOT, TRAIN_TMP_ROOT, DATASETS, FEW_SHOTS, SEEDS, EPISODES, AUG_K. Defaults for CONDA_PREFIX_PATH and HF_HOME point at the machine the experiments were run on — override them for your environment.

3. Aggregate per-run JSONs into the reported tables:

python scripts/aggregate_results.py

The no-augmentation baseline at k=100 is produced separately by scripts/run_baseline_matrix.sh (driving scripts/run_baseline.py).

Citation

If you use this code, please cite:

@article{DBLP:journals/datamine/PostiglioneVSSM26,
  author       = {Marco Postiglione and
                  Andrea Vignali and
                  Giancarlo Sperl{\'{\i}} and
                  Guido Secondulfo and
                  Vincenzo Moscato},
  title        = {Palauner: policy-based active learning to augment named entity recognition
                  datasets},
  journal      = {Data Min. Knowl. Discov.},
  volume       = {40},
  number       = {5},
  pages        = {75},
  year         = {2026},
  url          = {https://doi.org/10.1007/s10618-026-01239-2},
  doi          = {10.1007/S10618-026-01239-2},
  timestamp    = {Thu, 23 Jul 2026 11:33:36 +0200},
  biburl       = {https://dblp.org/rec/journals/datamine/PostiglioneVSSM26.bib},
  bibsource    = {dblp computer science bibliography, https://dblp.org}
}