pepe-cli 1.0.1__tar.gz → 1.0.3__tar.gz

{pepe_cli-1.0.1/src/pepe_cli.egg-info → pepe_cli-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.0.1
+Version: 1.0.3
 Summary: Pipeline for Easy Protein Embedding - Extract embeddings and attention matrices from protein sequences
 Home-page: https://github.com/csi-greifflab/pepe-cli
 Author: Jahn Zhong

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/__init__.py

@@ -2,7 +2,7 @@ import logging
 import sys
 
 # Package metadata - single source of truth
-__version__ = "1.0.1"
+__version__ = "1.0.3"
 __package_name__ = "pepe-cli"
 __module_name__ = "pepe"
 __author__ = "Jahn Zhong"

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/__main__.py

@@ -1,5 +1,4 @@
 import logging
-from pepe.model_selecter import select_model
 from pepe.parse_arguments import parse_arguments
 
 logger = logging.getLogger("src.__main__")
@@ -8,6 +7,9 @@ logger = logging.getLogger("src.__main__")
 def main():
     args = parse_arguments()
 
+    # Lazy import to avoid loading heavy dependencies during --help
+    from pepe.model_selecter import select_model
+
     selected_model = select_model(args.model_name)
 
     embedder = selected_model(args)

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/embedders/esm_embedder.py

@@ -1,9 +1,23 @@
 import logging
 import torch
-from esm import pretrained
 from pepe.embedders.base_embedder import BaseEmbedder
 import pepe.utils
 
+
+# Lazy imports to avoid loading heavy dependencies at import time
+def _import_esm():
+    """Lazy import of ESM components to avoid loading issues."""
+    try:
+        from esm import pretrained
+
+        return pretrained
+    except ImportError as e:
+        logger.error(f"Failed to import ESM: {e}")
+        raise ImportError(
+            "Failed to import ESM. Please ensure fair-esm is installed: pip install fair-esm"
+        ) from e
+
+
 logger = logging.getLogger("src.embedders.esm_embedder")
 
 
@@ -36,6 +50,10 @@ class ESMEmbedder(BaseEmbedder):
         """Initialize the model, tokenizer"""
         #  Loading the pretrained model and alphabet for tokenization
         logger.info("Loading model...")
+
+        # Lazy import ESM components
+        pretrained = _import_esm()
+
         # 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

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/embedders/huggingface_embedder.py

@@ -3,11 +3,30 @@ 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,
-)
+
+
+# Lazy imports to avoid loading heavy dependencies at import time
+def _import_transformers():
+    """Lazy import of transformers components to avoid loading issues."""
+    try:
+        from transformers import T5EncoderModel, T5Tokenizer
+        from transformers import RoFormerTokenizer, RoFormerModel
+        from transformers.models.roformer.modeling_roformer import (
+            RoFormerSinusoidalPositionalEmbedding,
+        )
+
+        return (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+        )
+    except ImportError as e:
+        logger.error(f"Failed to import transformers: {e}")
+        raise ImportError(
+            "Failed to import transformers. Please ensure transformers is installed: pip install transformers"
+        ) from e
 
 
 # Set max_split_size_mb
@@ -141,6 +160,16 @@ class Antiberta2Embedder(HuggingfaceEmbedder):
         else:
             device = torch.device("cpu")
             logger.info("No GPU available, using CPU")
+
+        # Lazy import transformers components
+        (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+        ) = _import_transformers()
+
         tokenizer = RoFormerTokenizer.from_pretrained(model_link, use_fast=True)
         model = RoFormerModel.from_pretrained(model_link).to(device)  # type: ignore
         model.eval()
@@ -191,6 +220,16 @@ class T5Embedder(HuggingfaceEmbedder):
         else:
             device = torch.device("cpu")
             logger.info("No GPU available, using CPU")
+
+        # Lazy import transformers components
+        (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+        ) = _import_transformers()
+
         tokenizer = T5Tokenizer.from_pretrained(model_link, use_fast=True)
         model = T5EncoderModel.from_pretrained(model_link).to(device)  # type: ignore
         model.eval()

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/model_selecter.py

@@ -1,19 +1,32 @@
-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 _get_esm_embedder():
+    """Lazy import of ESM embedder to avoid loading heavy dependencies."""
+    from pepe.embedders.esm_embedder import ESMEmbedder
+
+    return ESMEmbedder
+
+
+def _get_huggingface_embedders():
+    """Lazy import of HuggingFace embedders to avoid loading heavy dependencies."""
+    from pepe.embedders.huggingface_embedder import T5Embedder, Antiberta2Embedder
+
+    return T5Embedder, Antiberta2Embedder
+
+
 def select_model(model_name):
     if "esm2" in model_name.lower():
-        return ESMEmbedder
+        return _get_esm_embedder()
     elif "esm1" in model_name.lower():
-        return ESMEmbedder
+        return _get_esm_embedder()
     # elif "antiberta2" in model_name.lower() and model_name.startswith("alchemab"):
+    #    T5Embedder, Antiberta2Embedder = _get_huggingface_embedders()
     #    return Antiberta2Embedder
     # elif "t5" in model_name.lower() and model_name.startswith("Rostlab"):
+    #    T5Embedder, Antiberta2Embedder = _get_huggingface_embedders()
     #    return T5Embedder
     elif (
         model_name.endswith(".pt")
@@ -35,8 +48,10 @@ def select_model(model_name):
             model_type = config.model_type.lower()
 
             if model_type in ["t5", "mt5"]:
+                T5Embedder, Antiberta2Embedder = _get_huggingface_embedders()
                 return T5Embedder
             elif model_type in ["roformer"]:
+                T5Embedder, Antiberta2Embedder = _get_huggingface_embedders()
                 return Antiberta2Embedder
             elif model_type in ["bert"]:
                 # For BERT-like models, we could potentially use a generic embedder

{pepe_cli-1.0.1 → pepe_cli-1.0.3}/src/pepe/parse_arguments.py

@@ -1,5 +1,4 @@
 import argparse
-from pepe.model_selecter import supported_models
 
 
 def str2bool(value):

{pepe_cli-1.0.1 → pepe_cli-1.0.3/src/pepe_cli.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.0.1
+Version: 1.0.3
 Summary: Pipeline for Easy Protein Embedding - Extract embeddings and attention matrices from protein sequences
 Home-page: https://github.com/csi-greifflab/pepe-cli
 Author: Jahn Zhong