diff --git a/docs/tutorials/bedbase-training-dataset.md b/docs/tutorials/bedbase-training-dataset.md new file mode 100644 index 00000000..de8f7a10 --- /dev/null +++ b/docs/tutorials/bedbase-training-dataset.md @@ -0,0 +1,112 @@ +# Build a tokenized training dataset from bedbase + +`geniml.dataset` turns a set of region sets into a tokenized training dataset for +region models (atacformer and any other tokenized-region model). It separates two +concerns: + +- a **source** — *where* the region sets come from (bedbase, a local file list, a + `BedSet`, ...), and +- a source-agnostic **dataset** — tokenize + window + batch. + +Adding a new data origin is a small source, not a new dataset. Bedbase is the first +concrete source: it turns a *live* bedbase selection into a tokenized dataset. + +## The pieces + +| Object | Role | +| --- | --- | +| `RegionSetSource` | Protocol: iterate `(RegionSet, meta)` tuples. | +| `BedbaseSource` | Yields region sets from a bedbase selection via `BBClient`. | +| `FileListSource` | Yields region sets from a list/text-file of local BED paths. | +| `BedSetSource` | Wraps a `geniml.io.BedSet` (or `BBClient.load_bedset` output). | +| `select_bedbase_samples` | Query the bedbase API → a dated `SampleSelection` manifest. | +| `TokenizedRegionDataset` | Tokenize+window any source into a training dataset. | + +## Bedbase → tokenized dataset + +```python +from geniml.atacformer import TrainingTokenizer +from geniml.dataset import ( + select_bedbase_samples, + BedbaseSource, + TokenizedRegionDataset, +) + +# 1. Select samples from the live bedbase API (replaces a hand-curated CSV). +# The selection is a dated, serializable manifest — provenance for "all of bedbase". +selection = select_bedbase_samples(genome="hg38", qc="good", limit=500) + +# 2. Wrap the selection as a source. Region sets are read via BBClient — from the +# local bbcache mirror if present, else downloaded and cached. +source = BedbaseSource(selection) + +# 3. Tokenize + window into a training dataset against a universe. +tokenizer = TrainingTokenizer("path/to/universe.bed") +ds = TokenizedRegionDataset( + source, + tokenizer, + context_size=8192, + mode="materialize", # tokenize once into an on-disk parquet (cached) + cache_dir="/scratch/bb_cache", + universe="path/to/universe.bed", # lets num_proc>1 rebuild the tokenizer in workers + num_proc=16, +).build() # -> a datasets.Dataset ready for a HF Trainer +``` + +Each row is `{"input_ids": [...], **metadata}`, where the metadata columns +(`description`, `species_name`, `cell_type`, `cell_line`, `tissue`, `assay`, +`antibody`, `target`, `treatment`, ...) are carried through for conditioning and +analysis. Drop them before training if your collator doesn't want them (or rely on +the HF `Trainer`'s `remove_unused_columns`). + +### Materialize vs stream + +- **`mode="materialize"`** (default) tokenizes the whole selection once into an + on-disk `datasets.Dataset`. The cache is keyed on *(selection, universe, context + size)*, so re-running the same selection reloads the parquet instead of + re-tokenizing. Best for multi-epoch training. +- **`mode="stream"`** returns a `torch` `IterableDataset` that tokenizes on the fly, + with an optional `.gtok` cache so epoch 2+ skips re-tokenization. Best for "all of + bedbase, always current, no giant parquet". + +```python +stream_ds = TokenizedRegionDataset( + source, tokenizer, context_size=8192, mode="stream", cache_dir="/scratch/gtok" +).build() +``` + +## Other sources, same dataset + +A local file list (what region2vec's `BEDDataset` did, on the shared interface): + +```python +from geniml.dataset import FileListSource, TokenizedRegionDataset + +source = FileListSource("bedfiles.txt") # or FileListSource(["a.bed", "b.bed"]) +ds = TokenizedRegionDataset(source, tokenizer).build() +``` + +An existing `BedSet` (including `BBClient.load_bedset` output): + +```python +from geniml.bbclient import BBClient +from geniml.dataset import BedSetSource, TokenizedRegionDataset + +bedset = BBClient().load_bedset("your-bedset-id") +ds = TokenizedRegionDataset(BedSetSource(bedset), tokenizer).build() +``` + +## Reproducibility + +Serialize the selection so a run is dated and rebuildable: + +```python +selection.to_json("bedbase_selection.json") +# later: +from geniml.dataset import SampleSelection +selection = SampleSelection.from_json("bedbase_selection.json") +``` + +`geniml.dataset.write_provenance(...)` writes a `dataset_provenance.json` (selection +hash, universe, context size, snapshot date) alongside a trained model, so an "all of +bedbase" run records exactly which bedbase snapshot produced it. diff --git a/geniml/dataset/__init__.py b/geniml/dataset/__init__.py new file mode 100644 index 00000000..6f14b02e --- /dev/null +++ b/geniml/dataset/__init__.py @@ -0,0 +1,67 @@ +"""Source-agnostic tokenized datasets for region-based models. + +Separates **source** (where region sets come from) from **dataset** (tokenize + +window + batch), built on the shared ``RegionSet``/``BedSet`` abstraction. A new +data origin is a small :class:`RegionSetSource`, not a new dataset -- so bedbase, +local file lists, and bedsets all feed the same +:class:`TokenizedRegionDataset`. + +Typical use (bedbase -> tokenized training dataset):: + + from geniml.atacformer import TrainingTokenizer + from geniml.dataset import ( + select_bedbase_samples, BedbaseSource, TokenizedRegionDataset, + ) + + selection = select_bedbase_samples(genome="hg38", qc="good", limit=500) + source = BedbaseSource(selection) + tokenizer = TrainingTokenizer("path/to/universe.bed") + ds = TokenizedRegionDataset(source, tokenizer, context_size=8192, + cache_dir="/scratch/bb_cache").build() + +The heavy deps (``datasets``, ``torch``) are imported lazily inside the dataset's +methods, so importing this package works on a base geniml install; only building a +dataset requires ``geniml[ml]``. +""" + +from .source import ( + BedSetSource, + FileListSource, + RegionSetItem, + RegionSetSource, +) +from .tokenize import sample_and_remove, tokenize_regionset +from .bedbase import ( + METADATA_COLUMNS, + BedbaseSource, + SampleRef, + SampleSelection, + select_bedbase_samples, +) +from .dataset import StreamingTokenizedRegionDataset, TokenizedRegionDataset +from .manifest import cache_key, source_manifest, universe_signature, write_provenance + +__all__ = [ + # source interface + "RegionSetSource", + "RegionSetItem", + "BedSetSource", + "FileListSource", + # tokenization core + "tokenize_regionset", + "sample_and_remove", + # bedbase source + selection + "select_bedbase_samples", + "BedbaseSource", + "SampleRef", + "SampleSelection", + "METADATA_COLUMNS", + # dataset + "TokenizedRegionDataset", + "StreamingTokenizedRegionDataset", + # manifest / provenance + "cache_key", + "source_manifest", + "universe_signature", + "write_provenance", +] diff --git a/geniml/dataset/bedbase.py b/geniml/dataset/bedbase.py new file mode 100644 index 00000000..d07db758 --- /dev/null +++ b/geniml/dataset/bedbase.py @@ -0,0 +1,312 @@ +"""Bedbase as a data source. + +The first concrete :class:`~geniml.dataset.RegionSetSource`: turn a *live* bedbase +selection into region sets. This replaces the lost, hand-curated +``all_hg38_good_meta.csv`` -- the sample list is now resolved from the bedbase API +at selection time and serialized as a dated, reproducible manifest. + +Two pieces: + +- :func:`select_bedbase_samples` -- query the bedbase API for beds of a given genome + (with a light QC filter) and return :class:`SampleRef`s carrying the metadata + columns the training pipeline conditions on. +- :class:`BedbaseSource` -- wrap a selection and yield ``(RegionSet, meta)`` by + reading each bed through :class:`geniml.bbclient.BBClient` (from the local bbcache + mirror when present, else download+cache). It never queries the bbcache SQLite + directly. + +On QC: the bedbase ``/v1/bed/list`` endpoint exposes ``genome`` and +``bed_compliance`` filters but no explicit "QC-good" flag, so ``qc="good"`` here +means: a real, processed bed for the requested genome (``is_universe`` false) that +meets ``min_compliance``. A stricter region-count filter is available via +``min_regions``/``max_regions`` but requires a per-bed metadata fetch, so it is +opt-in. +""" + +import os +from dataclasses import dataclass, asdict, field +from logging import getLogger +from typing import Dict, Iterator, List, Optional + +import requests + +from .source import RegionSetItem + +_LOGGER = getLogger("geniml.dataset") + +# Defined locally (not imported from geniml.bbclient.const) so selecting samples +# does not pull in the bbclient package's heavy s3/zarr import chain. Kept in sync +# with geniml.bbclient.const.DEFAULT_BEDBASE_API. +DEFAULT_BEDBASE_API = os.getenv("BEDBASE_API") or "https://api.bedbase.org" + +#: bedbase metadata columns preserved on each sample (parity with the frozen parquet). +METADATA_COLUMNS = ( + "description", + "species_name", + "cell_type", + "cell_line", + "tissue", + "assay", + "antibody", + "target", + "treatment", +) + +#: bed_compliance strings, weakest to strongest, for the ``min_compliance`` filter. +_COMPLIANCE_ORDER = ["bed2+0", "bed3+0", "bed4+0", "bed5+0", "bed6+0"] + + +@dataclass +class SampleRef: + """A single selected bedbase sample: its id plus the metadata columns.""" + + id: str + name: Optional[str] = None + genome: Optional[str] = None + number_of_regions: Optional[int] = None + meta: Dict = field(default_factory=dict) + + def as_meta(self) -> Dict: + """Flatten to a metadata dict suitable for a dataset row.""" + out = {"id": self.id, "name": self.name, "genome": self.genome} + out.update(self.meta) + return out + + +@dataclass +class SampleSelection: + """A dated, serializable manifest of a bedbase selection (provenance). + + Serialize with :meth:`to_json` so an "all of bedbase" run is reproducible: the + exact ids, the genome, the QC filter, the API, and the snapshot date. + """ + + genome: str + qc: str + bedbase_api: str + snapshot_date: Optional[str] + samples: List[SampleRef] + + def ids(self) -> List[str]: + return [s.id for s in self.samples] + + def __len__(self) -> int: + return len(self.samples) + + def to_json(self, path: str) -> str: + import json + + payload = { + "genome": self.genome, + "qc": self.qc, + "bedbase_api": self.bedbase_api, + "snapshot_date": self.snapshot_date, + "samples": [asdict(s) for s in self.samples], + } + with open(path, "w") as fh: + json.dump(payload, fh, indent=2) + return path + + @classmethod + def from_json(cls, path: str) -> "SampleSelection": + import json + + with open(path, "r") as fh: + payload = json.load(fh) + samples = [SampleRef(**s) for s in payload["samples"]] + return cls( + genome=payload["genome"], + qc=payload["qc"], + bedbase_api=payload["bedbase_api"], + snapshot_date=payload.get("snapshot_date"), + samples=samples, + ) + + +def _compliance_ok(value: Optional[str], minimum: str) -> bool: + if not minimum: + return True + try: + return _COMPLIANCE_ORDER.index(value) >= _COMPLIANCE_ORDER.index(minimum) + except ValueError: + # unknown compliance string -> keep it (don't silently drop unfamiliar beds) + return True + + +def _record_to_sampleref(record: Dict) -> SampleRef: + annotation = record.get("annotation") or {} + meta = {col: annotation.get(col) for col in METADATA_COLUMNS} + # bedbase calls the organism 'organism'; the training pipeline expects species_name + if meta.get("species_name") is None: + meta["species_name"] = annotation.get("organism") + return SampleRef( + id=record["id"], + name=record.get("name"), + genome=record.get("genome_alias"), + meta=meta, + ) + + +def select_bedbase_samples( + genome: str = "hg38", + qc: str = "good", + limit: Optional[int] = None, + bedbase_api: str = DEFAULT_BEDBASE_API, + min_compliance: str = "bed3+0", + min_regions: Optional[int] = None, + max_regions: Optional[int] = None, + page_size: int = 1000, + snapshot_date: Optional[str] = None, +) -> SampleSelection: + """Select bedbase samples for a genome, returning a dated selection manifest. + + Args: + genome: genome alias to filter on (e.g. ``"hg38"``). + qc: QC policy label recorded in the manifest. ``"good"`` applies the + ``is_universe``/``min_compliance`` filter; ``"all"`` keeps everything + for the genome. + limit: cap the number of selected samples (``None`` = all). + bedbase_api: bedbase API base URL. + min_compliance: minimum ``bed_compliance`` to keep (see ``_COMPLIANCE_ORDER``). + min_regions: if set, drop beds with fewer regions (requires a per-bed + metadata fetch -- slower). + max_regions: if set, drop beds with more regions (requires a per-bed fetch). + page_size: API page size for pagination. + snapshot_date: ISO date string stamped into the manifest for provenance. If + ``None``, caller should stamp it (kept out of here so the function is + deterministic/testable). + + Returns: + SampleSelection: the selection manifest. + """ + strict = qc != "all" + need_stats = min_regions is not None or max_regions is not None + + samples: List[SampleRef] = [] + offset = 0 + session = requests.Session() + while True: + url = f"{bedbase_api}/v1/bed/list" + params = {"genome": genome, "limit": page_size, "offset": offset} + resp = session.get(url, params=params, timeout=60) + resp.raise_for_status() + payload = resp.json() + results = payload.get("results", []) + if not results: + break + + for record in results: + if strict: + if record.get("is_universe"): + continue + if not _compliance_ok(record.get("bed_compliance"), min_compliance): + continue + ref = _record_to_sampleref(record) + if need_stats: + n = _fetch_region_count(session, bedbase_api, ref.id) + ref.number_of_regions = n + if min_regions is not None and (n is None or n < min_regions): + continue + if max_regions is not None and (n is None or n > max_regions): + continue + samples.append(ref) + if limit is not None and len(samples) >= limit: + break + + if limit is not None and len(samples) >= limit: + break + offset += page_size + if offset >= payload.get("count", 0): + break + + _LOGGER.info("Selected %d bedbase samples for genome=%s (qc=%s)", len(samples), genome, qc) + return SampleSelection( + genome=genome, + qc=qc, + bedbase_api=bedbase_api, + snapshot_date=snapshot_date, + samples=samples, + ) + + +def _fetch_region_count(session, bedbase_api: str, bed_id: str) -> Optional[int]: + try: + resp = session.get( + f"{bedbase_api}/v1/bed/{bed_id}/metadata", params={"full": "true"}, timeout=60 + ) + resp.raise_for_status() + stats = resp.json().get("stats") or {} + n = stats.get("number_of_regions") + return int(n) if n is not None else None + except Exception: # pragma: no cover - network best-effort + return None + + +class BedbaseSource: + """A :class:`~geniml.dataset.RegionSetSource` over a bedbase selection. + + Yields ``(RegionSet, meta)`` by loading each selected bed through a + :class:`~geniml.bbclient.BBClient` -- from the local bbcache mirror if present, + otherwise downloading and caching it. Reads only via ``BBClient.load_bed``; it + never touches the bbcache SQLite directly. + """ + + def __init__( + self, + selection: SampleSelection, + bbclient=None, + skip_errors: bool = True, + ): + """Initialize a BedbaseSource. + + Args: + selection: a :class:`SampleSelection` (from :func:`select_bedbase_samples` + or loaded from a manifest). + bbclient: a configured ``geniml.bbclient.BBClient``; a default one is + created if omitted (imported lazily, since it pulls in s3/zarr deps). + skip_errors: if True, samples that fail to load are logged and skipped + rather than aborting the whole iteration. + """ + if bbclient is None: + from ..bbclient.bbclient import BBClient + + bbclient = BBClient() + self.selection = selection + self.bbclient = bbclient + self.skip_errors = skip_errors + + def __len__(self) -> int: + return len(self.selection) + + def __iter__(self) -> Iterator[RegionSetItem]: + for ref in self.selection.samples: + try: + region_set = self.bbclient.load_bed(ref.id) + except Exception as exc: # network / cache miss + if self.skip_errors: + _LOGGER.warning("Skipping bed %s: %s", ref.id, exc) + continue + raise + yield region_set, ref.as_meta() + + def manifest(self) -> List[str]: + """Stable list of selected bed ids for cache keying/provenance.""" + return self.selection.ids() + + def path_items(self): + """Return ``[(bed_path, meta), ...]``, loading (downloading) beds as needed. + + Used by the dataset's parallel materialize path, which tokenizes from BED + paths in worker processes. + """ + out = [] + for ref in self.selection.samples: + try: + rs = self.bbclient.load_bed(ref.id) + out.append((rs.path, ref.as_meta())) + except Exception as exc: + if self.skip_errors: + _LOGGER.warning("Skipping bed %s: %s", ref.id, exc) + continue + raise + return out diff --git a/geniml/dataset/dataset.py b/geniml/dataset/dataset.py new file mode 100644 index 00000000..aa6da071 --- /dev/null +++ b/geniml/dataset/dataset.py @@ -0,0 +1,337 @@ +"""Source-agnostic tokenized training dataset. + +:class:`TokenizedRegionDataset` is the join between the *read* half (a +:class:`~geniml.dataset.RegionSetSource`) and the *train* half (a tokenizer + a +HuggingFace ``Trainer``). It tokenizes and windows every region set from a source +into training rows (``input_ids`` + preserved metadata), in one of two modes: + +- ``materialize`` -- tokenize once into an on-disk parquet/Arrow dataset and return + a ``datasets.Dataset``. Cache is keyed on (source selection, universe, context + size), so re-running the same selection reloads instead of re-tokenizing. Best for + multi-epoch training. +- ``stream`` -- a ``torch.utils.data.IterableDataset`` that tokenizes on the fly, + with an optional ``.gtok`` cache so epoch 2+ skips re-tokenization. Best for "all + of bedbase, always current, no giant parquet". + +Rows match what the atacformer RTD collator expects (``input_ids``; masks come from +the collator). ``datasets`` / ``torch`` are imported lazily so importing the source +and tokenization halves works on a base install. +""" + +import os +import random +from logging import getLogger +from typing import Dict, List, Optional + +from .manifest import cache_key +from .source import RegionSetSource +from .tokenize import sample_and_remove, tokenize_regionset + +_LOGGER = getLogger("geniml.dataset") + + +class TokenizedRegionDataset: + """Tokenize + window any :class:`RegionSetSource` into a training dataset. + + Args: + source: where region sets come from (a :class:`RegionSetSource`). + tokenizer: a ``gtars`` tokenizer / ``TrainingTokenizer`` bound to a universe. + context_size: tokens per training window. + mode: ``"materialize"`` (default) or ``"stream"``. + cache_dir: directory for the materialized parquet cache and/or gtok cache. + max_windows_per_file: skip a file producing more than this many windows. + num_proc: worker processes for materialize-mode tokenization. Parallelism + requires ``universe`` (workers rebuild the tokenizer) and path-backed + items; otherwise it falls back to sequential. + seed: RNG seed for reproducible windowing. + drop_unk: drop ``unk`` tokens before windowing. + universe: universe path/id used to rebuild the tokenizer in worker processes + (enables ``num_proc > 1``). + """ + + def __init__( + self, + source: RegionSetSource, + tokenizer, + context_size: int = 8192, + mode: str = "materialize", + cache_dir: Optional[str] = None, + max_windows_per_file: int = 10, + num_proc: int = 1, + seed: int = 42, + drop_unk: bool = True, + universe: Optional[str] = None, + ): + if mode not in ("materialize", "stream"): + raise ValueError(f"mode must be 'materialize' or 'stream', got {mode!r}") + self.source = source + self.tokenizer = tokenizer + self.context_size = context_size + self.mode = mode + self.cache_dir = cache_dir + self.max_windows_per_file = max_windows_per_file + self.num_proc = max(1, num_proc) + self.seed = seed + self.drop_unk = drop_unk + self.universe = universe + + # -- public API ------------------------------------------------------- + + def build(self): + """Build the dataset according to ``mode``. + + Returns: + A ``datasets.Dataset`` (materialize) or a + :class:`StreamingTokenizedRegionDataset` (stream). + """ + if self.mode == "materialize": + return self.materialize() + return self.stream() + + def materialize(self): + """Tokenize the whole source into an on-disk ``datasets.Dataset`` (cached).""" + try: + from datasets import Dataset + except ImportError as exc: # pragma: no cover + raise ImportError( + "materialize mode requires the `datasets` package (pip install geniml[ml])." + ) from exc + + out_path = self._cache_parquet_path() + if out_path and os.path.exists(out_path): + _LOGGER.info("Loading cached tokenized dataset from %s", out_path) + return Dataset.from_parquet(out_path) + + rows = self._tokenize_all() + _LOGGER.info("Tokenized %d windows from source", len(rows)) + ds = Dataset.from_list(rows) + + if out_path: + os.makedirs(os.path.dirname(out_path), exist_ok=True) + ds.to_parquet(out_path) + _LOGGER.info("Wrote tokenized cache to %s", out_path) + return ds + + def stream(self) -> "StreamingTokenizedRegionDataset": + """Return a streaming ``IterableDataset`` over the source.""" + return StreamingTokenizedRegionDataset( + source=self.source, + tokenizer=self.tokenizer, + context_size=self.context_size, + max_windows_per_file=self.max_windows_per_file, + cache_dir=self.cache_dir, + seed=self.seed, + drop_unk=self.drop_unk, + ) + + # -- internals -------------------------------------------------------- + + def _cache_parquet_path(self) -> Optional[str]: + if not self.cache_dir: + return None + key = cache_key(self.source, self.tokenizer, self.context_size) + if key is None: + return None + return os.path.join(self.cache_dir, f"tokenized_{key}.parquet") + + def _tokenize_all(self) -> List[Dict]: + if self.num_proc > 1: + parallel_rows = self._tokenize_parallel() + if parallel_rows is not None: + return parallel_rows + _LOGGER.warning( + "num_proc>1 requested but source is not path-backed or `universe` " + "was not given; falling back to sequential tokenization." + ) + return self._tokenize_sequential() + + def _tokenize_sequential(self) -> List[Dict]: + rows: List[Dict] = [] + rng = random.Random(self.seed) + for region_set, meta in self.source: + windows = tokenize_regionset( + region_set, + self.tokenizer, + context_size=self.context_size, + max_windows=self.max_windows_per_file, + rng=rng, + drop_unk=self.drop_unk, + ) + for window in windows: + rows.append({"input_ids": window, **meta}) + return rows + + def _tokenize_parallel(self) -> Optional[List[Dict]]: + """Parallel tokenization from BED paths; None if not applicable.""" + if not self.universe: + return None + items = self._path_items() + if items is None: + return None + + from functools import partial + from multiprocessing import Pool + + worker = partial( + _tokenize_path_worker, + universe=self.universe, + context_size=self.context_size, + max_windows=self.max_windows_per_file, + drop_unk=self.drop_unk, + ) + # deterministic per-item seed derived from the base seed + index + tasks = [(i, path, meta, self.seed + i) for i, (path, meta) in enumerate(items)] + + rows: List[Dict] = [] + with Pool(processes=self.num_proc) as pool: + for windows, meta in pool.imap_unordered(worker, tasks, chunksize=8): + for window in windows: + rows.append({"input_ids": window, **meta}) + return rows + + def _path_items(self): + """Return ``[(path, meta), ...]`` if every item is path-backed, else None.""" + path_items_fn = getattr(self.source, "path_items", None) + if callable(path_items_fn): + return path_items_fn() + items = [] + for region_set, meta in self.source: + path = getattr(region_set, "path", None) or meta.get("path") + if not path: + return None + items.append((path, meta)) + return items + + +def _tokenize_path_worker(task, universe, context_size, max_windows, drop_unk): + """Top-level worker: (idx, bed_path, meta, seed) -> (windows, meta). + + Rebuilds the tokenizer from ``universe`` in each process (gtars tokenizers are + not picklable across processes). + """ + idx, path, meta, seed = task + from gtars.models import RegionSet + from gtars.tokenizers import Tokenizer + + tokenizer = Tokenizer.from_bed(universe) if os.path.isfile(universe) else Tokenizer(universe) + windows = tokenize_regionset( + RegionSet(path), + tokenizer, + context_size=context_size, + max_windows=max_windows, + rng=random.Random(seed), + drop_unk=drop_unk, + ) + return windows, meta + + +_STREAM_CLS = None + + +def _streaming_class(): + """Build (once, lazily) the concrete IterableDataset subclass. + + ``torch`` is optional, so the class can't be declared at module top; the + DataLoader relies on ``isinstance(ds, IterableDataset)``, so streaming must be a + genuine subclass rather than a duck-typed stand-in. + """ + global _STREAM_CLS + if _STREAM_CLS is not None: + return _STREAM_CLS + + try: + from torch.utils.data import IterableDataset + except ImportError as exc: # pragma: no cover + raise ImportError("stream mode requires torch (pip install geniml[ml]).") from exc + + class _StreamingTokenizedRegionDataset(IterableDataset): + """Tokenizes a source on the fly, yielding ``{"input_ids": [...], **meta}``. + + With ``cache_dir`` set, each item's filtered token ids are cached to a + ``.gtok`` on first pass so later epochs skip re-tokenization (windowing is + re-done each epoch -- cheap, and keeps windows fresh). Composes with the HF + ``Trainer`` and the RTD ``DataCollator``. + """ + + def __init__( + self, + source, + tokenizer, + context_size=8192, + max_windows_per_file=10, + cache_dir=None, + seed=42, + drop_unk=True, + ): + super().__init__() + self.source = source + self.tokenizer = tokenizer + self.context_size = context_size + self.max_windows_per_file = max_windows_per_file + self.cache_dir = cache_dir + self.seed = seed + self.drop_unk = drop_unk + self._epoch = 0 + if cache_dir: + os.makedirs(cache_dir, exist_ok=True) + + def __iter__(self): + import torch + + worker_info = torch.utils.data.get_worker_info() + num_workers = worker_info.num_workers if worker_info else 1 + worker_id = worker_info.id if worker_info else 0 + + rng = random.Random(self.seed + self._epoch) + self._epoch += 1 + + for i, (region_set, meta) in enumerate(self.source): + if i % num_workers != worker_id: + continue # shard across DataLoader workers + ids = self._cached_ids(region_set, meta) + if not ids: + continue + if ( + self.max_windows_per_file is not None + and len(ids) > self.max_windows_per_file * self.context_size + ): + continue + for window in sample_and_remove(ids, self.context_size, rng=rng): + yield {"input_ids": window, **meta} + + def _cached_ids(self, region_set, meta): + item_id = meta.get("id") or getattr(region_set, "identifier", None) + gtok_path = ( + os.path.join(self.cache_dir, f"{item_id}.ctx.gtok") + if (self.cache_dir and item_id) + else None + ) + if gtok_path and os.path.exists(gtok_path): + from gtars.utils import read_tokens_from_gtok + + return list(read_tokens_from_gtok(gtok_path)) + + ids = self.tokenizer(region_set)["input_ids"] + if self.drop_unk: + unk_id = getattr(self.tokenizer, "unk_token_id", None) + if unk_id is not None: + ids = [i for i in ids if i != unk_id] + + if gtok_path and ids: + from gtars.utils import write_tokens_to_gtok + + write_tokens_to_gtok(gtok_path, ids) + return ids + + _STREAM_CLS = _StreamingTokenizedRegionDataset + return _STREAM_CLS + + +def StreamingTokenizedRegionDataset(*args, **kwargs): + """Construct a streaming tokenized dataset (a ``torch`` ``IterableDataset``). + + A factory rather than a class so the torch dependency stays lazy; the returned + instance is a genuine ``IterableDataset`` subclass, so the HF ``Trainer`` / + ``DataLoader`` treat it as iterable-style. + """ + return _streaming_class()(*args, **kwargs) diff --git a/geniml/dataset/manifest.py b/geniml/dataset/manifest.py new file mode 100644 index 00000000..fa69c5e0 --- /dev/null +++ b/geniml/dataset/manifest.py @@ -0,0 +1,99 @@ +"""Manifest, cache keying, and provenance. + +Two jobs: + +- **Cache keying**: a materialized dataset is keyed on (source selection, universe, + context size) so re-running the same selection is incremental -- it reloads the + parquet instead of re-tokenizing. +- **Provenance**: emit a small JSON alongside a trained model recording *exactly* + which bedbase selection, universe, and snapshot date produced it, so an "all of + bedbase" run is dated and rebuildable. +""" + +import hashlib +import json +import os +from typing import List, Optional + + +def hash_ids(ids: List[str]) -> str: + """Order-independent hash of a list of item ids (the source selection).""" + h = hashlib.sha256() + for item in sorted(str(i) for i in ids): + h.update(item.encode("utf-8")) + h.update(b"\n") + return h.hexdigest()[:16] + + +def universe_signature(tokenizer) -> str: + """Stable hash of a tokenizer's vocabulary (identifies the universe).""" + try: + vocab = tokenizer.get_vocab() + blob = json.dumps(vocab, sort_keys=True) + except Exception: + blob = repr(getattr(tokenizer, "vocab_size", tokenizer)) + return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16] + + +def source_manifest(source) -> Optional[List[str]]: + """Return a source's stable item-id list, or ``None`` if it can't provide one. + + A source opts in to caching by defining ``manifest()``; sources that can't + (unknown/streaming) return ``None`` and the dataset simply tokenizes fresh. + """ + fn = getattr(source, "manifest", None) + if callable(fn): + return fn() + return None + + +def cache_key(source, tokenizer, context_size: int) -> Optional[str]: + """Compute the cache key for a (source, universe, context_size) combination. + + Returns ``None`` when the source has no manifest (caching disabled). + """ + ids = source_manifest(source) + if ids is None: + return None + return f"{hash_ids(ids)}.{universe_signature(tokenizer)}.ctx{context_size}" + + +def write_provenance( + directory: str, + *, + source_ids: Optional[List[str]], + universe: str, + context_size: int, + snapshot_date: Optional[str] = None, + extra: Optional[dict] = None, +) -> str: + """Write a ``dataset_provenance.json`` describing a materialized run. + + Args: + directory: directory to write into (created if missing) -- typically the + trained model's output dir. + source_ids: the exact list of selected item ids (the manifest). + universe: universe identifier/path used for tokenization. + context_size: window size used. + snapshot_date: date the source selection was made (ISO string). + extra: any additional fields to record (bedbase api, qc filter, etc.). + + Returns: + str: path to the written JSON file. + """ + os.makedirs(directory, exist_ok=True) + payload = { + "universe": universe, + "universe_hash": None, + "context_size": context_size, + "snapshot_date": snapshot_date, + "n_items": len(source_ids) if source_ids is not None else None, + "source_ids_hash": hash_ids(source_ids) if source_ids is not None else None, + "source_ids": source_ids, + } + if extra: + payload.update(extra) + path = os.path.join(directory, "dataset_provenance.json") + with open(path, "w") as fh: + json.dump(payload, fh, indent=2) + return path diff --git a/geniml/dataset/source.py b/geniml/dataset/source.py new file mode 100644 index 00000000..676ea935 --- /dev/null +++ b/geniml/dataset/source.py @@ -0,0 +1,135 @@ +"""Region-set sources. + +A *source* is anything that yields :class:`gtars.models.RegionSet` objects plus +per-item metadata. It is the pluggable half of the dataset: where the region sets +come from. A source does **no** tokenizing -- it only produces region sets. Adding +a new data origin (local files, bedbase, anndata, ...) means writing a small source, +not a new dataset. + +The interface is deliberately minimal (see :class:`RegionSetSource`): + +- ``__iter__`` yields ``(RegionSet, meta)`` tuples (:data:`RegionSetItem`). +- ``__len__`` returns the number of items, or ``None`` if unknown/streaming. +- ``manifest`` (optional) returns a stable list of item ids so the dataset can + key an on-disk cache on the exact selection. +""" + +import os +from logging import getLogger +from typing import Dict, Iterator, List, Optional, Tuple, Union +from typing_extensions import Protocol, runtime_checkable + +from gtars.models import RegionSet + +_LOGGER = getLogger("geniml.dataset") + +#: A single item produced by a source: a region set plus its metadata dict. +RegionSetItem = Tuple[RegionSet, Dict] + + +@runtime_checkable +class RegionSetSource(Protocol): + """Protocol for anything that yields region sets plus metadata. + + A source is the pluggable input to :class:`~geniml.dataset.TokenizedRegionDataset`. + Any object that is iterable over ``(RegionSet, meta)`` tuples satisfies it -- + ``geniml.io.BedSet`` already very nearly does. + """ + + def __iter__(self) -> Iterator[RegionSetItem]: # pragma: no cover - protocol + ... + + def __len__(self) -> Optional[int]: # pragma: no cover - protocol + ... + + +class BedSetSource: + """Wrap a :class:`geniml.io.BedSet` (or ``BBClient.load_bedset`` output) as a source. + + ``BedSet`` already holds a list of ``RegionSet``s and is iterable, so this + adapter is thin: it just attaches metadata (the region set's identifier and + path) to each item. Use it to feed an existing bedset -- including one returned + by :meth:`geniml.bbclient.BBClient.load_bedset` -- straight into the dataset. + """ + + def __init__(self, bedset, meta: Optional[Dict] = None): + """Initialize a BedSetSource. + + Args: + bedset: a ``geniml.io.BedSet`` (anything iterable over RegionSets with + a ``__len__``). + meta: optional metadata merged into every item's metadata dict (e.g. + the bedset identifier), useful for provenance/conditioning. + """ + self.bedset = bedset + self.meta = dict(meta or {}) + identifier = getattr(bedset, "identifier", None) + if identifier and "bedset_id" not in self.meta: + self.meta["bedset_id"] = identifier + + def __len__(self) -> Optional[int]: + try: + return len(self.bedset) + except TypeError: + return None + + def __iter__(self) -> Iterator[RegionSetItem]: + for region_set in self.bedset: + meta = dict(self.meta) + meta.setdefault("id", getattr(region_set, "identifier", None)) + meta.setdefault("path", getattr(region_set, "path", None)) + yield region_set, meta + + def manifest(self) -> Optional[List[str]]: + """Stable list of item ids (region-set identifiers) for cache keying.""" + try: + return [getattr(rs, "identifier", None) or getattr(rs, "path") for rs in self.bedset] + except Exception: # pragma: no cover - defensive + return None + + +class FileListSource: + """A source over a list of local BED file paths. + + This is exactly what region2vec's ``BEDDataset`` did (a text file listing local + BED paths), expressed on the shared source interface. Accepts either a path to a + text file (one BED path per line) or an in-memory list of paths. + """ + + def __init__(self, files: Union[str, os.PathLike, List[str]], root: Optional[str] = None): + """Initialize a FileListSource. + + Args: + files: path to a text file listing one BED path per line, OR a list of + BED file paths. + root: optional directory prepended to each (relative) path. + """ + if isinstance(files, (str, os.PathLike)) and os.path.isfile(files): + with open(files, "r") as fh: + paths = [line.strip() for line in fh if line.strip()] + elif isinstance(files, (list, tuple)): + paths = [str(p) for p in files] + else: + raise ValueError( + "`files` must be a path to a text file listing BED paths, or a list of paths." + ) + if root: + paths = [p if os.path.isabs(p) else os.path.join(root, p) for p in paths] + self.paths = paths + + def __len__(self) -> int: + return len(self.paths) + + def __iter__(self) -> Iterator[RegionSetItem]: + for path in self.paths: + rs = RegionSet(path) + meta = { + "id": getattr(rs, "identifier", None), + "path": path, + "name": os.path.basename(path), + } + yield rs, meta + + def manifest(self) -> List[str]: + """Stable list of item ids (the sorted file paths).""" + return sorted(self.paths) diff --git a/geniml/dataset/tokenize.py b/geniml/dataset/tokenize.py new file mode 100644 index 00000000..b964e9f5 --- /dev/null +++ b/geniml/dataset/tokenize.py @@ -0,0 +1,92 @@ +"""Tokenization + windowing core. + +The pure, side-effect-free half of the dataset: turn one region set into a list of +fixed-size token windows. Generalizes the offline atacformer pretokenization +(``benchmarking/bedbase_bulk/tokenization/pretokenize.py``): + +1. tokenize the region set against the universe, +2. drop ``unk`` tokens (regions that hit no universe interval), +3. window the remaining ids into ``context_size`` chunks by *sample-and-remove*, +4. skip pathologically large files (more than ``max_windows`` windows' worth). + +No source, no dataset, no I/O -- just ids in, windows out -- so it is trivially +unit-testable and reusable by any model. +""" + +import random +from typing import List, Optional + + +def sample_and_remove( + ids: List[int], context_size: int, rng: Optional[random.Random] = None +) -> List[List[int]]: + """Window a flat token list into ``context_size`` chunks by sampling without replacement. + + Repeatedly draw ``context_size`` tokens uniformly at random (without + replacement) and remove them, until fewer than ``context_size`` remain; the + final partial window is kept if non-empty. This randomizes which regions + co-occur in a window across the whole file, which is what we want for a + set-based (order-invariant) model -- it removes the positional bias that plain + front-to-back chunking would bake in. Implemented as a shuffle-then-chunk, which + is distributionally identical to iterated sample-and-remove but O(n) instead of + O(n^2). + + Args: + ids: flat list of token ids for one region set. + context_size: number of tokens per window. + rng: optional ``random.Random`` for reproducible shuffling. If ``None``, a + fresh unseeded generator is used. + + Returns: + List[List[int]]: the windows. Each has length ``context_size`` except + possibly the last, which may be shorter (the collator pads it). + """ + if context_size <= 0: + raise ValueError("context_size must be a positive integer") + if not ids: + return [] + pool = list(ids) + (rng or random).shuffle(pool) + return [pool[i : i + context_size] for i in range(0, len(pool), context_size)] + + +def tokenize_regionset( + region_set, + tokenizer, + context_size: int = 8192, + max_windows: Optional[int] = 10, + rng: Optional[random.Random] = None, + drop_unk: bool = True, +) -> List[List[int]]: + """Tokenize one region set into a list of fixed-size token windows. + + Args: + region_set: a ``gtars.models.RegionSet`` (or anything the tokenizer accepts). + tokenizer: a ``gtars.tokenizers.Tokenizer`` (or ``TrainingTokenizer``); called + as ``tokenizer(region_set)["input_ids"]``. + context_size: tokens per window. + max_windows: skip the file if it would produce more than this many windows + (i.e. if the filtered token count exceeds ``max_windows * context_size``). + ``None`` disables the cap. + rng: optional ``random.Random`` for reproducible windowing. + drop_unk: drop ``unk`` tokens (regions with no universe overlap) before + windowing. + + Returns: + List[List[int]]: token windows for this region set (empty if the file was + skipped or produced no tokens). + """ + ids = tokenizer(region_set)["input_ids"] + + if drop_unk: + unk_id = getattr(tokenizer, "unk_token_id", None) + if unk_id is not None: + ids = [i for i in ids if i != unk_id] + + if not ids: + return [] + + if max_windows is not None and len(ids) > max_windows * context_size: + return [] + + return sample_and_remove(ids, context_size, rng=rng) diff --git a/geniml/region2vec/region_shuffling.py b/geniml/region2vec/region_shuffling.py index eb407660..8ccf8051 100644 --- a/geniml/region2vec/region_shuffling.py +++ b/geniml/region2vec/region_shuffling.py @@ -24,17 +24,19 @@ class BEDDataset: def __init__(self, file_list: str) -> None: """Initializes a BEDDataset object. + The file-list handling is expressed on the shared dataset interface via + :class:`geniml.dataset.FileListSource`; the sentence-generation methods + below remain region2vec-specific (word2vec hard tokenization). + Args: file_list (str): A file storing a list of BED file names that should be included in the dataset. """ - self.filename_list = [] - with open(file_list, "r") as f: - for idx, line in enumerate(f): - filename = line.strip() - self.filename_list.append(filename) + from geniml.dataset import FileListSource - self.nfiles = len(self.filename_list) + self._source = FileListSource(file_list) + self.filename_list = list(self._source.paths) + self.nfiles = len(self._source) def regions2sentences_sampling(self, src_path: str, dst_path: str) -> None: """Constructs a sentence by sampling regions from a BED file. diff --git a/mkdocs.yml b/mkdocs.yml index 07bc68b6..03b6c3d3 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -19,6 +19,7 @@ nav: - Cell-type prediction using KNN: tutorials/cell-type-annotation-with-knn.md - Tokenization: tutorials/tokenization.md - Tokenize a BED file on the command line: tutorials/cli-tokenization.md + - Build a training dataset from bedbase: tutorials/bedbase-training-dataset.md - Create consensus peaks: tutorials/create-consensus-peaks.md - Fine-tune embeddings: tutorials/fine-tune-region2vec-model.md - Randomize bed files: tutorials/bedshift.md diff --git a/setup.py b/setup.py index ce7f7d66..339520f5 100755 --- a/setup.py +++ b/setup.py @@ -54,6 +54,7 @@ def _read_reqs(path): "geniml.assess", "geniml.bedspace", "geniml.bedshift", + "geniml.dataset", "geniml.eval", "geniml.likelihood", "geniml.models", diff --git a/tests/test_dataset.py b/tests/test_dataset.py new file mode 100644 index 00000000..e2bc7b02 --- /dev/null +++ b/tests/test_dataset.py @@ -0,0 +1,199 @@ +"""Tests for geniml.dataset: source interface, windowing, selection, dataset.""" + +import os +import random + +import pytest + +from gtars.models import RegionSet +from gtars.tokenizers import Tokenizer + +from geniml.dataset import ( + BedSetSource, + FileListSource, + RegionSetSource, + SampleRef, + SampleSelection, + TokenizedRegionDataset, + cache_key, + sample_and_remove, + select_bedbase_samples, + tokenize_regionset, +) +from geniml.dataset.bedbase import _compliance_ok, _record_to_sampleref +from geniml.io import BedSet + +DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") +UNIVERSE = os.path.join(DATA, "universe.bed") +BED1 = os.path.join(DATA, "to_tokenize.bed") +BED2 = os.path.join(DATA, "to_tokenize2.bed") + + +@pytest.fixture +def tokenizer(): + return Tokenizer.from_bed(UNIVERSE) + + +# ---------------------------------------------------------------- windowing + + +def test_sample_and_remove_windows_cover_all_tokens(): + ids = list(range(10)) + windows = sample_and_remove(ids, context_size=4, rng=random.Random(0)) + assert [len(w) for w in windows] == [4, 4, 2] + # every token appears exactly once across windows + assert sorted(t for w in windows for t in w) == ids + + +def test_sample_and_remove_empty_and_bad_context(): + assert sample_and_remove([], 8) == [] + with pytest.raises(ValueError): + sample_and_remove([1, 2], 0) + + +def test_sample_and_remove_deterministic_with_seed(): + ids = list(range(20)) + a = sample_and_remove(ids, 5, rng=random.Random(42)) + b = sample_and_remove(ids, 5, rng=random.Random(42)) + assert a == b + + +def test_tokenize_regionset_filters_unk(tokenizer): + rs = RegionSet(BED1) + windows = tokenize_regionset(rs, tokenizer, context_size=8192, max_windows=10) + unk = tokenizer.unk_token_id + assert all(unk not in w for w in windows) + assert sum(len(w) for w in windows) > 0 + + +def test_tokenize_regionset_skips_huge_files(tokenizer): + rs = RegionSet(BED1) + # max_windows tiny + context_size 1 forces the skip branch + windows = tokenize_regionset(rs, tokenizer, context_size=1, max_windows=1) + assert windows == [] + + +# ------------------------------------------------------------------ sources + + +def test_filelist_source_from_list_iterates_and_manifest(): + source = FileListSource([BED1, BED2]) + assert isinstance(source, RegionSetSource) + assert len(source) == 2 + items = list(source) + assert len(items) == 2 + rs, meta = items[0] + assert isinstance(rs, RegionSet) + assert meta["name"] == "to_tokenize.bed" + assert source.manifest() == sorted([BED1, BED2]) + + +def test_filelist_source_from_textfile(tmp_path): + listing = tmp_path / "beds.txt" + listing.write_text(f"{BED1}\n{BED2}\n") + source = FileListSource(str(listing)) + assert len(source) == 2 + + +def test_bedset_source_wraps_bedset(): + bs = BedSet([BED1, BED2]) + source = BedSetSource(bs) + assert isinstance(source, RegionSetSource) + assert len(source) == 2 + items = list(source) + assert len(items) == 2 + assert all("id" in meta for _, meta in items) + + +# --------------------------------------------------------- bedbase selection + + +def test_compliance_filter(): + assert _compliance_ok("bed6+0", "bed3+0") + assert not _compliance_ok("bed2+0", "bed3+0") + assert _compliance_ok("weird-format", "bed3+0") # unknown -> kept + + +def test_record_to_sampleref_maps_metadata(): + record = { + "id": "abc", + "name": "sample1", + "genome_alias": "hg38", + "annotation": { + "organism": "Homo sapiens", + "cell_type": "Tcell", + "assay": "ATAC-seq", + }, + } + ref = _record_to_sampleref(record) + assert ref.id == "abc" + assert ref.genome == "hg38" + assert ref.meta["cell_type"] == "Tcell" + assert ref.meta["species_name"] == "Homo sapiens" # organism -> species_name + assert "assay" in ref.meta + + +def test_sample_selection_roundtrip(tmp_path): + sel = SampleSelection( + genome="hg38", + qc="good", + bedbase_api="https://api.bedbase.org", + snapshot_date="2026-07-18", + samples=[SampleRef(id="a", name="A", genome="hg38", meta={"assay": "ATAC-seq"})], + ) + path = tmp_path / "manifest.json" + sel.to_json(str(path)) + loaded = SampleSelection.from_json(str(path)) + assert loaded.ids() == ["a"] + assert loaded.samples[0].meta["assay"] == "ATAC-seq" + assert loaded.snapshot_date == "2026-07-18" + + +# ------------------------------------------------------------- cache keying + + +def test_cache_key_is_stable_and_selection_sensitive(tokenizer): + s1 = FileListSource([BED1, BED2]) + s2 = FileListSource([BED2, BED1]) # order-independent + s3 = FileListSource([BED1]) + k1 = cache_key(s1, tokenizer, 8192) + k2 = cache_key(s2, tokenizer, 8192) + k3 = cache_key(s3, tokenizer, 8192) + assert k1 == k2 + assert k1 != k3 + assert cache_key(s1, tokenizer, 4096) != k1 # context size matters + + +# ------------------------------------------------------- materialize dataset + + +def test_materialize_dataset_and_cache(tmp_path, tokenizer): + pytest.importorskip("datasets") + source = FileListSource([BED1, BED2]) + ds_builder = TokenizedRegionDataset( + source, tokenizer, context_size=2, cache_dir=str(tmp_path), max_windows_per_file=100 + ) + ds = ds_builder.materialize() + assert len(ds) > 0 + assert "input_ids" in ds.column_names + # a parquet cache was written and reloads + cached = [f for f in os.listdir(tmp_path) if f.endswith(".parquet")] + assert cached + ds2 = ds_builder.materialize() + assert len(ds2) == len(ds) + + +# ----------------------------------------------------------- network (opt-in) + + +@pytest.fixture +def bedbase(request): + if not request.config.getoption("--bedbase"): + pytest.skip("use --bedbase to run bedbase network tests") + + +def test_select_bedbase_samples_live(bedbase): + selection = select_bedbase_samples(genome="hg38", qc="good", limit=5) + assert len(selection) == 5 + assert all(s.genome == "hg38" for s in selection.samples) + assert all(s.id for s in selection.samples) diff --git a/tests/test_scembed.py b/tests/test_scembed.py index 3ec39631..6720e3e3 100644 --- a/tests/test_scembed.py +++ b/tests/test_scembed.py @@ -3,6 +3,11 @@ import sys import pytest + +# scanpy (and scembed, which imports it) require the `sc` optional-dep group; skip +# the whole module gracefully when it isn't installed (e.g. an [ml]-only CI run). +pytest.importorskip("scanpy") + import scanpy as sc from geniml.region2vec.utils import Region2VecDataset from geniml.scembed.main import ScEmbed diff --git a/tests/test_search.py b/tests/test_search.py index 612f40e9..6462e799 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -4,6 +4,12 @@ import numpy as np import pytest + +# geniml.search.backends imports qdrant_client at module load; it lives in the +# `search` optional-dep group. Skip the whole module when it isn't installed +# (execution is further gated by --qdrant/--huggingface). +pytest.importorskip("qdrant_client") + from geniml.io import RegionSet from geniml.region2vec.main import Region2VecExModel from geniml.search import BED2BEDSearchInterface, BED2Vec, Text2BEDSearchInterface, Text2Vec diff --git a/tests/test_text2bednn.py b/tests/test_text2bednn.py index c31917c0..e85e8dfb 100644 --- a/tests/test_text2bednn.py +++ b/tests/test_text2bednn.py @@ -2,6 +2,12 @@ import numpy as np import pytest + +# geniml.search.backends imports qdrant_client at module load; it lives in the +# `search` optional-dep group. Skip the whole module when it isn't installed +# (execution is further gated by --huggingface). +pytest.importorskip("qdrant_client") + from geniml.search.backends import HNSWBackend from geniml.text2bednn.text2bednn import Vec2VecFNN from geniml.text2bednn.utils import metadata_dict_from_csv