| """Extract a bounded, reproducible sample from one published bucket object.""" |
| import argparse |
| from datetime import datetime, timezone |
| import hashlib |
| import json |
| from pathlib import Path |
| import tempfile |
|
|
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| BUCKET = "HuggingFaceBio/genbank-annotations" |
| DEFAULT_SOURCE = "annotations/fungi/s0000/s0000.chunk-00510.parquet" |
| ROOT = Path(__file__).resolve().parent |
|
|
|
|
| def extract(source, output, source_path, max_rows=24, max_bp=2_000_000, xet_hash=None): |
| if max_rows < 1 or max_bp < 1: |
| raise ValueError("Sample limits must be positive.") |
| pf = pq.ParquetFile(source) |
| selected, bases = [], 0 |
| |
| lengths = pf.read(columns=["segment_bp_length"]).column(0).to_pylist() |
| for i, length in enumerate(lengths): |
| if 0 < length <= max_bp - bases: |
| selected.append(i) |
| bases += length |
| if len(selected) >= max_rows: |
| break |
| if not selected: |
| raise ValueError("No complete segments fit the sample budget.") |
| pieces, offset = [], 0 |
| for group in range(pf.num_row_groups): |
| count = pf.metadata.row_group(group).num_rows |
| indices = [i - offset for i in selected if offset <= i < offset + count] |
| if indices: |
| pieces.append(pf.read_row_group(group).take(pa.array(indices))) |
| offset += count |
| sample = pa.concat_tables(pieces) |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=True) |
| target = output / "sample.parquet" |
| |
| pq.write_table(sample, target, compression="zstd", row_group_size=1) |
| metadata = sample.select([c for c in sample.column_names if not c.startswith("pred_prob_") and c != "sequence"]).to_pylist() |
| manifest = { |
| "bucket_id": BUCKET, "source_path": source_path, "source_xet_hash": xet_hash, |
| "created_at": datetime.now(timezone.utc).isoformat(), |
| "selection": "First complete segments fitting the row and base budgets, in source order", |
| "max_rows": max_rows, "max_bp": max_bp, "rows": sample.num_rows, |
| "bases": bases, "source_rows": pf.metadata.num_rows, "source_row_indices": selected, |
| "assemblies": sorted({r["assembly_accession"] for r in metadata}), |
| "sha256": hashlib.sha256(target.read_bytes()).hexdigest(), |
| } |
| (output / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") |
| print(json.dumps(manifest, indent=2)) |
| return manifest |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--source", default=DEFAULT_SOURCE) |
| parser.add_argument("--local-file", type=Path) |
| parser.add_argument("--output", type=Path, default=ROOT / "data") |
| parser.add_argument("--max-rows", type=int, default=24) |
| parser.add_argument("--max-bp", type=int, default=2_000_000) |
| args = parser.parse_args() |
| if args.local_file: |
| extract(args.local_file, args.output, args.source, args.max_rows, args.max_bp) |
| else: |
| from huggingface_hub import HfApi |
| api = HfApi() |
| entries = api.list_bucket_tree(BUCKET, args.source, recursive=False) |
| entry = next((e for e in entries if e.path == args.source), None) |
| if entry is None: |
| raise ValueError(f"Bucket object not found: {args.source}") |
| if entry.size > 100_000_000: |
| raise ValueError("Source exceeds the 100 MB download budget. Choose a smaller object or use --local-file.") |
| with tempfile.TemporaryDirectory() as temp: |
| local = Path(temp) / "source.parquet" |
| api.download_bucket_files(BUCKET, [(entry, local)], raise_on_missing_files=True) |
| extract(local, args.output, args.source, args.max_rows, args.max_bp, entry.xet_hash) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
