UniVHH

Code for UniVHH training and inference, with an IgBert checkpoint (seed 345).

Setup

Run commands from codes_publish/. Use an existing Python environment with the dependencies in requirements.txt. Scripts use .venv/bin/python or ../.venv/bin/python; set PY to use another environment.

Prepare:

• Local Hugging Face model and tokenizer caches (scripts run offline).
• Prepared datasets at clean/vhhhold_locked/data/<task>_<backbone>_pre/.

Each dataset is a saved Hugging Face DatasetDict with train, validation, and test splits. Inference only needs test. Required columns: input_ids, attention_mask, antigen_embeddings, task_ids, and labels.

If starting from cleaned CSV files, place them at clean/vhhhold_locked/<task>/{train,val,test}.csv and provide data/NBbench-antigen-embeddings.pt, then run:

bash run_preprocess.sh

CSV columns:

Antigen-related tasks also require Ag_label matching the embedding dictionary. Data cleaning, split assignment, and Paratope antigen ID assignment should already be complete.

Train

bash run_train.sh

Runs five backbones with seeds 234, 345, and 456 using 8 GPUs by default. Set NPROC_PER_NODE to change the GPU count:

NPROC_PER_NODE=1 bash run_train.sh

Models and predictions are saved under outputs_2/main/<SWANLAB_PROJECT>/<backbone>_<seed>_rerun/.

Inference

Run the included IgBert checkpoint:

bash run_infer.sh

To use a newly trained model, update --ckpt_dir and the matching --palm_type in run_infer.sh.

Results are saved as test_<task>_results.json and predictions_<task>.csv in the configured output directory.

Optional CDR metrics:

../.venv/bin/python scripts/compute_cdr_metrics.py outputs/infer_IgBert_345

Files

• src/: model, data loading, and metrics.
• scripts/: input preparation, training, inference, and CDR evaluation.
• checkpoints/: included IgBert weights.