pepe-cli 1.0.0__py3-none-any.whl

pepe/embedders/esm_embedder.py

@@ -0,0 +1,164 @@
+import logging
+import torch
+from esm import pretrained
+from pepe.embedders.base_embedder import BaseEmbedder
+import pepe.utils
+
+logger = logging.getLogger("src.embedders.esm_embedder")
+
+
+class ESMEmbedder(BaseEmbedder):
+    def __init__(self, args):
+        super().__init__(args)
+        self.sequences = pepe.utils.fasta_to_dict(args.fasta_path)
+        self.num_sequences = len(self.sequences)
+        (
+            self.model,
+            self.alphabet,
+            self.num_heads,
+            self.num_layers,
+            self.embedding_size,
+            self.prepend_bos,
+            self.append_eos,
+        ) = self._initialize_model(self.model_name)
+        self.valid_tokens = set(self.alphabet.all_toks)
+        pepe.utils.check_input_tokens(
+            self.valid_tokens, self.sequences, self.model_name
+        )
+        self.special_tokens = self.get_special_tokens()
+        self.layers = self._load_layers(self.layers)
+        self.data_loader, self.max_length = self._load_data(
+            self.sequences, self.substring_dict
+        )  # tokenize and batch sequences and update max_length
+        self._set_output_objects()
+
+    def _initialize_model(self, model_name):
+        """Initialize the model, tokenizer"""
+        #  Loading the pretrained model and alphabet for tokenization
+        logger.info("Loading model...")
+        # model, alphabet = pretrained.load_model_and_alphabet(model_name)
+        model, alphabet = pretrained.load_model_and_alphabet_hub(model_name)
+        model.eval()  # Setting the model to evaluation mode
+        if not self.disable_special_tokens:
+            model.append_eos = True if not model_name.startswith("esm1") else False  # type: ignore
+            model.prepend_bos = True  # type: ignore
+        else:
+            model.append_eos = False  # type: ignore
+            model.prepend_bos = False  # type: ignore
+
+        num_heads = model.layers[0].self_attn.num_heads  # type: ignore
+        num_layers = len(model.layers)  # type: ignore
+        embedding_size = (
+            model.embed_tokens.embedding_dim  # type: ignore
+            if model_name.startswith("esm1")
+            else model.embed_dim
+        )
+
+        # Moving the model to GPU if available for faster processing
+        if torch.cuda.is_available() and self.device == "cuda":
+            model = model.cuda()
+            logger.info("Transferred model to GPU")
+        else:
+            logger.info("No GPU available, using CPU")
+        return (
+            model,
+            alphabet,
+            num_heads,
+            num_layers,
+            embedding_size,
+            model.prepend_bos,
+            model.append_eos,
+        )
+
+    def get_special_tokens(self):
+        special_tokens = self.alphabet.all_special_tokens
+        special_token_ids = torch.tensor(
+            [self.alphabet.tok_to_idx[tok] for tok in special_tokens],
+            device=self.device,
+            dtype=torch.int8,
+        )
+        return special_token_ids
+
+    def _load_layers(self, layers):
+        if not layers:
+            layers = list(range(1, self.model.num_layers + 1))  # type: ignore
+            return layers
+        # Checking if the specified representation layers are valid
+        assert all(
+            -(self.model.num_layers + 1) <= i <= self.model.num_layers for i in layers  # type: ignore
+        )
+        layers = [
+            (i + self.model.num_layers + 1) % (self.model.num_layers + 1)  # type: ignore
+            for i in layers
+        ]
+        return layers
+
+    def _load_data(self, sequences, substring_dict=None):
+        # Creating a dataset from the input fasta file
+        dataset = pepe.utils.ESMDataset(
+            sequences,
+            substring_dict,
+            self.context,
+            self.alphabet,
+            self.max_length,
+            self.prepend_bos,  # type: ignore
+            self.append_eos,  # type: ignore
+        )
+        # Generating batch indices based on token count
+        logger.info("Generating batches...")
+        batches = pepe.utils.TokenBudgetBatchSampler(dataset, self.batch_size)
+        # DataLoader to iterate through batches efficiently
+        data_loader = torch.utils.data.DataLoader(
+            dataset, batch_sampler=batches, collate_fn=dataset.safe_collate
+        )
+        logger.info("Data loaded")
+        # Getting the maximum sequence length from the dataset
+        max_length = dataset.get_max_encoded_length()
+        return data_loader, max_length
+
+    def _compute_outputs(
+        self,
+        model,
+        toks,
+        attention_mask,
+        return_embeddings,
+        return_contacts,
+        return_logits,
+    ):
+        outputs = model(
+            toks,
+            repr_layers=self.layers,
+            return_contacts=return_contacts,
+        )
+        if return_logits:
+            logits = (
+                outputs["logits"]
+                .to(dtype=self._precision_to_dtype(self.precision, "torch"))
+                .permute(2, 0, 1)
+                .cpu()
+            )  # permute to match the shape of the representations
+            torch.cuda.empty_cache()
+        else:
+            logits = None
+
+        if return_contacts:
+            attention_matrices = (
+                outputs["attentions"]
+                .to(dtype=self._precision_to_dtype(self.precision, "torch"))
+                .permute(1, 0, 2, 3, 4)
+            ).cpu()  # permute to match the shape of the representations
+            torch.cuda.empty_cache()
+        else:
+            attention_matrices = None
+        # Extracting layer representations and moving them to CPU
+        if return_embeddings:
+            representations = {
+                layer: t.to(
+                    dtype=self._precision_to_dtype(self.precision, "torch")
+                ).cpu()
+                for layer, t in outputs["representations"].items()
+            }
+            torch.cuda.empty_cache()
+        else:
+            representations = None
+        return logits, representations, attention_matrices

