File size: 945 Bytes
c87881a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | #!/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()
|