Skills › Scientific Agent Skills › waypoint-bio › vocab_coverage.py
vocab_coverage.py
257 lines · skills/waypoint-bio/scripts/vocab_coverage.py at 5f2bc03
1#!/usr/bin/env python3
2"""Report how much of a dataset a Waypoint tokenizer can actually see.
4The Waypoint vocabulary is frozen at pretraining time from the Atlas corpus.
5Taxa outside it resolve to ``<unk>`` and are dropped: ``waypoint embed`` skips
6them before ordering, and ``filter_unk_taxa: true`` does the same during
7fine-tuning and benchmarking. Neither warns you. The upstream paper names this
8as the models' main limitation, so measure it before trusting a downstream
9number.
11Two coverage figures are reported per sample:
13* **taxon coverage** -- fraction of a sample's taxa that map to a real token.
14* **abundance coverage** -- fraction of a sample's *relative abundance* carried
15 by those taxa. This measures retained mass, not biological importance:
16 low-abundance taxa can still matter for a particular research question.
18Examples
19--------
20 python vocab_coverage.py --model outpost-bio/Waypoint-6m --data dataset.parquet
21 python vocab_coverage.py --model outputs/pretrain/best_model --data dataset.parquet \
22 --report-missing 20 --output coverage.csv
23"""
25from __future__ import annotations
27import argparse
28import ast
29import math
30import sys
31from collections import Counter
32from pathlib import Path
34import pandas as pd
37def _parse_list_cell(value):
38 """Parse one CSV cell back into a list.
40 ``ast.literal_eval`` only evaluates Python literals and cannot execute
41 code, but it still raises on malformed or pathologically nested input.
42 Turn that into a readable message instead of a traceback.
43 """
44 if not isinstance(value, str):
45 return value
46 try:
47 return ast.literal_eval(value)
48 except (ValueError, SyntaxError, MemoryError, RecursionError) as exc:
49 raise ValueError(
50 f"could not parse list cell {value[:60]!r}: {exc}. "
51 "Waypoint-format CSV stores Taxa and Relative Abundances as Python "
52 "list reprs -- prefer .parquet, which avoids this round-trip."
53 ) from exc
56def load_dataframe(path: Path) -> pd.DataFrame:
57 """Read a waypoint-format file, parsing list columns back from CSV/TSV."""
58 suffix = path.suffix.lower()
59 if suffix == ".parquet":
60 frame = pd.read_parquet(path)
61 elif suffix in {".csv", ".tsv", ".tab"}:
62 sep = "\t" if suffix in {".tsv", ".tab"} else ","
63 frame = pd.read_csv(path, sep=sep)
64 # Our converter writes the sample-ID index as its first CSV column.
65 if frame.columns[0] == "sample_id" or str(frame.columns[0]).startswith("Unnamed:"):
66 id_column = frame.columns[0]
67 frame = pd.read_csv(path, sep=sep, dtype={id_column: str}).set_index(id_column)
68 # Lists round-trip through CSV as their repr. Check each value rather
69 # than the column dtype: pandas 3 uses a dedicated string dtype, so a
70 # ``dtype == object`` guard silently skips the parse.
71 for column in ("Taxa", "Relative Abundances"):
72 if column in frame.columns:
73 frame[column] = frame[column].map(_parse_list_cell)
74 else:
75 raise ValueError(f"unsupported format {suffix!r}; use .parquet, .csv, or .tsv")
77 for column in ("Taxa", "Relative Abundances"):
78 if column not in frame.columns:
79 raise ValueError(f"{path} is not waypoint format: missing {column!r} column")
80 return frame
83def load_tokenizer(model: str, revision: str | None = None):
84 """Load a Waypoint tokenizer from a Hub id or local checkpoint directory.
86 Imported lazily so ``--help`` works without transformers installed.
87 """
88 try:
89 from waypoint_bio.tokenizer import load_tokenizer as _load
90 except ImportError:
91 pass
92 else:
93 if revision is not None:
94 from transformers import AutoTokenizer
95 return AutoTokenizer.from_pretrained(model, revision=revision, trust_remote_code=True)
96 return _load(model)
98 try:
99 from transformers import AutoTokenizer
100 except ImportError as exc: # pragma: no cover - environment dependent
101 raise SystemExit(
102 "Neither waypoint-bio nor transformers is installed. "
103 "Install with: pip install waypoint-bio"
104 ) from exc
105 return AutoTokenizer.from_pretrained(model, revision=revision, trust_remote_code=True)
108def coverage_report(
109 frame: pd.DataFrame, tokenizer
110) -> tuple[pd.DataFrame, Counter[str]]:
111 """Per-sample taxon and abundance coverage, plus a missing-taxon counter."""
112 unk_id = tokenizer.unk_token_id
113 if frame.empty:
114 raise ValueError("dataset has no samples")
115 if not frame.index.is_unique:
116 raise ValueError("sample IDs must be unique")
117 missing: Counter[str] = Counter()
118 rows = []
120 for sample_id, row in frame.iterrows():
121 taxa = row["Taxa"]
122 abundances = row["Relative Abundances"]
123 if any(isinstance(values, (str, bytes, dict)) or not hasattr(values, "__len__")
124 for values in (taxa, abundances)):
125 raise ValueError(f"sample {sample_id!r}: taxa and abundances must be lists")
126 if len(taxa) != len(abundances):
127 raise ValueError(f"sample {sample_id!r}: taxa and abundance lengths differ")
129 n_total = len(taxa)
130 n_known = 0
131 abundance_total = 0.0
132 abundance_known = 0.0
134 for taxon, abundance in zip(taxa, abundances):
135 value = float(abundance)
136 if not math.isfinite(value) or value < 0:
137 raise ValueError(f"sample {sample_id!r}: abundances must be finite and nonnegative")
138 abundance_total += value
139 if tokenizer.convert_tokens_to_ids(str(taxon)) == unk_id:
140 missing[str(taxon)] += 1
141 else:
142 n_known += 1
143 abundance_known += value
145 rows.append(
146 {
147 "sample_id": sample_id,
148 "n_taxa": n_total,
149 "n_in_vocab": n_known,
150 "taxon_coverage": n_known / n_total if n_total else 0.0,
151 "abundance_coverage": (
152 abundance_known / abundance_total if abundance_total else 0.0
153 ),
154 }
155 )
157 return pd.DataFrame(rows).set_index("sample_id"), missing
160def main(argv: list[str] | None = None) -> int:
161 parser = argparse.ArgumentParser(
162 description="Report Waypoint tokenizer vocabulary coverage for a dataset.",
163 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
164 )
165 parser.add_argument(
166 "--revision", help="Reviewed full Hub commit SHA for the remote tokenizer code."
167 )
168 parser.add_argument(
169 "--model",
170 default="outpost-bio/Waypoint-6m",
171 help="Hub id or local checkpoint directory (gated repos need HF_TOKEN).",
172 )
173 parser.add_argument(
174 "--data", required=True, help="Waypoint-format .parquet / .csv / .tsv."
175 )
176 parser.add_argument(
177 "--threshold",
178 type=float,
179 default=0.8,
180 help="Flag samples whose abundance coverage falls below this.",
181 )
182 parser.add_argument(
183 "--report-missing",
184 type=int,
185 default=15,
186 help="Show this many of the most frequent out-of-vocabulary taxa.",
187 )
188 parser.add_argument(
189 "--output", default=None, help="Optional path to write the per-sample table."
190 )
191 args = parser.parse_args(argv)
192 if not 0 <= args.threshold <= 1:
193 parser.error("threshold must be between 0 and 1")
195 data_path = Path(args.data)
196 if not data_path.exists():
197 parser.error(f"data not found: {data_path}")
199 frame = load_dataframe(data_path)
200 tokenizer = load_tokenizer(args.model, revision=args.revision)
201 report, missing = coverage_report(frame, tokenizer)
203 taxon = report["taxon_coverage"]
204 abundance = report["abundance_coverage"]
205 print(f"Model: {args.model} (vocab size {len(tokenizer.get_vocab())})")
206 print(f"Dataset: {data_path} ({len(report)} samples)")
207 print()
208 print(f"{'':22}{'median':>10}{'mean':>10}{'min':>10}")
209 print(
210 f"{'taxon coverage':22}{taxon.median():>10.3f}"
211 f"{taxon.mean():>10.3f}{taxon.min():>10.3f}"
212 )
213 print(
214 f"{'abundance coverage':22}{abundance.median():>10.3f}"
215 f"{abundance.mean():>10.3f}{abundance.min():>10.3f}"
216 )
218 below = report[abundance < args.threshold]
219 empty = report[report["n_in_vocab"] == 0]
220 print()
221 print(
222 f"{len(below)}/{len(report)} samples below the "
223 f"{args.threshold:.2f} abundance-coverage threshold"
224 )
225 if len(empty):
226 print(
227 f"{len(empty)} samples have NO in-vocabulary taxa -- these encode as "
228 "[BOS][EOS] and their embeddings are meaningless",
229 file=sys.stderr,
230 )
232 if missing and args.report_missing > 0:
233 print()
234 print(f"Most frequent out-of-vocabulary taxa ({len(missing)} distinct):")
235 for taxon_name, count in missing.most_common(args.report_missing):
236 print(f" {count:>6} {taxon_name}")
237 print()
238 print(
239 "Common causes: a different taxonomy database (SILVA/GTDB vs NCBI naming), "
240 "a 'd__' domain prefix the tokenizer ignores, '|' separators instead of ';', "
241 "or genuinely novel taxa."
242 )
244 if args.output:
245 out = Path(args.output)
246 out.parent.mkdir(parents=True, exist_ok=True)
247 if out.suffix.lower() == ".parquet":
248 report.to_parquet(out)
249 else:
250 report.to_csv(out)
251 print(f"\nWrote per-sample coverage to {out}")
253 return 0
256if __name__ == "__main__":
257 raise SystemExit(main())