extract_embedding.py

import os
import sys as _sys
import numpy as np
import torch
from pathlib import Path
from torch.nn.utils.rnn import pad_sequence
from load_model import load_pretrained_model
from utils.dataset import CodonMambaDataset
from tqdm import tqdm
import argparse
from utils.config.mamba_config import *


class DynamicBatchCollator:
    def __init__(self, pad_idx):
        self.pad_idx = pad_idx

    def __call__(self, batch):
        rna_ids, input_ids = zip(*batch)

        input_ids = pad_sequence(
            [t.contiguous() for t in input_ids],
            batch_first=True,
            padding_value=self.pad_idx
        )
        return rna_ids, input_ids


def extract_embedding(model_path, fasta_file, batch_size, num_workers, output_path):
    """
    Extract sequence embeddings, remove PAD positions, and save all embeddings
    into a single .npz file.

    Each sequence is stored as one array with key = rna_id and
    value = numpy array of shape [L_i, H].
    """
    # 1. Load pretrained model and vocabulary
    rna_vocab, model = load_pretrained_model(model_path)
    model.eval()

    # 2. Build dataset and DataLoader
    extract_data = CodonMambaDataset(fasta_file, rna_vocab)
    collate_fn = DynamicBatchCollator(rna_vocab.pad_idx)

    data_loader: torch.utils.data.DataLoader = torch.utils.data.DataLoader(
        extract_data,
        batch_size=batch_size,
        num_workers=num_workers,
        shuffle=False,
        collate_fn=collate_fn
    )

    # 3. Select device
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # 4. Store embeddings for all sequences
    embeddings = {}  # {rna_id: np.ndarray [L_i, H]}

    with torch.no_grad():
        for batch in tqdm(data_loader, desc="Extracting embeddings"):
            rna_ids, input_ids = batch
            input_ids = input_ids.to(device)

            results = model(input_ids)

            # Usually we do not need logits here, only hidden_states
            # logits = results.logits.to(torch.float)
            hidden_states = results.hidden_states.to(torch.float)  # [B, L, H]

            # Remove padding
            pad_idx = rna_vocab.pad_idx
            pad_mask = (input_ids != pad_idx)  # [B, L], True indicates non-PAD

            # Remove PAD and save per sequence
            for i, rna_id in enumerate(rna_ids):
                valid_pos = pad_mask[i]                   # [L]
                seq_hidden = hidden_states[i][valid_pos] # [L_valid, H]
                print(f"the shape of embedding is{rna_id}:{seq_hidden.shape}")
                # Convert to numpy for saving
                embeddings[rna_id] = seq_hidden.cpu().numpy()

    # 5. Save as a single .npz file
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Save multiple arrays into one npz: embeddings[rna_id] is stored under key rna_id
    np.savez_compressed(output_path, **embeddings)
    print(f"Saved embeddings for {len(embeddings)} sequences to: {output_path}")


if __name__ == "__main__":
    current_directory = Path(__file__).parent.absolute()
    parent_dir = current_directory.parent
    print("Project root:", parent_dir)

    parser = argparse.ArgumentParser(
        description="Extract CodonMamba embeddings and save as a NumPy .npz file"
    )
    parser.add_argument(
        "--model_path",
        type=str,
        default=os.path.join(current_directory, "checkpoints/codonmamba_weights.ckpt"),
        help="Path to the pretrained model checkpoint"
    )
    parser.add_argument(
        "--fasta_file",
        type=str,
        required=True,
        help="Path to the input FASTA file"
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=2,
        help="Batch size for the DataLoader"
    )
    parser.add_argument(
        "--num_workers",
        type=int,
        default=4,
        help="Number of worker processes for the DataLoader"
    )
    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="Output .npz file path, e.g. ./embeddings.npz"
    )

    args = parser.parse_args()

    extract_embedding(
        model_path=args.model_path,
        fasta_file=args.fasta_file,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        output_path=args.output
    )