pepe/embedders/huggingface_embedder.py

@@ -0,0 +1,200 @@
+import logging
+import os
+import torch
+import pepe.utils
+from pepe.embedders.base_embedder import BaseEmbedder
+from transformers import T5EncoderModel, T5Tokenizer
+from transformers import RoFormerTokenizer, RoFormerModel
+from transformers.models.roformer.modeling_roformer import (
+    RoFormerSinusoidalPositionalEmbedding,
+)
+
+
+# Set max_split_size_mb
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
+
+logger = logging.getLogger("src.embedders.huggingface_embedder")
+
+
+class HuggingfaceEmbedder(BaseEmbedder):
+    def __init__(self, args):
+        super().__init__(args)
+        if self.return_logits:
+            logger.warning(
+                "Warning: Logits are not supported for this model. Setting to False."
+            )
+            self.return_logits = False
+            self.output_types.remove("logits")
+
+    def _load_layers(self, layers):
+        """Check if the specified representation layers are valid."""
+        if not layers:
+            layers = list(range(1, self.model.config.num_hidden_layers + 1))  # type: ignore
+            return layers
+        assert all(
+            -(self.model.config.num_hidden_layers + 1)  # type: ignore
+            <= i
+            <= self.model.config.num_hidden_layers  # type: ignore
+            for i in layers
+        )
+        layers = [
+            (i + self.model.config.num_hidden_layers + 1)  # type: ignore
+            % (self.model.config.num_hidden_layers + 1)  # type: ignore
+            for i in layers
+        ]
+        return layers
+
+    def _load_data(self, sequences, substring_dict):
+        """Tokenize sequences and create a DataLoader."""
+        # Tokenize sequences
+        dataset = pepe.utils.HuggingFaceDataset(
+            sequences,
+            substring_dict,
+            self.context,
+            self.tokenizer,  # type: ignore
+            self.max_length,
+            add_special_tokens=not self.disable_special_tokens,
+        )
+        logger.info("Batching sequences...")
+        batch_sampler = pepe.utils.TokenBudgetBatchSampler(
+            dataset=dataset, token_budget=self.batch_size
+        )
+        data_loader = torch.utils.data.DataLoader(
+            dataset, batch_sampler=batch_sampler, collate_fn=dataset.safe_collate
+        )
+        max_length = dataset.get_max_encoded_length()
+        logger.info("Finished tokenizing and batching sequences")
+
+        return data_loader, max_length
+
+    def _compute_outputs(
+        self,
+        model,
+        toks,
+        attention_mask,
+        return_embeddings,
+        return_contacts,
+        return_logits=False,
+    ):
+        outputs = model(
+            input_ids=toks,
+            attention_mask=attention_mask,
+            output_hidden_states=return_embeddings,
+            output_attentions=return_contacts,
+        )
+        if return_contacts:
+            attention_matrices = (
+                torch.stack(outputs.attentions)  # type: ignore
+                .to(self._precision_to_dtype(self.precision, "torch"))  # type: ignore
+                .cpu()
+            )  # stack attention matrices across layers
+            torch.cuda.empty_cache()
+        else:
+            attention_matrices = None
+        if return_embeddings:
+            representations = {
+                layer: outputs.hidden_states[layer]
+                .to(
+                    self._precision_to_dtype(self.precision, "torch"),
+                )
+                .cpu()
+                for layer in self.layers  # type: ignore
+            }
+            torch.cuda.empty_cache()
+        else:
+            representations = None
+        logits = None  # Model doesn't return logits
+        return logits, representations, attention_matrices
+
+
+class Antiberta2Embedder(HuggingfaceEmbedder):
+    def __init__(self, args):
+        super().__init__(args)
+        self.sequences = pepe.utils.fasta_to_dict(args.fasta_path)
+        self.num_sequences = len(self.sequences)
+        (
+            self.model,
+            self.tokenizer,
+            self.num_heads,
+            self.num_layers,
+            self.embedding_size,
+        ) = self._initialize_model(self.model_link)
+        self.valid_tokens = set(self.tokenizer.get_vocab().keys())
+        pepe.utils.check_input_tokens(
+            self.valid_tokens, self.sequences, self.model_name
+        )
+        self.special_tokens = torch.tensor(
+            self.tokenizer.all_special_ids, device=self.device, dtype=torch.int8
+        )
+        self.layers = self._load_layers(self.layers)
+        self.data_loader, self.max_length = self._load_data(
+            self.sequences, self.substring_dict
+        )
+        self._set_output_objects()
+        assert self.max_length <= 256, "AntiBERTa2 only supports max_length <= 256"
+
+    def _initialize_model(self, model_link="alchemab/antiberta2-cssp"):
+        """Initialize the model, tokenizer, and device."""
+        if torch.cuda.is_available() and self.device == "cuda":
+            device = torch.device("cuda")
+            logger.info("Transferred model to GPU")
+        else:
+            device = torch.device("cpu")
+            logger.info("No GPU available, using CPU")
+        tokenizer = RoFormerTokenizer.from_pretrained(model_link, use_fast=True)
+        model = RoFormerModel.from_pretrained(model_link).to(device)  # type: ignore
+        model.eval()
+        num_heads = model.config.num_attention_heads
+        num_layers = model.config.num_hidden_layers
+        embedding_size = model.config.hidden_size
+        return model, tokenizer, num_heads, num_layers, embedding_size
+
+
+class T5Embedder(HuggingfaceEmbedder):
+    def __init__(self, args):
+        super().__init__(args)
+        self.sequences = self.fasta_to_dict(args.fasta_path)  # type: ignore
+        self.num_sequences = len(self.sequences)
+        (
+            self.model,
+            self.tokenizer,
+            self.num_heads,
+            self.num_layers,
+            self.embedding_size,
+        ) = self._initialize_model(self.model_link)
+        self.valid_tokens = self.get_valid_tokens()
+        pepe.utils.check_input_tokens(
+            self.valid_tokens, self.sequences, self.model_name
+        )
+        self.special_tokens = torch.tensor(
+            self.tokenizer.all_special_ids, device=self.device, dtype=torch.int8
+        )
+        self.layers = self._load_layers(self.layers)
+        self.data_loader, self.max_length = self._load_data(
+            self.sequences, self.substring_dict
+        )
+        self._set_output_objects()
+
+    def get_valid_tokens(self):
+        valid_tokens = set(
+            k[1:] if k.startswith("▁") else k
+            for k in set(self.tokenizer.get_vocab().keys())
+        )
+        return valid_tokens
+
+    def _initialize_model(self, model_link="Rostlab/prot_t5_xl_half_uniref50-enc"):
+        """Initialize the model, tokenizer, and device."""
+
+        if torch.cuda.is_available() and self.device == "cuda":
+            device = torch.device("cuda")
+            logger.info("Transferred model to GPU")
+        else:
+            device = torch.device("cpu")
+            logger.info("No GPU available, using CPU")
+        tokenizer = T5Tokenizer.from_pretrained(model_link, use_fast=True)
+        model = T5EncoderModel.from_pretrained(model_link).to(device)  # type: ignore
+        model.eval()
+        num_heads = model.config.num_heads
+        num_layers = model.config.num_layers
+        embedding_size = model.config.hidden_size
+        return model, tokenizer, num_heads, num_layers, embedding_size

