pepe-cli 1.1__tar.gz → 1.3.0__tar.gz

{pepe_cli-1.1/src/pepe_cli.egg-info → pepe_cli-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.1
+Version: 1.3.0
 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
@@ -13,7 +13,7 @@ Project-URL: Documentation, https://github.com/csi-greifflab/pepe-cli#readme
 Keywords: protein,embeddings,bioinformatics,machine-learning,nlp,transformers
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Science/Research
-Classifier: License :: OSI Approved :: GNU Affero General Public License v3
+Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8
@@ -31,6 +31,7 @@ Requires-Dist: sentencepiece
 Requires-Dist: numpy
 Requires-Dist: protobuf
 Requires-Dist: alive_progress
+Requires-Dist: rjieba
 Dynamic: author
 Dynamic: home-page
 Dynamic: license-file
@@ -40,9 +41,16 @@ Dynamic: requires-python
 
 PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings and attention matrices from protein sequences using pre-trained models. This tool supports various configurations for extracting embeddings and attention matrices, including options for handling CDR3 sequences. Currently implemented models are ESM2 from the 2023 paper ["Evolutionary-scale prediction of atomic-level protein structure with a language model"](https://science.org/doi/10.1126/science.ade2574) and AntiBERTa2-CSSP from the 2023 conference paper ["Enhancing Antibody Language Models with Structural Information"](https://www.mlsb.io/papers_2023/Enhancing_Antibody_Language_Models_with_Structural_Information.pdf). PEPE also supports custom PLMs from local files or from Huggingface Hub addresses. 
 
+### Citation
+> **PEPE: Scalable extraction of multi-modal protein language model representations**
+> Jahn Zhong, Niccolò Cardente, Geir Kjetil Sandve, Habib Bashour, Maria Francesca Abbate, Victor Greiff
+> *bioRxiv* (2026) 
+> [DOI: 10.1101/2025.10.13.680902](https://doi.org/10.1101/2025.10.13.680902)
+
 ## Quick start
 
-1. Install PEPE \
+1. Install PEPE
+
     From PyPI:    
     ```sh
     pip install pepe-cli
@@ -57,7 +65,13 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
     cd pepe-cli
     pip install .
     ```
-2. Run the embedding script:\
+
+2. *(Optional)* For ESMC models (e.g. `biohub/ESMC-300M`), install Biohub's transformers fork:
+    ```sh
+    pip install git+https://github.com/Biohub/transformers.git@main
+    ```
+
+3. Extract embeddings:\
     Extract mean pooled embeddings from protein amino acid sequences in FASTA file:
     ```sh
     pepe --experiment_name <optional_string> --fasta_path <file_path> --output_path <directory> --model_name <model_name>
@@ -67,7 +81,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
 
 PEPE can also be used as a Python library. This allows for programmatic access to protein embeddings without using the command-line interface.
 
-1. Install PEPE (see [Quick start with CLI](#quick-start-with-cli) for details).
+1. Install PEPE (see above for details).
 2. Use the `pepe.embed()` function in your script:
 
 ```python
@@ -220,6 +234,10 @@ results = pepe.embed(
     - RoFormer models
         - alchemab/antiberta2-cssp
         - alchemab/antiberta2
+    - ESMC models (requires Biohub transformers fork; see install instructions above)
+        - biohub/ESMC-300M
+        - biohub/ESMC-600M
+        - biohub/ESMC-6B
     - Custom Hugging Face models
         - Any compatible model from Hugging Face Hub: `username/model-name`
         - Private models with authentication
@@ -234,6 +252,7 @@ results = pepe.embed(
 ### Required Arguments
 - **`--model_name`** (str): Name of model or link to model. Choose from [List of supported models](../README.md#list-of-supported-models) or use custom models:
   - ESM models: `esm2_t33_650M_UR50D`
+  - ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start)
   - Hugging Face models: `username/model-name`
   - Custom PyTorch models: `/path/to/model.pt` or `/path/to/model_directory/`
   - Local HF models: `/path/to/local_hf_directory/`

{pepe_cli-1.1 → pepe_cli-1.3.0}/README.md

@@ -2,9 +2,16 @@
 
 PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings and attention matrices from protein sequences using pre-trained models. This tool supports various configurations for extracting embeddings and attention matrices, including options for handling CDR3 sequences. Currently implemented models are ESM2 from the 2023 paper ["Evolutionary-scale prediction of atomic-level protein structure with a language model"](https://science.org/doi/10.1126/science.ade2574) and AntiBERTa2-CSSP from the 2023 conference paper ["Enhancing Antibody Language Models with Structural Information"](https://www.mlsb.io/papers_2023/Enhancing_Antibody_Language_Models_with_Structural_Information.pdf). PEPE also supports custom PLMs from local files or from Huggingface Hub addresses. 
 
+### Citation
+> **PEPE: Scalable extraction of multi-modal protein language model representations**
+> Jahn Zhong, Niccolò Cardente, Geir Kjetil Sandve, Habib Bashour, Maria Francesca Abbate, Victor Greiff
+> *bioRxiv* (2026) 
+> [DOI: 10.1101/2025.10.13.680902](https://doi.org/10.1101/2025.10.13.680902)
+
 ## Quick start
 
-1. Install PEPE \
+1. Install PEPE
+
     From PyPI:    
     ```sh
     pip install pepe-cli
@@ -19,7 +26,13 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
     cd pepe-cli
     pip install .
     ```
-2. Run the embedding script:\
+
+2. *(Optional)* For ESMC models (e.g. `biohub/ESMC-300M`), install Biohub's transformers fork:
+    ```sh
+    pip install git+https://github.com/Biohub/transformers.git@main
+    ```
+
+3. Extract embeddings:\
     Extract mean pooled embeddings from protein amino acid sequences in FASTA file:
     ```sh
     pepe --experiment_name <optional_string> --fasta_path <file_path> --output_path <directory> --model_name <model_name>
@@ -29,7 +42,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
 
 PEPE can also be used as a Python library. This allows for programmatic access to protein embeddings without using the command-line interface.
 
-1. Install PEPE (see [Quick start with CLI](#quick-start-with-cli) for details).
+1. Install PEPE (see above for details).
 2. Use the `pepe.embed()` function in your script:
 
 ```python
@@ -182,6 +195,10 @@ results = pepe.embed(
     - RoFormer models
         - alchemab/antiberta2-cssp
         - alchemab/antiberta2
+    - ESMC models (requires Biohub transformers fork; see install instructions above)
+        - biohub/ESMC-300M
+        - biohub/ESMC-600M
+        - biohub/ESMC-6B
     - Custom Hugging Face models
         - Any compatible model from Hugging Face Hub: `username/model-name`
         - Private models with authentication
@@ -196,6 +213,7 @@ results = pepe.embed(
 ### Required Arguments
 - **`--model_name`** (str): Name of model or link to model. Choose from [List of supported models](../README.md#list-of-supported-models) or use custom models:
   - ESM models: `esm2_t33_650M_UR50D`
+  - ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start)
   - Hugging Face models: `username/model-name`
   - Custom PyTorch models: `/path/to/model.pt` or `/path/to/model_directory/`
   - Local HF models: `/path/to/local_hf_directory/`
@@ -239,4 +257,4 @@ When False, all outputs are stored in RAM and written to disk at once after comp
 
 ### Performance Configuration
 - **`--num_workers`** (int, optional): Number of workers for asynchronous data writing. Only relevant when `--streaming_output` is enabled. Default is `8`.
-- **`--flush_batches_after`** (int, optional): Size (in MB) of outputs to accumulate in RAM per worker before flushing to disk. Default is `128`.
+- **`--flush_batches_after`** (int, optional): Size (in MB) of outputs to accumulate in RAM per worker before flushing to disk. Default is `128`.

{pepe_cli-1.1 → pepe_cli-1.3.0}/examples/model_selection.md

@@ -10,7 +10,7 @@ pepe \
 --output_path "src/tests/test_files/test_output"
 ```
 ## Models hosted on Huggingface
-The user can also choose from any PLM with the same input and output formats as the supported models hosted on    Huggingface Hub:
+The user can also choose from PLMs hosted on Huggingface Hub by passing the repository path as `--model_name`:
 ```sh
 pepe \
     --experiment_name "test" \
@@ -19,6 +19,20 @@ pepe \
     --output_path "src/tests/test_files/test_output"
 ```
 
+### ESMC models
+ESMC (ESM Cambrian) is a distinct architecture from ESM2. It requires Biohub's transformers fork, which coexists with ESM1's `fair-esm` package:
+```sh
+pip install git+https://github.com/Biohub/transformers.git@main
+pepe \
+    --experiment_name "esmc_test" \
+    --model_name "biohub/ESMC-300M" \
+    --fasta_path "src/tests/test_files/test.fasta" \
+    --output_path "src/tests/test_files/test_output" \
+    --extract_embeddings mean_pooled \
+    --device cpu \
+    --layers -1
+```
+
 ## Custom models
 To specify your own custom model, pass the path to a directory containing:
 - .pt or .pth file (model weights)

{pepe_cli-1.1 → pepe_cli-1.3.0}/pyproject.toml

@@ -15,7 +15,7 @@ requires-python = ">=3.8"
 classifiers = [
     "Development Status :: 3 - Alpha",
     "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: GNU Affero General Public License v3",
+    "License :: OSI Approved :: MIT License",
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.8",
@@ -33,6 +33,7 @@ dependencies = [
     "numpy",
     "protobuf",
     "alive_progress",
+    "rjieba",
 ]
 
 [project.urls]

{pepe_cli-1.1 → pepe_cli-1.3.0}/requirements.txt

@@ -5,4 +5,5 @@ fair-esm
 sentencepiece
 numpy
 protobuf
-alive_progress
+alive_progress
+rjieba

{pepe_cli-1.1 → pepe_cli-1.3.0}/setup.py

@@ -28,7 +28,7 @@ setup(
     classifiers=[
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Science/Research",
-        "License :: OSI Approved :: GNU Affero General Public License v3",
+        "License :: OSI Approved :: MIT License",
         "Operating System :: Linux",
         "Operating System :: macOS",
         "Programming Language :: Python :: 3",
@@ -47,6 +47,7 @@ setup(
         "numpy",
         "protobuf",
         "alive_progress",
+        "rjieba",
     ],
     entry_points={
         "console_scripts": [

{pepe_cli-1.1 → pepe_cli-1.3.0}/src/pepe/__init__.py

@@ -7,7 +7,7 @@ except ImportError:
     pass
 
 # Package metadata - single source of truth
-__version__ = "1.1"
+__version__ = "1.3.0"
 __package_name__ = "pepe-cli"
 __module_name__ = "pepe"
 __author__ = "Jahn Zhong"

{pepe_cli-1.1 → pepe_cli-1.3.0}/src/pepe/embedders/custom_embedder.py

@@ -288,9 +288,10 @@ class CustomEmbedder(BaseEmbedder):
 
     def _load_layers(self, layers):
         """Process layer specification."""
+        if layers is None:
+            return list(range(1, self.num_layers + 1))
         if not layers:
-            layers = list(range(1, self.num_layers + 1))
-            return layers
+            layers = [-1]
 
         # Validate layer indices
         assert all(

{pepe_cli-1.1 → pepe_cli-1.3.0}/src/pepe/embedders/esm_embedder.py

@@ -102,9 +102,10 @@ class ESMEmbedder(BaseEmbedder):
         return special_token_ids
 
     def _load_layers(self, layers):
+        if layers is None:
+            return list(range(1, self.model.num_layers + 1))  # type: ignore
         if not layers:
-            layers = list(range(1, self.model.num_layers + 1))  # type: ignore
-            return layers
+            layers = [-1]
         # 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

{pepe_cli-1.1 → pepe_cli-1.3.0}/src/pepe/embedders/huggingface_embedder.py

@@ -15,6 +15,7 @@ def _import_transformers():
             RoFormerSinusoidalPositionalEmbedding,
         )
         from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
+        from transformers import AutoModelForMaskedLM
 
         return (
             T5EncoderModel,
@@ -25,6 +26,7 @@ def _import_transformers():
             AutoModel,
             AutoTokenizer,
             AutoModelForCausalLM,
+            AutoModelForMaskedLM,
         )
     except ImportError as e:
         logger.error(f"Failed to import transformers: {e}")
@@ -51,20 +53,13 @@ class HuggingfaceEmbedder(BaseEmbedder):
 
     def _load_layers(self, layers):
         """Check if the specified representation layers are valid."""
+        num_layers = self.num_layers  # type: ignore
+        if layers is None:
+            return list(range(1, num_layers + 1))
         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
-        ]
+            layers = [-1]
+        assert all(-(num_layers + 1) <= i <= num_layers for i in layers)
+        layers = [(i + num_layers + 1) % (num_layers + 1) for i in layers]
         return layers
 
     def _load_data(self, sequences, substring_dict, bracket_type):
@@ -116,14 +111,18 @@ class HuggingfaceEmbedder(BaseEmbedder):
         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
-            }
+            dtype = self._precision_to_dtype(self.precision, "torch")
+            hidden_states = outputs.hidden_states
+            if isinstance(hidden_states, torch.Tensor):
+                representations = {
+                    layer: hidden_states[layer].to(dtype).cpu()
+                    for layer in self.layers  # type: ignore
+                }
+            else:
+                representations = {
+                    layer: hidden_states[layer].to(dtype).cpu()
+                    for layer in self.layers  # type: ignore
+                }
             torch.cuda.empty_cache()
         else:
             representations = None
@@ -182,6 +181,7 @@ class Antiberta2Embedder(HuggingfaceEmbedder):
             AutoModel,
             AutoTokenizer,
             AutoModelForCausalLM,
+            AutoModelForMaskedLM,
         ) = _import_transformers()
 
         tokenizer = RoFormerTokenizer.from_pretrained(model_link, use_fast=True)
@@ -246,6 +246,7 @@ class T5Embedder(HuggingfaceEmbedder):
             AutoModel,
             AutoTokenizer,
             AutoModelForCausalLM,
+            AutoModelForMaskedLM,
         ) = _import_transformers()
 
         tokenizer = T5Tokenizer.from_pretrained(model_link, use_fast=True)
@@ -257,6 +258,272 @@ class T5Embedder(HuggingfaceEmbedder):
         return model, tokenizer, num_heads, num_layers, embedding_size
 
 
+def _resolve_esm2_model_link(model_name):
+    if "/" in model_name:
+        return model_name
+    return f"facebook/{model_name}"
+
+
+class ESM2Embedder(HuggingfaceEmbedder):
+    """ESM-2 embedder using HuggingFace transformers instead of fair-esm."""
+
+    def __init__(self, args):
+        BaseEmbedder.__init__(self, 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_name)
+        self.valid_tokens = self._get_valid_tokens()
+        self.bracket_type = pepe.utils.get_bracket_type(self.tokenizer)
+        self._check_max_input_length()
+        pepe.utils.check_input_tokens(
+            self.valid_tokens,
+            self.sequences,
+            self.model_name,
+            split_long_sequences=self.split_long_sequences,
+        )
+        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_input_length = self._load_data(
+            self.sequences, self.substring_dict, self.bracket_type
+        )
+        self._set_output_objects()
+
+    def _get_valid_tokens(self):
+        return {tok for tok in self.tokenizer.get_vocab().keys() if len(tok) == 1}
+
+    def _load_data(self, sequences, substring_dict, bracket_type):
+        dataset = pepe.utils.HuggingFaceDataset(
+            sequences,
+            substring_dict,
+            self.context,
+            bracket_type,
+            self.tokenizer,
+            self.max_input_length,
+            add_special_tokens=not self.disable_special_tokens,
+        )
+        dataset.pad_token_id = self.tokenizer.pad_token_id
+        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 _initialize_model(self, model_name):
+        model_link = _resolve_esm2_model_link(model_name)
+        if torch.cuda.is_available() and self.device.type == "cuda":
+            device = torch.device("cuda")
+            logger.info("Transferred model to GPU")
+        else:
+            device = torch.device("cpu")
+            logger.info("No GPU available, using CPU")
+
+        (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
+            AutoModelForMaskedLM,
+        ) = _import_transformers()
+
+        logger.info(f"Loading ESM-2 model from HuggingFace: {model_link}")
+        tokenizer = AutoTokenizer.from_pretrained(model_link)
+        model_kwargs = {}
+        if self.return_contacts:
+            model_kwargs["attn_implementation"] = "eager"
+
+        if self.return_logits:
+            model = AutoModelForMaskedLM.from_pretrained(
+                model_link, **model_kwargs
+            ).to(device)
+        else:
+            model = AutoModel.from_pretrained(model_link, **model_kwargs).to(device)
+        model.eval()
+
+        config = model.config
+        num_heads = config.num_attention_heads
+        num_layers = config.num_hidden_layers
+        embedding_size = config.hidden_size
+        return model, tokenizer, num_heads, num_layers, embedding_size
+
+    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_logits:
+            logits = (
+                outputs.logits
+                .to(dtype=self._precision_to_dtype(self.precision, "torch"))
+                .permute(2, 0, 1)
+                .cpu()
+            )
+            torch.cuda.empty_cache()
+        else:
+            logits = None
+
+        if return_contacts:
+            attention_matrices = (
+                torch.stack(outputs.attentions)  # type: ignore
+                .to(self._precision_to_dtype(self.precision, "torch"))  # type: ignore
+                .cpu()
+            )
+            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
+
+        return logits, representations, attention_matrices
+
+
+def _get_config_attr(config, *names, default=None):
+    for name in names:
+        if hasattr(config, name):
+            return getattr(config, name)
+    return default
+
+
+class ESMCEmbedder(HuggingfaceEmbedder):
+    """ESMC embedder using Biohub transformers fork (model_type esmc)."""
+
+    def __init__(self, args):
+        BaseEmbedder.__init__(self, 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 = self._get_valid_tokens()
+        self.bracket_type = pepe.utils.get_bracket_type(self.tokenizer)
+        self._check_max_input_length()
+        pepe.utils.check_input_tokens(
+            self.valid_tokens,
+            self.sequences,
+            self.model_name,
+            split_long_sequences=self.split_long_sequences,
+        )
+        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_input_length = self._load_data(
+            self.sequences, self.substring_dict, self.bracket_type
+        )
+        self._set_output_objects()
+
+    def _get_valid_tokens(self):
+        return {tok for tok in self.tokenizer.get_vocab().keys() if len(tok) == 1}
+
+    def _load_data(self, sequences, substring_dict, bracket_type):
+        dataset = pepe.utils.HuggingFaceDataset(
+            sequences,
+            substring_dict,
+            self.context,
+            bracket_type,
+            self.tokenizer,
+            self.max_input_length,
+            add_special_tokens=not self.disable_special_tokens,
+            gapped_sequences=False,
+        )
+        dataset.pad_token_id = self.tokenizer.pad_token_id
+        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 _initialize_model(self, model_link):
+        if torch.cuda.is_available() and self.device.type == "cuda":
+            device = torch.device("cuda")
+            logger.info("Transferred model to GPU")
+        else:
+            device = torch.device("cpu")
+            logger.info("No GPU available, using CPU")
+
+        (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
+            AutoModelForMaskedLM,
+        ) = _import_transformers()
+
+        logger.info(f"Loading ESMC model from HuggingFace: {model_link}")
+        tokenizer = AutoTokenizer.from_pretrained(model_link)
+        model_kwargs = {}
+        if self.return_contacts:
+            model_kwargs["attn_implementation"] = "eager"
+
+        if self.return_logits:
+            model = AutoModelForMaskedLM.from_pretrained(
+                model_link, **model_kwargs
+            ).to(device)
+        else:
+            model = AutoModel.from_pretrained(model_link, **model_kwargs).to(device)
+        model.eval()
+
+        config = model.config
+        num_heads = _get_config_attr(
+            config, "num_attention_heads", "n_heads", "num_heads"
+        )
+        num_layers = _get_config_attr(
+            config, "num_hidden_layers", "n_layers", "num_layers"
+        )
+        embedding_size = _get_config_attr(
+            config, "hidden_size", "d_model", "embed_dim"
+        )
+        return model, tokenizer, num_heads, num_layers, embedding_size
+
+
 class GenericHuggingFaceEmbedder(HuggingfaceEmbedder):
     """Generic HuggingFace embedder that can handle models with unknown architectures using AutoModel and AutoTokenizer."""
 
@@ -320,6 +587,7 @@ class GenericHuggingFaceEmbedder(HuggingfaceEmbedder):
             AutoModel,
             AutoTokenizer,
             AutoModelForCausalLM,
+            AutoModelForMaskedLM,
         ) = _import_transformers()
 
         # For models that commonly require custom code, use trust_remote_code=True immediately