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

{pepe_cli-1.0.3 → pepe_cli-1.0.4}/MANIFEST.in

@@ -8,3 +8,9 @@ recursive-exclude * __pycache__
 recursive-exclude * *.py[co]
 recursive-exclude * *.so
 recursive-exclude * .DS_Store
+recursive-exclude .github *
+recursive-exclude .git *
+recursive-exclude .venv *
+recursive-exclude * .pytest_cache
+recursive-exclude * .coverage
+exclude .gitignore

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.0.3
+Version: 1.0.4
 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.3 → pepe_cli-1.0.4}/src/pepe/__init__.py

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

{pepe_cli-1.0.3 → pepe_cli-1.0.4}/src/pepe/embedders/custom_embedder.py

@@ -53,8 +53,9 @@ class CustomEmbedder(BaseEmbedder):
 
         # Set up tokenizer and validate tokens
         self.valid_tokens = self._get_valid_tokens()
+        self.bracket_type = pepe.utils.get_bracket_type(self.tokenizer)
         pepe.utils.check_input_tokens(
-            self.valid_tokens, self.sequences, self.model_name
+            self.valid_tokens, self.sequences, self.model_name, 
         )
 
         # Set up special tokens
@@ -65,7 +66,7 @@ class CustomEmbedder(BaseEmbedder):
 
         # Load and tokenize data
         self.data_loader, self.max_length = self._load_data(
-            self.sequences, self.substring_dict
+            self.sequences, self.substring_dict, self.bracket_type
         )
 
         # Initialize output objects
@@ -297,13 +298,14 @@ class CustomEmbedder(BaseEmbedder):
         layers = [(i + self.num_layers + 1) % (self.num_layers + 1) for i in layers]
         return layers
 
-    def _load_data(self, sequences, substring_dict=None):
+    def _load_data(self, sequences, substring_dict=None, bracket_type=None):
         """Load and tokenize sequences."""
         # Create dataset
-        dataset = CustomDataset(
+        dataset = pepe.utils.CustomDataset(
             sequences,
             substring_dict,
             self.context,
+            bracket_type,
             self.tokenizer,
             self.max_length,
             add_special_tokens=not self.disable_special_tokens,
@@ -543,105 +545,3 @@ class DefaultProteinTokenizer:
             }
 
         return tokenized[0] if len(tokenized) == 1 else tokenized
-
-
-class CustomDataset(pepe.utils.SequenceDictDataset):
-    """
-    Dataset class for custom embedder.
-    """
-
-    def __init__(
-        self,
-        sequences,
-        substring_dict,
-        context,
-        tokenizer,
-        max_length,
-        add_special_tokens=True,
-    ):
-        super().__init__(sequences, substring_dict, context)
-        self.tokenizer = tokenizer
-        self.max_length = max_length
-        self.add_special_tokens = add_special_tokens
-        self.pad_token_id = getattr(tokenizer, "pad_token_id", 0)
-
-        # Encode sequences
-        self.encoded_data = self._encode_sequences(
-            self.data, tokenizer, max_length, add_special_tokens
-        )
-
-        # Handle substring sequences if provided
-        if self.substring_dict:
-            logger.info("Tokenizing substrings...")
-            self.encoded_substring_data = self._encode_sequences(
-                self.filtered_substring_data,
-                tokenizer,
-                "max_length",
-                add_special_tokens=False,
-            )
-            self.substring_masks = self._get_substring_masks()
-
-    def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens):
-        """Encode sequences using the tokenizer."""
-        labels, strs = zip(*data)
-
-        # Convert max_length to int if needed
-        if max_length == "max_length":
-            max_length = max(len(s) for s in strs) + (2 if add_special_tokens else 0)
-
-        encoded_sequences = []
-        for label, seq in zip(labels, strs):
-            # Tokenize the sequence
-            tokens = tokenizer(
-                seq,
-                truncation=True,
-                padding="max_length",
-                max_length=max_length,
-                add_special_tokens=add_special_tokens,
-                return_tensors="pt",
-            )
-
-            encoded_sequences.append(
-                (
-                    label,
-                    seq,
-                    tokens["input_ids"].squeeze(0),
-                    tokens["attention_mask"].squeeze(0),
-                )
-            )
-
-        return encoded_sequences
-
-    def get_max_encoded_length(self):
-        """Get maximum encoded sequence length."""
-        return max(len(toks) for _, _, toks, _ in self.encoded_data)
-
-    def safe_collate(self, batch):
-        """Collate function for data loader."""
-        if self.substring_dict:
-            labels, seqs, toks, attn_masks, substring_masks = zip(*batch)
-            return (
-                list(labels),
-                list(seqs),
-                torch.stack(toks),
-                torch.stack(attn_masks),
-                torch.stack(substring_masks),
-            )
-        else:
-            labels, seqs, toks, attn_masks = zip(*batch)
-            return (
-                list(labels),
-                list(seqs),
-                torch.stack(toks),
-                torch.stack(attn_masks),
-                None,
-            )
-
-    def __getitem__(self, idx):
-        """Get item from dataset."""
-        labels, seqs, toks, attn_mask = self.encoded_data[idx]
-        if self.substring_dict:
-            substring_masks = self.substring_masks[idx]
-            return labels, seqs, toks, attn_mask, substring_masks
-        else:
-            return labels, seqs, toks, attn_mask

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