pepe/model_selecter.py

@@ -0,0 +1,99 @@
+from pepe.embedders.esm_embedder import ESMEmbedder
+from pepe.embedders.huggingface_embedder import T5Embedder, Antiberta2Embedder
+from pepe.embedders.custom_embedder import CustomEmbedder
+
+
+import os
+
+
+def select_model(model_name):
+    if "esm2" in model_name.lower():
+        return ESMEmbedder
+    elif "esm1" in model_name.lower():
+        return ESMEmbedder
+    # elif "antiberta2" in model_name.lower() and model_name.startswith("alchemab"):
+    #    return Antiberta2Embedder
+    # elif "t5" in model_name.lower() and model_name.startswith("Rostlab"):
+    #    return T5Embedder
+    elif (
+        model_name.endswith(".pt")
+        or model_name.endswith(".pth")
+        or model_name.startswith("custom:")
+        or (
+            os.path.exists(model_name)
+            and (os.path.isfile(model_name) or os.path.isdir(model_name))
+        )
+    ):
+        return CustomEmbedder
+    elif "/" in model_name:
+        # Assume it's a Hugging Face model (username/model-name format)
+        # Try to determine the architecture automatically
+        try:
+            from transformers import AutoConfig
+
+            config = AutoConfig.from_pretrained(model_name)
+            model_type = config.model_type.lower()
+
+            if model_type in ["t5", "mt5"]:
+                return T5Embedder
+            elif model_type in ["roformer"]:
+                return Antiberta2Embedder
+            elif model_type in ["bert"]:
+                # For BERT-like models, we could potentially use a generic embedder
+                # but for now, suggest using CustomEmbedder or creating a specific one
+                raise ValueError(
+                    f"BERT-like models are not yet directly supported. Consider using a PyTorch version of the model with CustomEmbedder."
+                )
+            else:
+                # For other architectures, you might want to add more specific embedders
+                # or use a generic HuggingfaceEmbedder
+                raise ValueError(
+                    f"Model architecture '{model_type}' not yet supported for custom Hugging Face models"
+                )
+        except Exception as e:
+            # Check if it's a Keras/TensorFlow model
+            error_msg = str(e)
+            if "Unrecognized model" in error_msg or "model_type" in error_msg:
+                raise ValueError(
+                    f"Model {model_name} appears to be a Keras/TensorFlow model or has an unsupported architecture. EmbedAIRR currently supports PyTorch models only. Consider using a PyTorch version or converting the model."
+                )
+            else:
+                raise ValueError(
+                    f"Could not determine model architecture for {model_name}: {e}"
+                )
+    else:
+        raise ValueError(f"Model {model_name} not supported")
+
+
+supported_models = [
+    # ESM models
+    "esm1_t34_670M_UR50S",
+    "esm1_t34_670M_UR50D",
+    "esm1_t34_670M_UR100",
+    "esm1_t12_85M_UR50S",
+    "esm1_t6_43M_UR50S",
+    "esm1b_t33_650M_UR50S",
+    #'esm_msa1_t12_100M_UR50S',
+    #'esm_msa1b_t12_100M_UR50S',
+    "esm1v_t33_650M_UR90S_1",
+    "esm1v_t33_650M_UR90S_2",
+    "esm1v_t33_650M_UR90S_3",
+    "esm1v_t33_650M_UR90S_4",
+    "esm1v_t33_650M_UR90S_5",
+    #'esm_if1_gvp4_t16_142M_UR50',
+    "esm2_t6_8M_UR50D",
+    "esm2_t12_35M_UR50D",
+    "esm2_t30_150M_UR50D",
+    "esm2_t33_650M_UR50D",
+    "esm2_t36_3B_UR50D",
+    "esm2_t48_15B_UR50D",
+    # Pre-defined Hugging Face models
+    "Rostlab/prot_t5_xl_half_uniref50-enc",
+    "Rostlab/ProstT5",
+    "alchemab/antiberta2-cssp",
+    "alchemab/antiberta2",
+    # Custom models examples:
+    # - PyTorch models: "/path/to/model.pt", "/path/to/model_directory/", "custom:/path/to/model.pt"
+    # - Hugging Face models: "username/model-name", "./local_hf_model"
+    # - See documentation for details on custom model requirements
+]

