Repository navigation
Add RTP-CODEC training, tokenizers and evaluation benchmarks - #7
Merged
Merged
Conversation
- requirements_no_mirpy.txt: same deps minus mirpy-lib/tcrtrie (no C++ toolchain available, not needed for forward/inverse) - scripts/check_standardizer_equivalence.py: verifies direct np.mean/std matches batch_cache.py's online sum/sum_sq formula on train split before relying on it in the no-cache path
…ataset Rename encode import to default_char_encode; ForwardDataset and InverseDataset now accept optional encode_fn (defaults to char encode), enabling WordPiece tokenizer swap without touching model/collate code. No behavior change for existing char-baseline callers.
scripts/prepare_wordpiece_corpus.py: builds train-only CDR3 corpus (artifacts/tokenizers/train_cdr3.txt) using the same read_airr_table + split_indices(seed=42) pattern as batch_cache.py, preventing val/test leakage into tokenizer training. Also saves split_indices.json (positional indices, clone_id when available) and corpus_stats.json. Verified on trb_15k.tsv: train=12000/val=1500/test=1500, train corpus mean_length=13.95 (matches full-dataset mean=13.94).
…ity-only
- save per vocab_size under wordpiece_vocab_{N}/ (tokenizer.json + vocab.txt)
instead of flat wordpiece_vocab{N}.json, matching required project layout
- export vocab.txt via tokenizer.model.save() (BERT-style, one token per
line, ordered by id)
- tag each vocab_size in wordpiece_training_summary.json with role=
sanity_check (100) vs main_grid (1000/2000/5000/10000), so it's
excluded from the main ablation comparison
- vocab_size=20000 stays opt-in via --vocab-sizes, not in default grid
…ache batch_nocache.py: full in-memory load via dataio.load_airr_with_embeddings, no .npz shards/uuid cache dirs (streaming reader is broken in this environment). Same split_indices(seed) call as batch_cache.py -> identical train/val/test split. Standardizer switched to two-pass np.mean/std (verified equivalent to the old online sum/sum_sq formula: max mean diff=0.0, max std diff relative to std=4.6e-5, prior discrepancy was float cancellation, not a logic error). Public API (prepare_cached_training_data/cleanup_batch_cache/save_training_metadata) kept identical, so train_forward.py needs no changes to import this instead. Not wired into train_forward.py yet (--no-cache flag deferred).
scripts/calc_tokenization_stats.py: computes per-split (train/val/test) length/UNK/pad/compression stats, top-token reports, cross-vocab-size example table, and three alarm checks (oversized vocab, misconfigured tokenizer, insufficient compression benefit) for each trained WordPiece tokenizer. Outputs JSON + colored console summary + plain-text summary (tokenization_stats_summary.txt) + two diagnostic PNG plots (train memorization vs vocab_size, val/test compression plateau), matplotlib import guarded/optional via --skip-plots. Empirical finding: at min_frequency=1, vocab_size approaching train corpus size causes WordPiece to memorize whole training sequences as single vocab entries (train fraction_encoded_as_1_token spikes while val/test stays ~0). Confirmed on both the 12k-row and 80k-row corpora, with the safe vocab_size ceiling scaling roughly with corpus size.
…roken) batch_cache.py's shard-based/streaming loader is currently broken in this environment, so train_forward.py cannot run end-to-end against it. Temporarily import prepare_cached_training_data/cleanup_batch_cache/ save_training_metadata from batch_nocache.py instead (same public API, no other changes needed downstream). TODO: replace this direct import swap with a --no-cache/--cache-mode flag so both loaders stay selectable, and switch back to batch_cache.py once its streaming reader is fixed.
…fail loudly on missing tokenizer Adds --training-summary-path (single source of truth from train_wordpiece_tokenizers.py's output) and --allow-missing-tokenizers. Previously a missing tokenizer.json was silently skipped with a WARNING, causing undetected drift between --vocab-sizes and what was actually trained.
Converts human.trb.aa.txt (VDJtools pool format) to the junction_aa/v_call/ j_call/locus schema expected by prepare_wordpiece_corpus.py. Dedupes by cdr3aa, filters VALID_AA/max_len, subsamples to --target-size AFTER dedup (pool file is sorted by descending count, so sampling raw rows would bias toward frequent clonotypes).
Adds --top-k-summary (default 10) and a human-readable top_tokens_summary.txt alongside tokenization_stats_summary.txt, showing top-K by frequency/longest/ non-single-amino-acid per vocab_size. Full top-n (default 50) remains in tokenization_stats.json only - this file is a compact slice, not a replacement.
…ion_stats.json Parameterized (STATS_JSON_PATH cell, papermill-compatible) so it can be re-run against any experiment's stats without editing the notebook body. Recomputes all 6 diagnostic plots directly from the JSON rather than depending on the PNG files the script may or may not have generated.
Ran and cached log10 Pgen
PAD/UNK/BOS/EOS now use the same ids in both tokenizers, and decoding no longer depends on their numeric order.
Same PAD/UNK/BOS/EOS ids as the char tokenizer, so both can be used interchangeably without touching the dataset or model code.
DataLoader(shuffle=True) requires len(dataset) to build a RandomSampler. Without __len__ this raised TypeError before any training could start.
Adds --tokenizer-type/--tokenizer-path/--vocab-size and wires the chosen tokenizer's vocab size into the model instead of the old hardcoded default. Char stays the default, so existing runs don't change.
Both training scripts will call into this instead of each reimplementing the same validation, so they can't quietly diverge.
Trains each (tokenizer, seed) combo in-process, times training and inference, and writes per-run results plus mean +/- std across seeds into one table.
Trains each (tokenizer, seed) combo in-process, times training and inference, and writes per-run results plus mean +/- std across seeds into one table.
mean/std were computed from the train split but never used - the model was training against raw, unstandardized embedding targets. Now the same train-derived transform is applied to train/val/test before the data reaches the model.
raw_results.csv now gets one row appended (and flushed) right after each (tokenizer, seed) run finishes, instead of everything being written in one go at the very end. Re-running with the same --output-root skips whatever's already in the CSV, so a long overnight run can be safely interrupted and picked back up.
New filter_vocab_by_max_token_length() drops WordPiece vocab entries longer than N amino acids
Off by default. Filtered and unfiltered runs are kept as separate configs, never mixed into the same average. Only affects this comparison script - the tokenizer files and real training scripts aren't touched.
A rebuilt tokenizer forgot to mark special tokens as special, so skip_special_tokens quietly stopped working. Doesn't affect any results so far - nothing decodes yet, only encodes.
Four steps on top of the anchored + max_token_len=9 + individual max_len base: narrower only, shallower only, both together, then narrower still - to isolate width vs depth as separate capacity levers, per the coworker's suggestion to try shrinking the network.
The pretraining question was unanswerable from three seeds: the IRRM arms spread by 0.03-0.06 RMSE across seeds, larger than the scratch-vs-pretrained difference itself. pgen_seeds_array adds seven seeds for those two arms on the full split, and collect_pgen now reports Welch's t-test alongside the difference of means. benchmark_speed times the full sequence-to-prediction path for every arm on one device. The speeds recorded during training covered only the trained head, which omits the encoder for the frozen arms - the dominant cost - and head and encoder timings came from different machines, so they could not be combined. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…tion and V/J cleaning
VEK239
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds the RTP-CODEC Transformer pipeline with a shared configurable latent representation and three training objectives: CDR3 reconstruction, TCRemP embedding prediction, and Pgen regression.
Includes:
Validation: 84 CPU tests passed; 146 Bash/Slurm scripts passed syntax checks; nine training/benchmark entrypoints passed --help; compilation and diff checks passed. GitHub CI passed for the final revision. Full GPU training and external-data experiments were not rerun.
Historical character-decoder checkpoints using the earlier special-token order require their original revision or explicit migration; compatibility details are documented in docs/publication.md.