@@ -14,6 +14,7 @@ def _import_transformers():
         from transformers.models.roformer.modeling_roformer import (
             RoFormerSinusoidalPositionalEmbedding,
         )
+        from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM
 
         return (
             T5EncoderModel,
@@ -21,6 +22,9 @@ def _import_transformers():
             RoFormerTokenizer,
             RoFormerModel,
             RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
         )
     except ImportError as e:
         logger.error(f"Failed to import transformers: {e}")
@@ -63,13 +67,14 @@ class HuggingfaceEmbedder(BaseEmbedder):
         ]
         return layers
 
-    def _load_data(self, sequences, substring_dict):
+    def _load_data(self, sequences, substring_dict, bracket_type):
         """Tokenize sequences and create a DataLoader."""
         # Tokenize sequences
         dataset = pepe.utils.HuggingFaceDataset(
             sequences,
             substring_dict,
             self.context,
+            bracket_type,
             self.tokenizer,  # type: ignore
             self.max_length,
             add_special_tokens=not self.disable_special_tokens,
@@ -139,6 +144,7 @@ class Antiberta2Embedder(HuggingfaceEmbedder):
             self.embedding_size,
         ) = self._initialize_model(self.model_link)
         self.valid_tokens = set(self.tokenizer.get_vocab().keys())
+        self.bracket_type = pepe.utils.get_bracket_type(self.tokenizer)
         pepe.utils.check_input_tokens(
             self.valid_tokens, self.sequences, self.model_name
         )
@@ -147,7 +153,7 @@ class Antiberta2Embedder(HuggingfaceEmbedder):
         )
         self.layers = self._load_layers(self.layers)
         self.data_loader, self.max_length = self._load_data(
-            self.sequences, self.substring_dict
+            self.sequences, self.substring_dict, self.bracket_type
         )
         self._set_output_objects()
         assert self.max_length <= 256, "AntiBERTa2 only supports max_length <= 256"
@@ -168,6 +174,9 @@ class Antiberta2Embedder(HuggingfaceEmbedder):
             RoFormerTokenizer,
             RoFormerModel,
             RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
         ) = _import_transformers()
 
         tokenizer = RoFormerTokenizer.from_pretrained(model_link, use_fast=True)
@@ -192,15 +201,16 @@ class T5Embedder(HuggingfaceEmbedder):
             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)
         pepe.utils.check_input_tokens(
-            self.valid_tokens, self.sequences, self.model_name
+            self.valid_tokens, self.sequences, self.model_name, self.bracket_type
         )
         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.sequences, self.substring_dict, self.bracket_type
         )
         self._set_output_objects()
 
@@ -228,6 +238,9 @@ class T5Embedder(HuggingfaceEmbedder):
             RoFormerTokenizer,
             RoFormerModel,
             RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
         ) = _import_transformers()
 
         tokenizer = T5Tokenizer.from_pretrained(model_link, use_fast=True)