pepe/parse_arguments.py

@@ -0,0 +1,167 @@
+import argparse
+from pepe.model_selecter import supported_models
+
+
+def str2bool(value):
+    if isinstance(value, bool):
+        return value
+    if value.lower() in ("true", "t", "yes", "y", "1"):
+        return True
+    elif value.lower() in ("false", "f", "no", "n", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError("Boolean value expected.")
+
+
+def str2ints(value):
+    """Convert a string to a list of integers."""
+    if isinstance(value, str):
+        if value.lower() == "all":
+            return None
+        elif value.lower() == "last":
+            return [-1]
+        else:
+            try:
+                return [int(x) for x in value.split(" ")]
+            except ValueError:
+                raise argparse.ArgumentTypeError(
+                    "Invalid input. Expected integer(s) or spaced list of integers or 'all' or 'last'."
+                )
+    elif isinstance(value, int):
+        return [value]
+
+
+# Parsing command-line arguments for input and output file paths
+def parse_arguments():
+    """Parse command-line arguments for input and output file paths."""
+    parser = argparse.ArgumentParser(description="Input path")
+    parser.add_argument(
+        "--experiment_name",
+        type=str,
+        default=None,
+        help="Name of the experiment. Will be used to name the output files. If not provided, the output files will be named after the input file.",
+    )
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        required=True,
+        help="Model name or path to custom model. For custom models, use path to .pt/.pth file or directory containing model files. For predefined models, use model name from supported list.",
+    )
+    parser.add_argument(
+        "--tokenizer_from",
+        type=str,
+        default=None,
+        help="Huggingface address of the tokenizer to use. If not provided, will use the tokenizer from the model. If using a custom model, provide the path to the tokenizer directory.",
+    )
+    parser.add_argument(
+        "--fasta_path",
+        type=str,
+        required=True,
+        help="Path to the input FASTA file. Required. If no experiment name is provided, the output files will be named after the input file.",
+    )
+    parser.add_argument(
+        "--output_path",
+        type=str,
+        required=True,
+        help="Directory for output files \n Will generate a subdirectory for outputs of each output_type.\n Will output multiple files if multiple layers are specified with '--layers'. Output file is a single tensor or a list of tensors when --pooling is False.",
+    )
+    parser.add_argument(
+        "--substring_path",
+        default=None,
+        type=str,
+        help=" Path to a CSV file with columns 'sequence_id' and 'substring'. Only required when selecting 'substring' embedding option.",
+    )
+    parser.add_argument(
+        "--context",
+        default=0,
+        type=int,
+        help="Only specify when including 'substring_pooling' in '--extract_embeddings' option. Number of amino acids to include before and after the substring sequence. Default is 0.",
+    )
+    parser.add_argument(
+        "--layers",
+        type=str2ints,
+        nargs="*",
+        default=[[-1]],
+        help="Representation layers to extract from the model. Default is the last layer. Example: argument '--layers -1 6' will output the last layer and the sixth layer.",
+    )
+    parser.add_argument(
+        "--extract_embeddings",
+        choices=[
+            "per_token",
+            "mean_pooled",
+            "substring_pooled",
+            "attention_head",
+            "attention_layer",
+            "attention_model",
+            "logits",
+        ],
+        default=["mean_pooled"],
+        nargs="+",
+        help="Set the embedding return types. Choose one or more from: 'per_token', 'mean_pooled', 'substring_pooled', 'attention_head', 'attention_layer', 'attention_model' and 'logits' (experimental). Default is 'pooled'.",
+    )
+    parser.add_argument(
+        "--batch_size",
+        type=int,
+        default=1024,
+        help="Number of tokens (not sequences!) per batch. Default is 1024.",
+    )
+    parser.add_argument(
+        "--discard_padding",
+        action="store_true",
+        help="Discard padding tokens from per_token embeddings output. Not compatible with --streaming_output Default is False.",
+    )
+    parser.add_argument(
+        "--max_length",
+        default="max_length",
+        help="Length to which sequences will be padded. Default is longest sequence.",
+    )
+    parser.add_argument(
+        "--streaming_output",
+        type=str2bool,
+        choices=[True, False],
+        default=True,
+        help="Preallocate output files and concurrently write embeddings to disk in batches while computing embeddings. Default is True.",
+    )
+    parser.add_argument(
+        "--num_workers",
+        type=int,
+        default=8,
+        help="Number of workers for asynchronous data writing. Only relevant when --batch_writing is enabled. Default is 8.",
+    )
+    parser.add_argument(
+        "--disable_special_tokens",
+        type=str2bool,
+        choices=[True, False],
+        default=False,
+        help="Disable special tokens in the model. Default is False.",
+    )
+    parser.add_argument(
+        "--flatten",
+        type=str2bool,
+        choices=[True, False],
+        default=False,
+        help="Flatten the output tensors. Default is False.",
+    )
+    parser.add_argument(
+        "--flush_batches_after",
+        type=int,
+        default=128,
+        help="Size (in MB) of outputs to accumulate in RAM per worker (--num_workers) before flushing to disk. Default is 128.",
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        default="32",
+        choices=["float16", "16", "half", "float32", "32", "full"],
+        help="Precision of the output data. Inference during embedding is not affected. Default is 'float32'.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str.lower,
+        default="cuda",
+        choices=["cuda", "cpu"],
+        help="Device to run the model on. Default is 'cuda'.",
+    )
+
+    args = parser.parse_args()
+    return args