diff --git a/config/tol_hybrid_to_wds_example.yaml b/config/tol_hybrid_to_wds_example.yaml new file mode 100644 index 0000000..45e023d --- /dev/null +++ b/config/tol_hybrid_to_wds_example.yaml @@ -0,0 +1,46 @@ +account: "PAS2136" + +# Input directory should point to the root of the parquet_to_hdf5 outputs. +path_to_input: "/fs/ess/PAS2136/TreeOfLife_HDF5/converted" +# Working directory for toolbox artifacts (logs, checkpoints, etc.). +path_to_output_folder: "/fs/scratch/PAS2136/TreeOfLife_hdf5_to_wds" + +tools_parameters: + num_workers: 8 + max_nodes: 2 + workers_per_node: 4 + cpu_per_worker: 32 + threshold_size: 224 + new_resize_size: 720 + +scripts: + tools_submitter: "scripts/tools_submit.sh" + tools_filter_script: "scripts/tools_filter.slurm" + tools_scheduling_script: "scripts/tools_scheduler.slurm" + tools_worker_script: "scripts/tools_worker.slurm" + tools_verification_script: "scripts/tools_verifier.slurm" + +output_structure: + urls_folder: "urls" + logs_folder: "logs" + images_folder: "images" + schedules_folder: "schedules" + profiles_table: "profiles.csv" + ignored_table: "ignored.csv" + inner_checkpoint_file: "inner_checkpoint.yaml" + tools_folder: "tools" + +tol_hybrid_to_wds: + shard_size: 10000 + shard_limit: 0 + metadata_glob: "**/*_metadata.parquet" + include_sources: [] + include_servers: [] + tar_output_root: "/fs/ess/PAS2136/TreeOfLife_WDS/shards" + resize_size: 224 + shard_prefix: "shard" + runner_timeout_seconds: 3600 + # Lookup selecting which images to convert: a parquet/CSV with a `uuid` column. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is required (an optional `hdf5_path` or `h5_file` column overrides it). + lookup_table_path: "/fs/ess/PAS2136/TreeOfLife/lookups/subset.parquet" + # Resolved-taxonomy parquet, joined by `uuid`; required for the text prompts. + taxa_glob: "/fs/ess/PAS2136/TreeOfLife/annotations/resolved_taxa/source=*/*.parquet" diff --git a/docs/README.md b/docs/README.md index f066e3c..4923c34 100644 --- a/docs/README.md +++ b/docs/README.md @@ -24,3 +24,9 @@ The dataset is divided into the following components, each with its own download 4. Monitor progress using the provided utilities. Each download component has its own README with detailed instructions. Click on the links above to access specific guidance for accessing images from each data provider. + +# Create webdataset + +Instructions for packaging the dataset into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards for model training, by storage format: + +- [Hybrid format (HDF5 + Parquet)](hybrid_to_webdataset_README.md) - create WebDataset shards from the hybrid HDF5-image + Parquet-metadata storage format. diff --git a/docs/hybrid_to_webdataset_README.md b/docs/hybrid_to_webdataset_README.md new file mode 100644 index 0000000..4417d89 --- /dev/null +++ b/docs/hybrid_to_webdataset_README.md @@ -0,0 +1,63 @@ +# Create a WebDataset from the hybrid HDF5 + Parquet format + +This guide covers packaging the Tree of Life dataset, stored in the hybrid format of WebP images in HDF5 (`*_images.h5`) alongside Parquet metadata (`*_metadata.parquet`), into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards for model training. It is implemented by the [`tol_hybrid_to_wds`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) tool. + +## What it produces + +One `shard-XXXXX.tar` per `shard_size` samples. Each sample is keyed by `uuid` and contains a JPEG plus ten text-prompt sidecars used for BioCLIP-style training: + +| | | +|---|---| +| `scientific_name.txt`, `common_name.txt`, `taxonomic_name.txt` | raw label values | +| `sci.txt`, `com.txt`, `taxon.txt`, `taxonTag.txt` | `a photo of {value}.` captions | +| `sci_com.txt`, `taxon_com.txt`, `taxonTag_com.txt` | the caption with ` with common name {common}.` appended | + +When no vernacular common name is available, the common name falls back to the scientific name (without author citation). + +## Inputs + +All three are joined by `uuid`: + +1. **Hybrid dataset**: the root holding `*_images.h5` (WebP bytes under an `/images/` group) and the paired `*_metadata.parquet`. +2. **Resolved taxonomy**: Parquet providing `scientific_name`, `common_name`, and the `kingdom` through `species` ranks per `uuid` (the text-label source). +3. **Lookup table**: Parquet/CSV with at least a `uuid` column; only listed UUIDs are converted. The `.h5` path is auto-derived as the sibling `*_images.h5` of each `*_metadata.parquet`, so no path column is needed; an optional `hdf5_path` (or `h5_file`) column overrides it. To build a lookup for a UUID subset from a converted dataset, use [`generate_hdf5_lookup.py`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py). + +## Prerequisites + +The pipeline runs on a cluster through Slurm, using Spark for the filter stage and MPI for the workers. This holds for any run size: for a small or low-resource job you scale the pipeline down through the config (see below) rather than switching to a different code path, so Spark and MPI are always required. + +Install the toolbox (see the top-level [installation instructions](../README.md#installation-instructions)): + +```bash +pip install -e . +``` + +The Slurm scripts encode site-specific choices you must adapt: + +- An **MPI** implementation with `mpi4py` installed against it. The workers run under MPI; if `mpi4py` is built against a different MPI than the one loaded at runtime, they fail to load it. +- A **Spark** runtime for the filter stage. +- The **partition** and **module** names in `scripts/tools_*.slurm`, and the virtual environment those scripts activate (by default the repo-root `.venv`, so install the toolbox there or point the scripts at your environment). + +Edit `scripts/tools_*.slurm` to load your cluster's MPI and Spark modules, target your partition, and activate your environment. As committed they target one cluster (OSC Cardinal) as a worked example to replace with your own. + +## Running the pipeline + +The toolbox runner reads one config and submits the whole pipeline (Spark filter, scheduler, MPI workers, verifier) with the right Slurm dependencies. Start from [`config/tol_hybrid_to_wds_example.yaml`](../config/tol_hybrid_to_wds_example.yaml) and set, for your run: + +- `account` and `tools_parameters` (nodes, workers, CPUs) for your allocation; +- `path_to_input` (the hybrid dataset root) and `path_to_output_folder` (the toolbox working directory); +- in the `tol_hybrid_to_wds` block: `tar_output_root`, `resize_size`, `taxa_glob` (required for the taxonomy/common-name prompts), and `lookup_table_path` (a parquet/CSV with a `uuid` column selecting which images to convert). + +For a small or low-resource run, scale `tools_parameters` down (for example `num_workers: 1`, `max_nodes: 1`, `workers_per_node: 1`); the execution path is identical. See the [tool README](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md) for the full key reference. + +Then run the pipeline with a single command: + +```bash +tree_of_life_toolbox config/tol_hybrid_to_wds_my_run.yaml tol_hybrid_to_wds +``` + +The runner reads the config, exports the environment it implies (account, node/worker counts, output folders), and submits the four chained stages. The **filter** discovers the `*_metadata.parquet` files under `path_to_input`, derives each sibling `*_images.h5`, joins the taxonomy, and restricts to the lookup UUIDs; the **scheduler** assigns shards to MPI ranks; the **workers** read the WebP bytes, resize to JPEG, write the ten prompts, and emit the tars; the **verifier** marks completion. The underlying Slurm scripts are described in [scripts_README](scripts_README.md). + +## Validating output + +[`validate_tars.py`](../src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py) checks shard integrity: that every tar is readable and each sample carries the JPEG and all ten text sidecars. diff --git a/pyproject.toml b/pyproject.toml index 02b2e82..792ccfe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,11 +30,13 @@ dependencies = [ "cramjam", "cython", "fsspec", + "h5py", "inflate64", "mpi4py", "multivolumefile", "opencv-python", "pandas", + "polars", "pathspec", "pillow", "psutil", @@ -51,6 +53,7 @@ dependencies = [ "texttable", "trove-classifiers", "typing-extensions", + "webdataset", "wheel" ] diff --git a/scripts/tools_filter.slurm b/scripts/tools_filter.slurm index 6aee3f6..f6d8c69 100644 --- a/scripts/tools_filter.slurm +++ b/scripts/tools_filter.slurm @@ -1,7 +1,8 @@ #!/bin/bash #SBATCH --job-name tool_filter -#SBATCH --mem=0 -#SBATCH --time=01:00:00 +#SBATCH --partition=cpu +#SBATCH --mem=200G +#SBATCH --time=2:00:00 if [ "$#" -eq 0 ]; then echo "Usage: $0 tool_name" @@ -10,17 +11,21 @@ fi tool_name=$1 -logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" +logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER:-${SCRATCH:-$PWD}/tools_logs}" mkdir -p "$logs_dir" driver_memory="110G" executor_memory="64G" -module load spark/3.4.1 -module load miniconda3/23.3.1-py310 +module load spark/3.5.1 +module load python/3.12 +if [ -z "$REPO_ROOT" ]; then + echo "REPO_ROOT is not set. Please export REPO_ROOT to the repository path." >&2 + exit 1 +fi source "${REPO_ROOT}/.venv/bin/activate" -pbs-spark-submit \ +spark-submit \ --driver-memory $driver_memory \ --executor-memory $executor_memory \ "${TOOLBOX_PATH}/main/filter.py" \ diff --git a/scripts/tools_scheduler.slurm b/scripts/tools_scheduler.slurm index ea35a32..8b159e4 100644 --- a/scripts/tools_scheduler.slurm +++ b/scripts/tools_scheduler.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_scheduler -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=00:05:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes=1 \ --ntasks-per-node=1 \ --cpus-per-task=1 \ - --mem=0 \ --output="${logs_dir}/tool_scheduler.log" \ python "${TOOLBOX_PATH}/main/scheduler.py" "${tool_name}" diff --git a/scripts/tools_submit.sh b/scripts/tools_submit.sh old mode 100644 new mode 100755 diff --git a/scripts/tools_verifier.slurm b/scripts/tools_verifier.slurm index 6a3b75e..ac332f5 100644 --- a/scripts/tools_verifier.slurm +++ b/scripts/tools_verifier.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_verifier -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=00:05:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes=1 \ --ntasks-per-node=1 \ --cpus-per-task=1 \ - --mem=0 \ --output="${logs_dir}/tool_verifier.log" \ python "${TOOLBOX_PATH}/main/verification.py" "${tool_name}" diff --git a/scripts/tools_worker.slurm b/scripts/tools_worker.slurm index 4856e62..601bd8d 100644 --- a/scripts/tools_worker.slurm +++ b/scripts/tools_worker.slurm @@ -1,6 +1,6 @@ #!/bin/bash #SBATCH --job-name tool_worker -#SBATCH --mem=0 +#SBATCH --partition=cpu #SBATCH --time=03:00:00 if [ "$#" -eq 0 ]; then @@ -14,17 +14,14 @@ logs_dir="${OUTPUT_TOOLS_LOGS_FOLDER}" mkdir -p "$logs_dir" module load intel/2021.10.0 -module load intelmpi/2021.10 -module load miniconda3/23.3.1-py310 +module load intel-oneapi-mpi/2021.10.0 source "${REPO_ROOT}/.venv/bin/activate" export PYARROW_IGNORE_TIMEZONE=1 -export I_MPI_JOB_RESPECT_PROCESS_PLACEMENT=0 srun \ --mpi=pmi2 \ --nodes="$TOOLS_MAX_NODES" \ --ntasks-per-node="$TOOLS_WORKERS_PER_NODE" \ --cpus-per-task="$TOOLS_CPU_PER_WORKER" \ - --mem=0 \ --output="${logs_dir}/tool_worker-%2t.log" \ python "${TOOLBOX_PATH}/main/runner.py" "${tool_name}" diff --git a/src/TreeOfLife_toolbox/__init__.py b/src/TreeOfLife_toolbox/__init__.py index c1a2a93..b1dd744 100644 --- a/src/TreeOfLife_toolbox/__init__.py +++ b/src/TreeOfLife_toolbox/__init__.py @@ -1,2 +1,7 @@ -from TreeOfLife_toolbox import tol200m_fathom_net_crop, tol200m_bioscan_data_transfer, data_transfer +from TreeOfLife_toolbox import ( + tol200m_fathom_net_crop, + tol200m_bioscan_data_transfer, + data_transfer, + tol_hybrid_to_wds, +) diff --git a/src/TreeOfLife_toolbox/main/filter.py b/src/TreeOfLife_toolbox/main/filter.py index ed526c5..b8c91a2 100644 --- a/src/TreeOfLife_toolbox/main/filter.py +++ b/src/TreeOfLife_toolbox/main/filter.py @@ -46,3 +46,10 @@ checkpoint["filtering_completed"] = True else: logger.info("Filtering was already completed") + + # Spark's driver JVM does not reliably exit after the context is stopped, which leaves the batch job lingering until its walltime even though the work is done. All output is already written and checkpointed, so stop the session and terminate immediately. + try: + tool_filter.spark.stop() + except Exception: + pass + os._exit(0) diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md new file mode 100644 index 0000000..756b15a --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/README.md @@ -0,0 +1,46 @@ +# tol_hybrid_to_wds + +Converts the hybrid Tree of Life dataset (HDF5 images + Parquet metadata) into [WebDataset](https://github.com/webdataset/webdataset) `.tar` shards that mirror the taxonomy-rich text sidecars produced by the previous Parquet-based pipeline. + +## Pipeline steps + +1. **Filter**: Scans the converted dataset for `*_metadata.parquet` files, enriches each row with the matching `*_images.h5` path, and writes shard-sized parquet partitions under `tools/tol_hybrid_to_wds/shard_metadata`. +2. **Scheduler**: Assigns shard identifiers to MPI ranks and writes `schedule.csv` for the toolbox runner scripts. +3. **Runner**: For every scheduled shard, loads the shard metadata, reads WebP-compressed bytes from the referenced HDF5 file, converts them to JPEG (with an optional resize), generates ten taxonomy/common-name text prompts, and emits `shard-XXXXX.tar` archives via `webdataset.TarWriter`. + +## Configuration highlights + +Add a `tol_hybrid_to_wds` block to the standard toolbox config template: + +```yaml +tol_hybrid_to_wds: + shard_size: 10000 # rows per shard + shard_limit: 0 # optional cap, 0 disables + metadata_glob: "**/*_metadata.parquet" + tar_output_root: "/fs/.../wds" # defaults to tools folder if omitted + resize_size: 224 # square resize in pixels (0 keeps original) + lookup_table_path: "/fs/.../subset.parquet" # optional UUID list (CSV or Parquet) + lookup_path_column: "path" # only needed for CSVs with raw paths + taxa_glob: "/fs/.../annotations/source=*/*.parquet" # external taxonomy parquet + # If the lookup includes an `hdf5_path` column it will override the auto-resolved path. +``` + +Set `path_to_input` to the root of the converted dataset (where the `*_metadata.parquet` and `*_images.h5` files live) and `path_to_output_folder` to the working directory for toolbox artifacts. + +If `lookup_table_path` is provided, it can be either CSV or Parquet and must contain at least a `uuid` column. Include a column named `hdf5_path` to override the derived file path, or a `path` column (configurable via `lookup_path_column`) when providing CSVs that still reference the original `data_*.parquet` names. Only UUIDs listed in the lookup are converted. + +Run the tool via the standard CLI: + +```bash +CONFIG_PATH=config/tol_hybrid_to_wds_example.yaml \ +tree_of_life_toolbox config/tol_hybrid_to_wds_example.yaml tol_hybrid_to_wds +``` + +Each shard sample contains the JPEG plus ten text prompts: `scientific_name.txt`, `taxonomic_name.txt`, `sci.txt`, `taxon.txt`, `taxonTag.txt`, `common_name.txt`, `com.txt`, `sci_com.txt`, `taxon_com.txt`, and `taxonTag_com.txt`. Files ending in `_com` append the resolved common name (or `Unknown` when no vernacular exists) to their prompt text. + +## Companion scripts + +Standalone helpers live in [`scripts/`](scripts/) beside this tool: + +- `generate_hdf5_lookup.py` builds a `uuid`-to-`h5` lookup table (Parquet) from a converted dataset, optionally sampling a subset, for use as `lookup_table_path`. +- `validate_tars.py` checks output shard integrity: every tar is readable and each sample carries the JPEG and all ten text sidecars. diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py new file mode 100644 index 0000000..636ae60 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/__init__.py @@ -0,0 +1,7 @@ +""" +Tooling for converting HDF5-backed Tree of Life data into WebDataset shards. + +Importing this module registers filter/scheduler/runner implementations with the toolbox registry so they can be scheduled through the generic CLI. +""" + +from TreeOfLife_toolbox.tol_hybrid_to_wds import classes # noqa: F401 diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py new file mode 100644 index 0000000..700080c --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/classes.py @@ -0,0 +1,413 @@ +""" +Toolbox integration for converting HDF5-backed TreeOfLife data into WebDatasets. +""" + +from __future__ import annotations + +import glob +import os +from pathlib import Path +from typing import Dict, List + +import h5py +import pandas as pd +import polars as pl +import pyspark.sql.functions as func +from pyspark.sql.window import Window + +from TreeOfLife_toolbox.main.config import Config +from TreeOfLife_toolbox.main.filters import FilterRegister, SparkFilterToolBase +from TreeOfLife_toolbox.main.runners import MPIRunnerTool, RunnerRegister +from TreeOfLife_toolbox.main.schedulers import DefaultScheduler, SchedulerRegister +from TreeOfLife_toolbox.tol_hybrid_to_wds.utils import ( + convert_webp_to_jpeg, + generate_text_files, + init_shard_writer, +) + + +def _sanitize_path_column(column): + return func.regexp_replace(column, "^file:", "") + + +@FilterRegister("tol_hybrid_to_wds") +class TolHDF5ToWDSFilter(SparkFilterToolBase): + """ + Builds shard-level metadata derived from the parquet_to_hdf5 outputs. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hybrid_to_wds" + tool_cfg = self.config.get("tol_hybrid_to_wds", {}) or {} + self.shard_size = int(tool_cfg.get("shard_size", 10000)) + self.shard_limit = int(tool_cfg.get("shard_limit", 0)) + self.metadata_glob = tool_cfg.get("metadata_glob", "**/*_metadata.parquet") + self.include_sources = tool_cfg.get("include_sources") or [] + self.include_servers = tool_cfg.get("include_servers") or [] + self.taxa_glob = tool_cfg.get("taxa_glob") + self.lookup_table_path = tool_cfg.get("lookup_table_path") + self.lookup_path_column = tool_cfg.get("lookup_path_column", "path") + self.shard_metadata_dir = os.path.join( + self.tools_path, self.filter_name, "shard_metadata" + ) + + if self.shard_size <= 0: + raise ValueError("tol_hybrid_to_wds.shard_size must be greater than zero") + + os.makedirs(self.shard_metadata_dir, exist_ok=True) + + def _discover_metadata_files(self) -> List[str]: + dataset_root = self.config.get("path_to_input") + if not dataset_root or not os.path.exists(dataset_root): + raise ValueError("path_to_input must point to converted HDF5 dataset") + + pattern = os.path.join(dataset_root, self.metadata_glob) + files = glob.glob(pattern, recursive=True) + return sorted(files) + + def run(self): + metadata_files = self._discover_metadata_files() + if not metadata_files: + raise ValueError( + f"No metadata parquet files found using pattern: {self.metadata_glob}" + ) + + self.logger.info("Discovered %d metadata files", len(metadata_files)) + + # Read only the UUID column to avoid schema conflicts from unused fields + metadata_df = self.spark.read.parquet(*metadata_files).select("uuid") + metadata_df = metadata_df.withColumn( + "metadata_file", _sanitize_path_column(func.input_file_name()) + ) + metadata_df = metadata_df.withColumn( + "hdf5_path", + func.regexp_replace("metadata_file", "_metadata\\.parquet$", "_images.h5"), + ) + metadata_df = metadata_df.withColumn( + "base_name", + func.regexp_extract("metadata_file", r"([^/]+)_metadata\.parquet$", 1), + ) + + if "source" not in metadata_df.columns: + metadata_df = metadata_df.withColumn( + "source", func.regexp_extract("metadata_file", r"source=([^/]+)", 1) + ) + if "server" not in metadata_df.columns: + metadata_df = metadata_df.withColumn( + "server", func.regexp_extract("metadata_file", r"server=([^/]+)", 1) + ) + + if self.taxa_glob: + taxa_files = glob.glob(self.taxa_glob) + if not taxa_files: + raise ValueError( + f"No taxonomy parquet files found for pattern: {self.taxa_glob}" + ) + taxa_df = self.spark.read.parquet(*taxa_files) + taxa_columns = [ + col_name + for col_name in [ + "uuid", + "scientific_name", + "provided_common_name", + "common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + if col_name in taxa_df.columns + ] + taxa_df = taxa_df.select(*taxa_columns) + metadata_df = metadata_df.join(taxa_df, on="uuid", how="left") + + # Drop a row only when every taxonomy field is null or blank, so no degenerate text sidecars are ever written. Robust to incomplete upstream taxonomy and a no-op on inputs that are already complete. + taxonomy_fields = [c for c in taxa_columns if c != "uuid"] + all_empty = None + for col_name in taxonomy_fields: + blank = func.col(col_name).isNull() | (func.trim(func.col(col_name)) == "") + all_empty = blank if all_empty is None else (all_empty & blank) + if all_empty is not None: + metadata_df = metadata_df.filter(~all_empty) + + if self.include_sources: + metadata_df = metadata_df.filter(func.col("source").isin(self.include_sources)) + if self.include_servers: + metadata_df = metadata_df.filter(func.col("server").isin(self.include_servers)) + + if self.lookup_table_path: + lookup_path_lower = self.lookup_table_path.lower() + if lookup_path_lower.endswith(".parquet"): + lookup_df = self.spark.read.parquet(self.lookup_table_path) + else: + lookup_df = ( + self.spark.read.option("header", True).csv(self.lookup_table_path) + ) + + if "uuid" not in lookup_df.columns: + raise ValueError("lookup table must contain a 'uuid' column") + + has_lookup_path = ( + self.lookup_path_column in lookup_df.columns + if self.lookup_path_column + else False + ) + has_hdf5_override = "hdf5_path" in lookup_df.columns + if not has_hdf5_override and "h5_file" in lookup_df.columns: + lookup_df = lookup_df.withColumnRenamed("h5_file", "hdf5_path") + has_hdf5_override = True + + if has_lookup_path: + lookup_df = lookup_df.withColumn( + "lookup_base_name", + func.regexp_replace( + func.regexp_extract( + func.col(self.lookup_path_column), r"([^/]+)$", 1 + ), + r"\.parquet$", + "", + ), + ) + else: + lookup_df = lookup_df.withColumn("lookup_base_name", func.lit(None)) + + if has_hdf5_override: + lookup_df = lookup_df.withColumnRenamed("hdf5_path", "lookup_hdf5_path") + + select_cols = ["uuid", "lookup_base_name"] + if has_hdf5_override: + select_cols.append("lookup_hdf5_path") + + lookup_df = lookup_df.select(*select_cols).dropDuplicates(["uuid"]) + metadata_df = metadata_df.join(lookup_df, on="uuid", how="inner") + + if has_lookup_path: + metadata_df = metadata_df.filter( + func.col("lookup_base_name").isNull() + | (func.col("lookup_base_name") == func.col("base_name")) + ) + + if has_hdf5_override: + if "hdf5_path" in metadata_df.columns: + metadata_df = metadata_df.drop("hdf5_path") + metadata_df = metadata_df.withColumnRenamed( + "lookup_hdf5_path", "hdf5_path" + ) + + metadata_df = metadata_df.drop("lookup_base_name") + + # Drop duplicate UUIDs to avoid double writes + metadata_df = metadata_df.dropDuplicates(["uuid"]) + + required_columns = [ + "uuid", + "hdf5_path", + "source", + "server", + "scientific_name", + "common_name", + "provided_common_name", + "kingdom", + "phylum", + "class", + "order", + "family", + "genus", + "species", + ] + + string_columns = set(required_columns) + + for col_name in required_columns: + if col_name not in metadata_df.columns: + default_value = func.lit(None) + if col_name in string_columns: + default_value = default_value.cast("string") + metadata_df = metadata_df.withColumn(col_name, default_value) + + ordering_window = Window.orderBy("metadata_file", "uuid") + metadata_df = ( + metadata_df.withColumn( + "row_number", func.row_number().over(ordering_window) + ) + .withColumn( + "shard_id", + func.floor((func.col("row_number") - 1) / self.shard_size).cast("int"), + ) + .drop("row_number") + ) + + if self.shard_limit > 0: + metadata_df = metadata_df.filter(func.col("shard_id") < self.shard_limit) + + if metadata_df.count() == 0: + raise ValueError("No metadata rows left after filtering/shard limiting") + + selected_columns = required_columns + ["shard_id"] + sharded_df = metadata_df.select(*selected_columns) + + ( + sharded_df.repartition("shard_id") + .write.partitionBy("shard_id") + .mode("overwrite") + .parquet(self.shard_metadata_dir) + ) + self.logger.info("Wrote shard metadata to %s", self.shard_metadata_dir) + + schedule_df = ( + sharded_df.select("shard_id") + .dropDuplicates() + .withColumn( + "metadata_path", + func.concat( + func.lit(self.shard_metadata_dir), + func.lit("/shard_id="), + func.col("shard_id"), + ), + ) + .withColumn("server_name", func.lit("tol_hybrid_to_wds")) + .withColumn("partition_id", func.col("shard_id")) + ) + + self.save_filter(schedule_df) + total_shards = schedule_df.count() + self.logger.info("Prepared %d shards", total_shards) + + +@SchedulerRegister("tol_hybrid_to_wds") +class TolHDF5ToWDSScheduler(DefaultScheduler): + """ + Scheduler that treats each metadata shard as an independent work unit. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hybrid_to_wds" + # Scheduling keys only. The data columns (metadata_path, shard_id) ride along in the filter table and are picked up by the runner's data_scheme; including them here duplicates them into schedule.csv and collides on the runner's merge. + self.scheme = ["server_name", "partition_id"] + + +@RunnerRegister("tol_hybrid_to_wds") +class TolHDF5ToWDSRunner(MPIRunnerTool): + """ + Converts shard metadata into WebDataset tar archives. + """ + + def __init__(self, cfg: Config): + super().__init__(cfg) + self.filter_name = "tol_hybrid_to_wds" + self.data_scheme = ["metadata_path", "shard_id", "server_name", "partition_id"] + self.verification_scheme = ["server_name", "partition_id"] + + params = self.config.get("tol_hybrid_to_wds", {}) or {} + self.tar_output_root = params.get("tar_output_root") or os.path.join( + self.tools_path, self.filter_name, "tar_dataset" + ) + self.resize_size = int(params.get("resize_size", 224)) + self.shard_prefix = params.get("shard_prefix", "shard") + self.total_time = int(params.get("runner_timeout_seconds", 3600)) + + os.makedirs(self.tar_output_root, exist_ok=True) + + def _load_shard_metadata(self, metadata_path: str) -> pl.DataFrame: + metadata_dir = Path(metadata_path) + if not metadata_dir.exists(): + raise FileNotFoundError(f"Shard metadata directory missing: {metadata_path}") + + files = sorted(metadata_dir.glob("*.parquet")) + if not files: + raise FileNotFoundError(f"No Parquet files found under {metadata_path}") + + return pl.read_parquet([str(path) for path in files]) + + def _iter_records(self, shard_df: pl.DataFrame): + for row in shard_df.iter_rows(named=True): + yield row + + def _open_hdf5(self, path: str, cache: Dict[str, h5py.File]) -> h5py.File: + if path not in cache: + if not os.path.exists(path): + raise FileNotFoundError(f"HDF5 file missing: {path}") + cache[path] = h5py.File(path, "r") + return cache[path] + + def apply_filter( + self, filtering_df: pd.DataFrame, server_name: str, partition_id: int + ) -> int: + shard_id = int(filtering_df.iloc[0]["shard_id"]) + metadata_path = filtering_df.iloc[0]["metadata_path"] + return self.process_shard(shard_id, metadata_path) + + def process_shard(self, shard_id: int, metadata_path: str) -> int: + self.is_enough_time() + shard_df = self._load_shard_metadata(metadata_path) + if shard_df.height == 0: + self.logger.warning("Shard %s is empty, skipping", shard_id) + return 0 + + tar_writer, tar_path = init_shard_writer( + self.tar_output_root, shard_id, self.shard_prefix + ) + hdf5_cache: Dict[str, h5py.File] = {} + written = 0 + failed = 0 + + try: + for row in self._iter_records(shard_df): + uuid = row.get("uuid") + hdf5_path = row.get("hdf5_path") + if not uuid or not hdf5_path: + failed += 1 + continue + + try: + h5_file = self._open_hdf5(hdf5_path, hdf5_cache) + images_group = h5_file.get("images") + if images_group is None or uuid not in images_group: + failed += 1 + continue + + dataset = images_group[uuid][:] + webp_bytes = dataset.tobytes() + jpeg_bytes = convert_webp_to_jpeg(webp_bytes, self.resize_size) + + taxon_dict = { + "scientific_name": row.get("scientific_name"), + "common_name": row.get("common_name"), + "kingdom": row.get("kingdom"), + "phylum": row.get("phylum"), + "class": row.get("class"), + "order": row.get("order"), + "family": row.get("family"), + "genus": row.get("genus"), + "species": row.get("species"), + } + + sample = {"__key__": uuid, "jpg": jpeg_bytes} + for ext, content in generate_text_files(taxon_dict).items(): + sample[ext] = content.encode("utf-8") + + tar_writer.write(sample) + written += 1 + except Exception as exc: # noqa: BLE001 + failed += 1 + self.logger.error( + "Failed to process uuid %s in shard %s: %s", uuid, shard_id, exc + ) + finally: + tar_writer.close() + for handle in hdf5_cache.values(): + handle.close() + + self.logger.info( + "Shard %s written to %s (%s records, %s failures)", + shard_id, + tar_path, + written, + failed, + ) + return written diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py new file mode 100644 index 0000000..c368d9b --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/generate_hdf5_lookup.py @@ -0,0 +1,122 @@ +#!/usr/bin/env python +""" +Generate a UUID-to-HDF5 lookup table from a converted Tree of Life dataset. + +Usage: + python generate_hdf5_lookup.py \ + --data-root /fs/scratch/PAS2136/TreeOfLife_test-wds/data \ + --output /fs/scratch/PAS2136/TreeOfLife_test-wds/lookup-tables/all/hdf5_lookup.parquet \ + --sample 100 +""" + +from __future__ import annotations + +import argparse +import glob +import os +from typing import Iterable, List + +import polars as pl + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Create a lookup Parquet of UUIDs to HDF5/metadata paths." + ) + parser.add_argument( + "--data-root", + required=True, + help="Root directory containing source=*/server=*/data_*_{metadata,images} files.", + ) + parser.add_argument( + "--output", + required=True, + help="Destination Parquet file (e.g. /.../lookup-tables/all/lookup.parquet).", + ) + parser.add_argument( + "--sample", + type=float, + default=100.0, + help="Percentage of UUIDs to keep (0 < sample <= 100). Default: 100.", + ) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Random seed used when down-sampling.", + ) + return parser.parse_args() + + +def discover_metadata_files(data_root: str) -> List[str]: + pattern = os.path.join(data_root, "source=*", "server=*", "data_*_metadata.parquet") + return sorted(glob.glob(pattern)) + + +def build_lazy_frames(metadata_files: Iterable[str]) -> List[pl.LazyFrame]: + lazy_frames: List[pl.LazyFrame] = [] + for meta_path in metadata_files: + source = _extract_between(meta_path, "source=", "/server=") + server = _extract_between(meta_path, "server=", "/data_") + base_name = os.path.basename(meta_path).replace("_metadata.parquet", "") + hdf5_path = meta_path.replace("_metadata.parquet", "_images.h5") + + lazy = ( + pl.scan_parquet(meta_path) + .select("uuid") + .with_columns( + pl.lit(meta_path).alias("metadata_path"), + pl.lit(hdf5_path).alias("hdf5_path"), + pl.lit(source).alias("source"), + pl.lit(server).alias("server"), + pl.lit(base_name).alias("base_name"), + ) + ) + lazy_frames.append(lazy) + return lazy_frames + + +def _extract_between(value: str, start_token: str, end_token: str) -> str: + start_idx = value.find(start_token) + if start_idx == -1: + return "" + start_idx += len(start_token) + end_idx = value.find(end_token, start_idx) + if end_idx == -1: + return value[start_idx:] + return value[start_idx:end_idx] + + +def main() -> None: + args = parse_args() + data_root = os.path.abspath(args.data_root) + metadata_files = discover_metadata_files(data_root) + + if not metadata_files: + raise SystemExit( + f"No metadata parquet files found under {data_root}. " + "Expected pattern source=*/server=*/data_*_metadata.parquet" + ) + + lazy_frames = build_lazy_frames(metadata_files) + lookup_lazy = pl.concat(lazy_frames) + + if args.sample <= 0 or args.sample > 100: + raise SystemExit("--sample must be in the interval (0, 100].") + + if args.sample < 100: + frac = args.sample / 100.0 + lookup_lazy = lookup_lazy.sample( + fraction=frac, with_replacement=False, shuffle=True, seed=args.seed + ) + + lookup_df = lookup_lazy.collect(streaming=True) + + output_path = os.path.abspath(args.output) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + lookup_df.write_parquet(output_path, compression="zstd") + print(f"Wrote {lookup_df.height} records to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py new file mode 100644 index 0000000..9d2e262 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/scripts/validate_tars.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +import os +import json +import argparse +from pathlib import Path +from concurrent.futures import ProcessPoolExecutor +from time import time +import tarfile +from collections import defaultdict + +# The 10 required sidecar suffixes +REQUIRED_SIDECARS = [ + 'com.txt', + 'common_name.txt', + 'sci.txt', + 'sci_com.txt', + 'scientific_name.txt', + 'taxon.txt', + 'taxonTag.txt', + 'taxonTag_com.txt', + 'taxon_com.txt', + 'taxonomic_name.txt', +] + +def validate_tar(tar_path: str) -> tuple[str, dict]: + """ + Validate that each JPG has all 10 required text sidecars AND they contain text. + Returns (tar_name, validation_result) + """ + tar_name = os.path.basename(tar_path) + + # Track all files by UUID + files_by_uuid = defaultdict(dict) # uuid -> {extension: has_content} + jpg_uuids = set() + + try: + with tarfile.open(tar_path, 'r|') as tar: + for member in tar: + if not member.isfile(): + continue + + name = member.name + basename = os.path.basename(name) + + if '.' not in basename: + continue + + # Split on first dot to get UUID + parts = basename.split('.', 1) + if len(parts) != 2: + continue + + uuid = parts[0] + extension = parts[1] + + # Track JPGs + if extension.lower() == 'jpg': + jpg_uuids.add(uuid) + files_by_uuid[uuid][extension] = True # JPGs don't need content check + + # For text files, check if they have content + elif extension in REQUIRED_SIDECARS: + has_content = False + try: + f = tar.extractfile(member) + if f: + content = f.read() + # Check if file has any non-whitespace content + has_content = len(content.strip()) > 0 + f.close() + except Exception: + has_content = False + + files_by_uuid[uuid][extension] = has_content + + # Now validate: check each JPG has all sidecars with content + issues = [] + for uuid in sorted(jpg_uuids): + extensions = files_by_uuid[uuid] + missing = [] + empty = [] + + for required in REQUIRED_SIDECARS: + if required not in extensions: + missing.append(required) + elif not extensions[required]: + empty.append(required) + + if missing or empty: + issue = {'uuid': uuid} + if missing: + issue['missing'] = missing + if empty: + issue['empty'] = empty + issues.append(issue) + + return (tar_name, { + 'valid': len(issues) == 0, + 'jpg_count': len(jpg_uuids), + 'issues': issues + }) + + except Exception as e: + return (tar_name, { + 'valid': False, + 'error': str(e), + 'jpg_count': 0, + 'issues': [] + }) + +def process_chunk(tar_paths: list[str]) -> list[tuple[str, dict]]: + """Process multiple tars in one worker.""" + return [validate_tar(p) for p in tar_paths] + +def chunks(lst, n): + """Yield successive n-sized chunks.""" + for i in range(0, len(lst), n): + yield lst[i:i + n] + +def validate_all_tars(source_dir: str, workers: int, chunk_size: int = 1) -> None: + src = Path(source_dir) + tar_files = sorted([str(p) for p in src.glob("*.tar")]) + n = len(tar_files) + + if n == 0: + raise SystemExit(f"No .tar files in {src}") + + print(f"Validating {n} shards with {workers} workers...", flush=True) + print("Checking: file existence + non-empty content", flush=True) + + results = {} + all_issues = [] + total_jpgs = 0 + valid_tars = 0 + total_missing = 0 + total_empty = 0 + + t0 = time() + processed = 0 + + # Split work into chunks + tar_chunks = list(chunks(tar_files, chunk_size)) + + with ProcessPoolExecutor(max_workers=workers) as ex: + for result_chunk in ex.map(process_chunk, tar_chunks): + for tar_name, validation in result_chunk: + results[tar_name] = validation + + jpg_count = validation.get('jpg_count', 0) + total_jpgs += jpg_count + + if validation['valid']: + valid_tars += 1 + else: + # Record issues for this tar + for issue in validation.get('issues', []): + issue_record = { + 'tar': tar_name, + 'uuid': issue['uuid'] + } + if 'missing' in issue: + issue_record['missing'] = issue['missing'] + total_missing += len(issue['missing']) + if 'empty' in issue: + issue_record['empty'] = issue['empty'] + total_empty += len(issue['empty']) + all_issues.append(issue_record) + + # Also note if there was an error + if 'error' in validation: + all_issues.append({ + 'tar': tar_name, + 'error': validation['error'] + }) + + processed += 1 + if processed % 500 == 0 or processed == n: + elapsed = time() - t0 + rate = processed / elapsed + eta = (n - processed) / rate if rate > 0 else 0 + print(f"[{processed}/{n}] {rate:.1f}/s | Valid: {valid_tars}/{processed} | Issues: {len(all_issues)} (missing: {total_missing}, empty: {total_empty}) | ETA: {eta:.0f}s", flush=True) + + # Write detailed results + out_path = src / "tar_validation_results.json" + summary = { + 'summary': { + 'total_tars': n, + 'valid_tars': valid_tars, + 'invalid_tars': n - valid_tars, + 'total_jpgs': total_jpgs, + 'total_issues': len(all_issues), + 'total_missing_files': total_missing, + 'total_empty_files': total_empty + }, + 'per_tar': results + } + + with open(out_path, "w") as f: + json.dump(summary, f, indent=2) + + # Write issue log (more human-readable) + if all_issues: + issues_path = src / "tar_validation_issues.log" + with open(issues_path, "w") as f: + f.write(f"Validation Issues - {len(all_issues)} total\n") + f.write(f"Missing files: {total_missing}, Empty files: {total_empty}\n") + f.write("=" * 80 + "\n\n") + + # Group by tar file + by_tar = defaultdict(list) + for issue in all_issues: + by_tar[issue['tar']].append(issue) + + for tar_name in sorted(by_tar.keys()): + issues = by_tar[tar_name] + f.write(f"TAR: {tar_name} ({len(issues)} issues)\n") + f.write("-" * 80 + "\n") + + for issue in issues: + if 'error' in issue: + f.write(f" ERROR: {issue['error']}\n") + else: + f.write(f" UUID: {issue['uuid']}\n") + if 'missing' in issue: + f.write(f" Missing: {', '.join(issue['missing'])}\n") + if 'empty' in issue: + f.write(f" Empty: {', '.join(issue['empty'])}\n") + + f.write("\n") + + print(f"\n⚠️ Issues found: {len(all_issues)}") + print(f" Missing files: {total_missing}") + print(f" Empty files: {total_empty}") + print(f" Wrote details to: {issues_path}") + + elapsed = time() - t0 + print(f"\n✓ Done in {elapsed:.1f}s ({n/elapsed:.1f} shards/s)") + print(f"✓ Valid tars: {valid_tars}/{n} ({100*valid_tars/n:.1f}%)") + print(f"✓ Total JPGs: {total_jpgs:,}") + print(f"✓ Wrote summary to: {out_path}") + +if __name__ == "__main__": + ap = argparse.ArgumentParser(description="Validate JPG sidecars in tar shards (checks content)") + ap.add_argument("--source_dir", required=True, help="Directory containing .tar files") + ap.add_argument("--workers", type=int, default=96, help="Number of parallel workers") + ap.add_argument("--chunk-size", type=int, default=1, help="Tars per worker batch") + args = ap.parse_args() + + validate_all_tars(args.source_dir, args.workers, args.chunk_size) + diff --git a/src/TreeOfLife_toolbox/tol_hybrid_to_wds/utils.py b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/utils.py new file mode 100644 index 0000000..14c04b8 --- /dev/null +++ b/src/TreeOfLife_toolbox/tol_hybrid_to_wds/utils.py @@ -0,0 +1,171 @@ +""" +Utility helpers for converting HDF5-backed metadata rows into WebDataset shards. +""" + +from __future__ import annotations + +import io +import logging +import os +from typing import Dict, Tuple + +import webdataset as wds +from PIL import Image + + +def convert_webp_to_jpeg(webp_bytes: bytes, resize_size: int = 0) -> bytes: + """ + Decode WebP bytes, optionally resize to a square, and re-encode as JPEG. + """ + image = Image.open(io.BytesIO(webp_bytes)).convert("RGB") + if resize_size and resize_size > 0: + resample = getattr(Image, "Resampling", Image).LANCZOS # Pillow>=9 compat + image = image.resize((resize_size, resize_size), resample) + + buffer = io.BytesIO() + image.save(buffer, format="JPEG", quality=95) + return buffer.getvalue() + + +def determine_most_specific_known_rank(taxon_dict: Dict[str, str | None]) -> str | None: + """ + Return the most specific rank (closest to species) whose value is not unknown. + """ + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + for rank in reversed(taxonomic_ranks): + value = taxon_dict.get(rank) + if value and value.strip().lower() != "unknown": + return rank + return None + + +def _clean_species_name(species_value: str | None) -> str | None: + if not species_value: + return None + parts = species_value.split() + if len(parts) > 1: + return parts[-1] + return species_value + + +def create_taxon_tag_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Compose a descriptive sentence enumerating the known taxonomy ranks. + """ + if not most_specific_rank: + return "a photo of unknown taxonomy." + + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + text_parts = [] + for rank in taxonomic_ranks: + value = taxon_dict.get(rank) + if not value or value.lower() == "unknown": + if rank == most_specific_rank: + break + continue + + if rank == "species": + cleaned = _clean_species_name(value) + if cleaned: + text_parts.append(f"{rank} {cleaned}") + else: + text_parts.append(f"{rank} {value}") + + if rank == most_specific_rank: + break + + return "a photo of " + " ".join(text_parts) + "." if text_parts else "a photo of unknown taxonomy." + + +def create_taxonomic_name_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Build a space-delimited taxonomy string up to the most specific known rank. + """ + if not most_specific_rank: + return "" + + taxonomic_ranks = ["kingdom", "phylum", "class", "order", "family", "genus", "species"] + parts = [] + for rank in taxonomic_ranks: + value = taxon_dict.get(rank) + if not value or value.lower() == "unknown": + if rank == most_specific_rank: + break + continue + + if rank == "species": + cleaned = _clean_species_name(value) + if cleaned: + parts.append(cleaned) + else: + parts.append(value) + + if rank == most_specific_rank: + break + + return " ".join(parts) + + +def create_scientific_name_text(taxon_dict: Dict[str, str | None], most_specific_rank: str | None) -> str: + """ + Build the scientific name from the most specific known rank. + """ + if not most_specific_rank: + return "Unknown" + + value = taxon_dict.get(most_specific_rank) + if not value or value.lower() == "unknown": + return "Unknown" + return value.rstrip(".") + + +def _get_common_name(taxon_dict: Dict[str, str | None]) -> str: + value = taxon_dict.get("common_name") + if value and value.strip(): + return value.strip().rstrip(".") + scientific = taxon_dict.get("scientific_name") + if scientific and scientific.strip(): + return scientific.strip().rstrip(".") + return "Unknown" + + +def generate_text_files(taxon_dict: Dict[str, str | None]) -> Dict[str, str]: + """ + Generate the five auxiliary text files used by downstream training pipelines. + """ + most_specific_rank = determine_most_specific_known_rank(taxon_dict) + taxonomic_name = create_taxonomic_name_text(taxon_dict, most_specific_rank) + scientific_name = create_scientific_name_text(taxon_dict, most_specific_rank) + taxon_for_common = dict(taxon_dict) + taxon_for_common["scientific_name"] = scientific_name + common_name = _get_common_name(taxon_for_common) + taxon_tag_text = create_taxon_tag_text(taxon_dict, most_specific_rank) + + files = { + "taxonomic_name.txt": taxonomic_name, + "scientific_name.txt": scientific_name, + "sci.txt": f"a photo of {scientific_name}.", + "taxon.txt": f"a photo of {taxonomic_name}.", + "taxonTag.txt": taxon_tag_text, + "common_name.txt": common_name, + "com.txt": f"a photo of {common_name}.", + "sci_com.txt": f"a photo of {scientific_name} with common name {common_name}.", + "taxon_com.txt": f"a photo of {taxonomic_name} with common name {common_name}.", + "taxonTag_com.txt": f"{taxon_tag_text.rstrip('.')} with common name {common_name}.", + } + + return files + + +def init_shard_writer(output_dir: str, shard_id: int, prefix: str = "shard") -> Tuple[wds.TarWriter, str]: + """ + Create (or truncate) the shard tarball and return the TarWriter plus path. + """ + os.makedirs(output_dir, exist_ok=True) + shard_name = f"{prefix}-{shard_id:05d}.tar" + shard_path = os.path.join(output_dir, shard_name) + if os.path.exists(shard_path): + os.remove(shard_path) + + logging.info("Creating shard %s at %s", shard_id, shard_path) + return wds.TarWriter(shard_path), shard_path