@@ -237,3 +250,157 @@ class T5Embedder(HuggingfaceEmbedder):
         num_layers = model.config.num_layers
         embedding_size = model.config.hidden_size
         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."""
+
+    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 = self._get_valid_tokens()
+        self.bracket_type = pepe.utils.get_bracket_type(self.tokenizer)
+        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.bracket_type
+        )
+        self._set_output_objects()
+
+    def _get_valid_tokens(self):
+        """Get valid tokens from the tokenizer."""
+        if hasattr(self.tokenizer, "get_vocab"):
+            vocab = self.tokenizer.get_vocab()
+            # Handle different tokenizer types
+            if hasattr(self.tokenizer, "decoder") and self.tokenizer.decoder:
+                # T5-style tokenizer
+                valid_tokens = set(
+                    k[1:] if k.startswith("▁") else k for k in vocab.keys()
+                )
+            else:
+                # Standard tokenizer
+                valid_tokens = set(vocab.keys())
+            return valid_tokens
+        return set()
+
+    def _initialize_model(self, model_link):
+        """Initialize the model, tokenizer, and device using AutoModel and AutoTokenizer."""
+        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")
+
+        # Lazy import transformers components
+        (
+            T5EncoderModel,
+            T5Tokenizer,
+            RoFormerTokenizer,
+            RoFormerModel,
+            RoFormerSinusoidalPositionalEmbedding,
+            AutoModel,
+            AutoTokenizer,
+            AutoModelForCausalLM,
+        ) = _import_transformers()
+
+        # For models that commonly require custom code, use trust_remote_code=True immediately
+        requires_custom_code = any(
+            pattern in model_link.lower()
+            for pattern in ["progen", "protgpt", "esm-fold", "esmfold", "alphafold"]
+        )
+
+        if requires_custom_code:
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_link, use_fast=True, trust_remote_code=True
+            )
+            try:
+                model = AutoModel.from_pretrained(
+                    model_link, trust_remote_code=True
+                ).to(device)
+            except ValueError as model_error:
+                # If AutoModel fails, try AutoModelForCausalLM
+                model = AutoModelForCausalLM.from_pretrained(
+                    model_link, trust_remote_code=True
+                ).to(device)
+        else:
+            # Try without trust_remote_code first, then with it if needed
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(model_link, use_fast=True)
+                try:
+                    model = AutoModel.from_pretrained(model_link).to(device)
+                except ValueError as model_error:
+                    # If AutoModel fails, try AutoModelForCausalLM
+                    model = AutoModelForCausalLM.from_pretrained(model_link).to(device)
+            except (OSError, ValueError) as e:
+                # If the model requires custom code, try with trust_remote_code=True
+                if (
+                    "trust_remote_code" in str(e).lower()
+                    or "custom code" in str(e).lower()
+                ):
+                    tokenizer = AutoTokenizer.from_pretrained(
+                        model_link, use_fast=True, trust_remote_code=True
+                    )
+                    try:
+                        model = AutoModel.from_pretrained(
+                            model_link, trust_remote_code=True
+                        ).to(device)
+                    except ValueError as model_error:
+                        # If AutoModel fails, try AutoModelForCausalLM
+                        model = AutoModelForCausalLM.from_pretrained(
+                            model_link, trust_remote_code=True
+                        ).to(device)
+                else:
+                    raise
+
+        # Handle tokenizer padding token
+        if tokenizer.pad_token is None:
+            if tokenizer.eos_token is not None:
+                tokenizer.pad_token = tokenizer.eos_token
+            elif tokenizer.unk_token is not None:
+                tokenizer.pad_token = tokenizer.unk_token
+            else:
+                # Add a special padding token
+                tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+                model.resize_token_embeddings(len(tokenizer))
+
+        model.eval()
+
+        # Get model configuration
+        config = model.config
+        num_heads = getattr(
+            config,
+            "num_attention_heads",
+            getattr(config, "num_heads", getattr(config, "n_head", 12)),
+        )
+        num_layers = getattr(
+            config,
+            "num_hidden_layers",
+            getattr(config, "num_layers", getattr(config, "n_layer", 12)),
+        )
+        embedding_size = getattr(
+            config,
+            "hidden_size",
+            getattr(config, "d_model", getattr(config, "embed_dim", 768)),
+        )
+
+        logger.info(f"Loaded generic HuggingFace model: {model_link}")
+        logger.info(f"Model type: {config.model_type}")
+        logger.info(
+            f"Number of heads: {num_heads}, layers: {num_layers}, embedding size: {embedding_size}"
+        )
+
+        return model, tokenizer, num_heads, num_layers, embedding_size

{pepe_cli-1.0.3 → pepe_cli-1.0.4}/src/pepe/utils.py

@@ -7,7 +7,7 @@ import time
 import gc
 import os
 import json
-from transformers import RoFormerTokenizer
+from transformers import RoFormerTokenizer, T5Tokenizer
 import threading, queue
 from alive_progress import alive_bar
 import shutil
@@ -43,7 +43,7 @@ class TokenBudgetBatchSampler:
 
 
 class SequenceDictDataset(Dataset):
-    def __init__(self, sequences, substring_dict, context):
+    def __init__(self, sequences, substring_dict, context, bracket_type = "square"):
         self.data = list(sequences.items())  # (label, seq)
         if substring_dict:
             self.substring_dict = substring_dict
@@ -53,7 +53,7 @@ class SequenceDictDataset(Dataset):
             self.substring_dict = None
             self.context = None
             self.filtered_substring_data = None
-        pass
+        self.bracket_type = bracket_type
 
     def __getitem__(self, idx):
         return self.data[idx]
@@ -128,13 +128,15 @@ class HuggingFaceDataset(SequenceDictDataset):
         sequences,
         substring_dict,
         context,
+        bracket_type,
         tokenizer,
         max_length,
         add_special_tokens=True,
     ):
-        super().__init__(sequences, substring_dict, context)
+        super().__init__(sequences, substring_dict, context, bracket_type)
+        
         self.encoded_data = self._encode_sequences(
-            self.data, tokenizer, max_length, add_special_tokens
+            self.data, tokenizer, max_length, add_special_tokens, gapped_sequences = True
         )  # (label, seq, toks, attention_mask)
         self.pad_token_id = tokenizer.pad_token_type_id
         if self.substring_dict:
@@ -144,16 +146,14 @@ class HuggingFaceDataset(SequenceDictDataset):
                 tokenizer,
                 "max_length",
                 add_special_tokens=False,
+                gapped_sequences=True
             )  # (label, seq, toks, attention_mask)
             self.substring_masks = self._get_substring_masks()
 
-    def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens):
+    def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gapped_sequences = True):
         labels, strs = zip(*data)
-        if isinstance(tokenizer, RoFormerTokenizer):
-            # RoFormerTokenizer requires space-separated tokenization
-            # for the input sequences
-            logger.info("Using RoFormerTokenizer, applying gap_sequence.")
-            strs = self._gap_sequence(strs)
+        if gapped_sequences:
+            strs = gap_sequence(strs, self.bracket_type)
             max_token_length = max(len(seq.split(" ")) for seq in strs)
         else:
             # For other tokenizers, use the default tokenization
@@ -187,11 +187,6 @@ class HuggingFaceDataset(SequenceDictDataset):
         attention_masks = torch.cat(loop_attention_mask, dim=0)
         return list(zip(labels, strs, list(toks), list(attention_masks)))
 
-    def _gap_sequence(self, sequences: Sequence[str]) -> Sequence[str]:
-        """Space-separated tokenization for RoFormer input."""
-        seqs = [" ".join(re.findall(r"\[.*?\]|.", sequence)) for sequence in sequences]
-        return seqs
-
     def get_max_encoded_length(self):
         return max(len(toks) for _, _, toks, _ in self.encoded_data)
 
@@ -319,21 +314,137 @@ class ESMDataset(SequenceDictDataset):
         else:
             return labels, seqs, toks, None, None
 
+class CustomDataset(SequenceDictDataset):
+    """
+    Dataset class for custom embedder.
+    """
+
+    def __init__(
+        self,
+        sequences,
+        substring_dict,
+        context,
+        bracket_type,
+        tokenizer,
+        max_length,
+        add_special_tokens=True,
+        gapped_sequences = True,
+    ):
+        super().__init__(sequences, substring_dict, context, bracket_type)
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.add_special_tokens = add_special_tokens
+        self.pad_token_id = getattr(tokenizer, "pad_token_id", 0)
+        self.gapped_sequences = gapped_sequences
+        self.bracket_type = bracket_type
+
+        # Encode sequences
+        self.encoded_data = self._encode_sequences(
+            self.data, tokenizer, max_length, add_special_tokens, self.gapped_sequences
+        )
+
+        # Handle substring sequences if provided
+        if self.substring_dict:
+            logger.info("Tokenizing substrings...")
+            self.encoded_substring_data = self._encode_sequences(
+                self.filtered_substring_data,
+                tokenizer,
+                "max_length",
+                add_special_tokens=False,
+            )
+            self.substring_masks = self._get_substring_masks()
+
+    def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gapped_sequences = True):
+        """Encode sequences using the tokenizer."""
+        labels, strs = zip(*data)
+
+        if gapped_sequences:
+            strs = gap_sequence(strs, self.bracket_type)
+            max_token_length = max(len(seq.split(" ")) for seq in strs)
+        else:
+            # For other tokenizers, use the default tokenization
+            max_token_length = max(len(seq) for seq in strs)
+
+        # Convert max_length to int if needed
+        if max_length == "max_length":
+            max_token_length = max(len(s) for s in strs) + (2 if add_special_tokens else 0)
+
+        encoded_sequences = []
+        for label, seq in zip(labels, strs):
+            # Tokenize the sequence
+            tokens = tokenizer(
+                seq,
+                truncation=True,
+                padding="max_length",
+                max_length=max_token_length,
+                add_special_tokens=add_special_tokens,
+                return_tensors="pt",
+            )
+
+            encoded_sequences.append(
+                (
+                    label,
+                    seq,
+                    tokens["input_ids"].squeeze(0),
+                    tokens["attention_mask"].squeeze(0),
+                )
+            )
+
+        return encoded_sequences
+
+    def get_max_encoded_length(self):
+        """Get maximum encoded sequence length."""
+        return max(len(toks) for _, _, toks, _ in self.encoded_data)
+
+    def safe_collate(self, batch):
+        """Collate function for data loader."""
+        if self.substring_dict:
+            labels, seqs, toks, attn_masks, substring_masks = zip(*batch)
+            return (
+                list(labels),
+                list(seqs),
+                torch.stack(toks),
+                torch.stack(attn_masks),
+                torch.stack(substring_masks),
+            )
+        else:
+            labels, seqs, toks, attn_masks = zip(*batch)
+            return (
+                list(labels),
+                list(seqs),
+                torch.stack(toks),
+                torch.stack(attn_masks),
+                None,
+            )
+
+    def __getitem__(self, idx):
+        """Get item from dataset."""
+        labels, seqs, toks, attn_mask = self.encoded_data[idx]
+        if self.substring_dict:
+            substring_masks = self.substring_masks[idx]
+            return labels, seqs, toks, attn_mask, substring_masks
+        else:
+            return labels, seqs, toks, attn_mask
+
+def check_input_tokens(valid_tokens, sequences, model_name, bracket_type = "square"):
+    def __str_to_list(sequence, bracket_type = "square"):
+        if bracket_type == "square":
+            return re.findall(r"\[.*?\]|.", sequence)
+        elif bracket_type == "angle":
+            return re.findall(r"<.*?>|.", sequence)
+        else:
+            raise ValueError(f"Invalid bracket type: {bracket_type}")
 
-def check_input_tokens(valid_tokens, sequences, model_name):
     with alive_bar(
         len(sequences), title="Checking input sequences for invalid tokens..."
     ) as bar:
         for label, sequence in sequences.items():
-            if "esm" not in model_name:
-                sequence = re.findall(r"\[.*?\]|.", sequence)
-                if "antiberta" in model_name:  # check for longest sequence
-                    assert (
-                        len(sequence) <= 256
-                    ), f"Antiberta2 does not support sequences longer than 256 tokens. Found {len(sequence)} tokens in sequence {label}."
+            sequence = __str_to_list(sequence, bracket_type)
+            if "antiberta" in model_name:  # check for longest sequence
+                assert (
+                    len(sequence) <= 256
+                ), f"Antiberta2 does not support sequences longer than 256 tokens. Found {len(sequence)} tokens in sequence {label}."
 
-            else:
-                sequence = re.findall(r"<.*?>|.", sequence)
             if not set(sequence).issubset(valid_tokens):
                 raise ValueError(
                     f"Invalid tokens found in sequence {label}: {set(sequence) - set(valid_tokens)}"
@@ -365,7 +476,25 @@ def fasta_to_dict(fasta_path):
     flush()
 
     return seq_dict
-
+def get_bracket_type(tokenizer):
+    try:
+        special_token_bracket = tokenizer.special_tokens_map["additional_special_tokens"][0][0]
+    except KeyError:
+        special_token_bracket = tokenizer.all_special_tokens[0][0]
+    if special_token_bracket == "<":
+        return "angle"
+    elif special_token_bracket == "[":
+        return "square"
+    else:
+        raise ValueError(f"Invalid special token bracket: {special_token_bracket}")
+def gap_sequence(sequences: Sequence[str], bracket_type = "square") -> Sequence[str]:
+    if bracket_type == "square":
+        seqs = [" ".join(re.findall(r"\[.*?\]|.", sequence)) for sequence in sequences]
+    elif bracket_type == "angle":
+        seqs = [" ".join(re.findall(r"<.*?>|.", sequence)) for sequence in sequences]
+    else:
+        raise ValueError(f"Invalid bracket type: {bracket_type}")
+    return seqs
 
 def flush_memmaps(obj):
     """Recursively flush memory maps."""

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.0.3
+Version: 1.0.4
 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.3 → pepe_cli-1.0.4}/src/tests/test_run.py

@@ -10,7 +10,7 @@ sys.argv = [
     "--experiment_name",
     "test",
     "--model_name",
-    "/home/jahn/embedairr/examples/custom_model/example_protein_model",
+    "ZYMScott/antiberty",
     "--fasta_path",
     "src/tests/test_files/test.fasta",
     "--output_path",