source/scripts/embed_external_proteins.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json

from bgc_retrieval.esm_embedding import generate_esm2_embeddings


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fasta", default="data/external/processed/external_proteins.fasta")
    parser.add_argument("--output", default="data/external/processed/external_esm2.h5")
    parser.add_argument("--metadata", default="data/external/processed/external_esm2_metadata.json")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--allow-cpu", action="store_true")
    args = parser.parse_args()
    metadata = generate_esm2_embeddings(
        args.fasta, args.output, args.metadata, args.batch_size, args.resume, args.allow_cpu
    )
    print(json.dumps(metadata, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()