affinity/extract_onnx_embeddings.py

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from tqdm.auto import tqdm

from affinity.data import load_dataset
from affinity.features import save_embedding_table
from affinity.onnx_embeddings import (
    create_onnx_embedder,
    detect_onnx_providers,
)


def extract_proteins(args, proteins: list[str], providers: list[str]) -> None:
    destination = Path(args.protein_output)
    if destination.exists() and not args.overwrite:
        print(f"Skipping existing protein cache: {destination}")
        return
    embedder = create_onnx_embedder(
        args.protein_encoder,
        args.protein_onnx,
        model_id=args.protein_model_id,
        max_length=args.protein_max_length,
        providers=providers,
    )
    print(
        f"Protein encoder ({args.protein_encoder}) providers: "
        f"{embedder.session.get_providers()}",
        flush=True,
    )
    blocks = []
    progress = tqdm(
        total=len(proteins),
        desc="Protein embeddings",
        unit="protein",
        dynamic_ncols=True,
    )
    for start in range(0, len(proteins), args.protein_batch_size):
        batch = proteins[start : start + args.protein_batch_size]
        blocks.append(embedder.encode(batch))
        progress.update(len(batch))
    progress.close()
    save_embedding_table(
        destination,
        proteins,
        np.concatenate(blocks, axis=0),
        args.protein_model_id,
        settings={
            "runtime": "onnxruntime",
            "encoder_type": args.protein_encoder,
            "pooling": (
                "mean_amino_acid_tokens"
                if args.protein_encoder == "esm2"
                else "attention_masked_mean_last_hidden_state"
            ),
            "max_length": args.protein_max_length,
            **(
                {"prompt": "[Determine superfamily] Seq=<{value}>"}
                if args.protein_encoder == "prollama"
                else {}
            ),
        },
    )


def extract_molecules(args, molecules: list[str], providers: list[str]) -> None:
    output = Path(args.molecule_output)
    output.mkdir(parents=True, exist_ok=True)
    embedder = create_onnx_embedder(
        args.molecule_encoder,
        args.molecule_onnx,
        model_id=args.molecule_model_id,
        max_length=args.molecule_max_length,
        providers=providers,
    )
    print(
        f"Molecule encoder ({args.molecule_encoder}) providers: "
        f"{embedder.session.get_providers()}",
        flush=True,
    )
    total_shards = (len(molecules) + args.molecule_shard_size - 1) // args.molecule_shard_size
    completed = 0
    if not args.overwrite:
        for shard_index in range(total_shards):
            start = shard_index * args.molecule_shard_size
            stop = min(start + args.molecule_shard_size, len(molecules))
            destination = output / f"molecules-{shard_index:05d}.npz"
            if destination.exists():
                completed += stop - start
    progress = tqdm(
        total=len(molecules),
        initial=completed,
        desc="Molecule embeddings",
        unit="molecule",
        dynamic_ncols=True,
    )
    for shard_index in range(total_shards):
        start = shard_index * args.molecule_shard_size
        stop = min(start + args.molecule_shard_size, len(molecules))
        destination = output / f"molecules-{shard_index:05d}.npz"
        if destination.exists() and not args.overwrite:
            continue
        values = molecules[start:stop]
        progress.set_postfix_str(f"shard {shard_index + 1}/{total_shards}")
        embeddings = []
        for batch_start in range(0, len(values), args.molecule_batch_size):
            batch = values[batch_start : batch_start + args.molecule_batch_size]
            embeddings.extend(embedder.encode(batch))
            progress.update(len(batch))
        save_embedding_table(
            destination,
            values,
            np.asarray(embeddings, dtype=np.float32),
            args.molecule_model_id,
            settings={
                "runtime": "onnxruntime",
                "encoder_type": args.molecule_encoder,
                "pooling": (
                    "pooler_output"
                    if args.molecule_encoder == "molformer"
                    else "mean_qformer_query_tokens"
                ),
                "max_length": args.molecule_max_length,
                **(
                    {"conformer_seed": 42}
                    if args.molecule_encoder == "mol_llama"
                    else {}
                ),
            },
        )
    progress.close()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Build resumable training caches with the deployed ONNX encoders"
    )
    parser.add_argument("--data", required=True)
    parser.add_argument("--protein-onnx")
    parser.add_argument("--molecule-onnx")
    parser.add_argument("--prollama-onnx")
    parser.add_argument("--mol-llama-onnx")
    parser.add_argument("--protein-output", required=True)
    parser.add_argument("--molecule-output", required=True)
    parser.add_argument(
        "--protein-encoder",
        choices=["esm2", "prollama"],
        default="esm2",
    )
    parser.add_argument(
        "--molecule-encoder",
        choices=["molformer", "mol_llama"],
        default="molformer",
    )
    parser.add_argument(
        "--protein-model-id",
        default="facebook/esm2_t12_35M_UR50D",
    )
    parser.add_argument(
        "--molecule-model-id",
        default="ibm-research/MoLFormer-XL-both-10pct",
    )
    parser.add_argument("--protein-batch-size", type=int, default=16)
    parser.add_argument("--molecule-batch-size", type=int, default=32)
    parser.add_argument("--protein-max-length", type=int, default=1024)
    parser.add_argument("--molecule-max-length", type=int, default=202)
    parser.add_argument("--molecule-shard-size", type=int, default=1000)
    parser.add_argument("--proteins-only", action="store_true")
    parser.add_argument("--molecules-only", action="store_true")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda"],
        default="auto",
        help="ONNX Runtime device selection",
    )
    args = parser.parse_args()
    if args.prollama_onnx:
        args.protein_onnx = args.protein_onnx or args.prollama_onnx
        args.protein_encoder = "prollama"
        if args.protein_model_id == "facebook/esm2_t12_35M_UR50D":
            args.protein_model_id = "GreatCaptainNemo/ProLLaMA"
    if args.mol_llama_onnx:
        args.molecule_onnx = args.molecule_onnx or args.mol_llama_onnx
        args.molecule_encoder = "mol_llama"
        if args.molecule_model_id == "ibm-research/MoLFormer-XL-both-10pct":
            args.molecule_model_id = "DongkiKim/Mol-Llama-3.1-8B-Instruct"
    if not args.protein_onnx or not args.molecule_onnx:
        parser.error("--protein-onnx and --molecule-onnx are required")
    if args.proteins_only and args.molecules_only:
        parser.error("--proteins-only and --molecules-only cannot be combined")

    print("Loading dataset...", flush=True)
    frame = load_dataset(args.data)
    proteins = frame["protein_sequence"].drop_duplicates().tolist()
    molecules = frame["compound_smiles"].drop_duplicates().tolist()
    print(
        f"Unique proteins: {len(proteins):,} | "
        f"Unique molecules: {len(molecules):,}",
        flush=True,
    )
    providers = detect_onnx_providers(args.device, verbose=True)
    if not args.molecules_only:
        print(f"Loading {args.protein_encoder} protein session...", flush=True)
        extract_proteins(args, proteins, providers)
    if not args.proteins_only:
        print(f"Loading {args.molecule_encoder} molecule session...", flush=True)
        extract_molecules(args, molecules, providers)


if __name__ == "__main__":
    main()