pepe-cli 1.0.0__tar.gz

pepe_cli-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Greiff Lab – Computational Systems Immunology
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pepe_cli-1.0.0/MANIFEST.in

@@ -0,0 +1,10 @@
+include README.md
+include LICENSE
+include requirements.txt
+include pyproject.toml
+recursive-include src *.py
+recursive-include examples *
+recursive-exclude * __pycache__
+recursive-exclude * *.py[co]
+recursive-exclude * *.so
+recursive-exclude * .DS_Store

pepe_cli-1.0.0/PKG-INFO

@@ -0,0 +1,137 @@
+Metadata-Version: 2.4
+Name: pepe-cli
+Version: 1.0.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
+Author-email: Jahn Zhong <jahn.zhong@medisin.uio.no>
+License: MIT
+Project-URL: Homepage, https://github.com/csi-greifflab/pepe-cli
+Project-URL: Bug Reports, https://github.com/csi-greifflab/pepe-cli/issues
+Project-URL: Source, https://github.com/csi-greifflab/pepe-cli
+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: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=1.9.0
+Requires-Dist: transformers>=4.20.0
+Requires-Dist: fair-esm
+Requires-Dist: sentencepiece
+Requires-Dist: numpy
+Requires-Dist: protobuf
+Requires-Dist: rjieba
+Requires-Dist: alive_progress
+Dynamic: author
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-python
+
+# PEPE
+
+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. 
+
+## Quick start
+
+1. Install PEPE \
+    From PyPI:    
+    ```sh
+    pip install pepe-cli
+    ```
+    Or install from the GitHub repository:    
+    ```sh
+    git clone https://github.com/csi-greifflab/pepe-cli
+    cd pepe-cli
+    pip install .
+    ```
+2. Run the embedding script:\
+    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>
+    ```
+
+## List of supported models:
+- ESM-family models
+    - ESM1:
+        - esm1_t34_670M_UR50S
+        - esm1_t34_670M_UR50D
+        - esm1_t34_670M_UR100
+        - esm1_t12_85M_UR50S
+        - esm1_t6_43M_UR50S
+        - esm1b_t33_650M_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
+    - ESM2:
+        - esm2_t6_8M_UR50D
+        - esm2_t12_35M_UR50D
+        - esm2_t30_150M_UR50D
+        - esm2_t33_650M_UR50D
+        - esm2_t36_3B_UR50D
+        - esm2_t48_15B_UR50D
+- Huggingface Transformer models
+    - T5 transformer models
+        - Rostlab/prot_t5_xl_half_uniref50-enc
+        - Rostlab/ProstT5
+    - RoFormer models
+        - alchemab/antiberta2-cssp
+        - alchemab/antiberta2
+    - Custom Hugging Face models
+        - Any compatible model from Hugging Face Hub: `username/model-name`
+        - Private models with authentication
+        - Local Hugging Face models
+- Custom Models
+    - Load your own PyTorch models with custom tokenizers
+    - Create example with: `python examples/custom_model/create_example_custom_model.py`
+
+
+## Arguments
+
+### 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`
+  - 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/`
+- **`--fasta_path`** (str): Path to the input FASTA file. If no experiment name is provided, the output files will be named after the input file.
+- **`--output_path`** (str): Directory for output files. Will generate a subdirectory for outputs of each output type.
+
+### Model Configuration
+- **`--tokenizer_from`** (str, optional): Huggingface address of the tokenizer to use. If not provided, will attempt to search for tokenizer packaged with model. If using a custom model, provide the path to the tokenizer directory.
+- **`--disable_special_tokens`** (bool, optional): When True, PEPE disables pre- and appending BOS/CLS and EOS/SEP tokens before embedding. Default is `False`.
+- **`--device`** (str, optional): Device to run the model on. Choose from `cuda` or `cpu`. Default is `cuda`.
+
+### Embedding Configuration
+- **`--layers`** (str, optional): Representation layers to extract from the model. Default is the last layer. Example: `--layers -1 6`.
+- **`--extract_embeddings`** (str, optional): 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 `mean_pooled`.
+- **`--substring_path`** (str, optional): Path to a CSV file with columns "sequence_id" and "substring". Only required when selecting "substring_pooled" option.
+- **`--context`** (int, optional): Only specify when including "substring_pooled" in `--extract_embeddings` option. Number of amino acids to include before and after the substring sequence. Default is `0`.
+
+### Processing Configuration
+- **`--batch_size`** (int, optional): Batch size for loading sequences. Default is `1024`. Decrease if encountering out-of-memory errors.
+- **`--max_length`** (int, optional): Length to which sequences will be padded. Default is length of longest sequence in input file. If shorter than longest sequence, will forcefully default to length of longest sequence.
+- **`--discard_padding`** (bool, optional): Discard padding tokens from per_token embeddings output. Default is `False`.
+
+### Output Configuration
+- **`--experiment_name`** (str, optional): Prefix for names of output files. If not provided, name of input file will be used for prefix.
+- **`--streaming_output`** (bool, optional): PEPE preallocates the required disk space and writes each batch of outputs concurrently. Can pose issues with file systems that do not support memory mapping (such as some distributed file systems.)
+When False, all outputs are stored in RAM and written to disk at once after computation has finished. Default is `True`.
+- **`--precision`** (str, optional): Precision of the output data. Choose from `float16`, `16`, `half`, `float32`, `32`, `full`. Inference during embedding is not affected. Default is `float32`.
+- **`--flatten`** (bool, optional): Flatten 2D output arrays (per_token embeddings or attention weights) to 1D arrays per input sequence. Default is `False`.
+
+### 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`.

pepe_cli-1.0.0/README.md

@@ -0,0 +1,97 @@
+# PEPE
+
+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. 
+
+## Quick start
+
+1. Install PEPE \
+    From PyPI:    
+    ```sh
+    pip install pepe-cli
+    ```
+    Or install from the GitHub repository:    
+    ```sh
+    git clone https://github.com/csi-greifflab/pepe-cli
+    cd pepe-cli
+    pip install .
+    ```
+2. Run the embedding script:\
+    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>
+    ```
+
+## List of supported models:
+- ESM-family models
+    - ESM1:
+        - esm1_t34_670M_UR50S
+        - esm1_t34_670M_UR50D
+        - esm1_t34_670M_UR100
+        - esm1_t12_85M_UR50S
+        - esm1_t6_43M_UR50S
+        - esm1b_t33_650M_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
+    - ESM2:
+        - esm2_t6_8M_UR50D
+        - esm2_t12_35M_UR50D
+        - esm2_t30_150M_UR50D
+        - esm2_t33_650M_UR50D
+        - esm2_t36_3B_UR50D
+        - esm2_t48_15B_UR50D
+- Huggingface Transformer models
+    - T5 transformer models
+        - Rostlab/prot_t5_xl_half_uniref50-enc
+        - Rostlab/ProstT5
+    - RoFormer models
+        - alchemab/antiberta2-cssp
+        - alchemab/antiberta2
+    - Custom Hugging Face models
+        - Any compatible model from Hugging Face Hub: `username/model-name`
+        - Private models with authentication
+        - Local Hugging Face models
+- Custom Models
+    - Load your own PyTorch models with custom tokenizers
+    - Create example with: `python examples/custom_model/create_example_custom_model.py`
+
+
+## Arguments
+
+### 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`
+  - 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/`
+- **`--fasta_path`** (str): Path to the input FASTA file. If no experiment name is provided, the output files will be named after the input file.
+- **`--output_path`** (str): Directory for output files. Will generate a subdirectory for outputs of each output type.
+
+### Model Configuration
+- **`--tokenizer_from`** (str, optional): Huggingface address of the tokenizer to use. If not provided, will attempt to search for tokenizer packaged with model. If using a custom model, provide the path to the tokenizer directory.
+- **`--disable_special_tokens`** (bool, optional): When True, PEPE disables pre- and appending BOS/CLS and EOS/SEP tokens before embedding. Default is `False`.
+- **`--device`** (str, optional): Device to run the model on. Choose from `cuda` or `cpu`. Default is `cuda`.
+
+### Embedding Configuration
+- **`--layers`** (str, optional): Representation layers to extract from the model. Default is the last layer. Example: `--layers -1 6`.
+- **`--extract_embeddings`** (str, optional): 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 `mean_pooled`.
+- **`--substring_path`** (str, optional): Path to a CSV file with columns "sequence_id" and "substring". Only required when selecting "substring_pooled" option.
+- **`--context`** (int, optional): Only specify when including "substring_pooled" in `--extract_embeddings` option. Number of amino acids to include before and after the substring sequence. Default is `0`.
+
+### Processing Configuration
+- **`--batch_size`** (int, optional): Batch size for loading sequences. Default is `1024`. Decrease if encountering out-of-memory errors.
+- **`--max_length`** (int, optional): Length to which sequences will be padded. Default is length of longest sequence in input file. If shorter than longest sequence, will forcefully default to length of longest sequence.
+- **`--discard_padding`** (bool, optional): Discard padding tokens from per_token embeddings output. Default is `False`.
+
+### Output Configuration
+- **`--experiment_name`** (str, optional): Prefix for names of output files. If not provided, name of input file will be used for prefix.
+- **`--streaming_output`** (bool, optional): PEPE preallocates the required disk space and writes each batch of outputs concurrently. Can pose issues with file systems that do not support memory mapping (such as some distributed file systems.)
+When False, all outputs are stored in RAM and written to disk at once after computation has finished. Default is `True`.
+- **`--precision`** (str, optional): Precision of the output data. Choose from `float16`, `16`, `half`, `float32`, `32`, `full`. Inference during embedding is not affected. Default is `float32`.
+- **`--flatten`** (bool, optional): Flatten 2D output arrays (per_token embeddings or attention weights) to 1D arrays per input sequence. Default is `False`.
+
+### 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`.

pepe_cli-1.0.0/examples/custom_model/create_example_custom_model.py

@@ -0,0 +1,374 @@
+#!/usr/bin/env python3
+"""
+Example script showing how to create and use a custom model with EmbedAIRR.
+"""
+
+import torch
+import torch.nn as nn
+import os
+import json
+from transformers import AutoTokenizer
+import numpy as np
+
+
+class ExampleProteinModel(nn.Module):
+    """
+    Example protein model for demonstration.
+    This is a simple transformer-based model for protein sequences.
+    """
+
+    def __init__(
+        self, vocab_size=25, hidden_size=768, num_layers=6, num_heads=12, max_length=512
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+
+        # Embedding layers
+        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
+        self.position_embedding = nn.Embedding(max_length, hidden_size)
+
+        # Transformer layers
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_size,
+            nhead=num_heads,
+            dim_feedforward=hidden_size * 4,
+            dropout=0.1,
+            batch_first=True,
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
+
+        # Output head
+        self.output_head = nn.Linear(hidden_size, vocab_size)
+
+        # Layer normalization
+        self.layer_norm = nn.LayerNorm(hidden_size)
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        output_hidden_states=False,
+        output_attentions=False,
+        return_dict=True,
+    ):
+        """
+        Forward pass compatible with EmbedAIRR custom embedder.
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Create position ids
+        position_ids = (
+            torch.arange(seq_len, device=input_ids.device)
+            .unsqueeze(0)
+            .expand(batch_size, -1)
+        )
+
+        # Embeddings
+        token_embeddings = self.token_embedding(input_ids)
+        position_embeddings = self.position_embedding(position_ids)
+        hidden_states = self.layer_norm(token_embeddings + position_embeddings)
+
+        # Store hidden states if requested
+        all_hidden_states = []
+        if output_hidden_states:
+            all_hidden_states.append(hidden_states)
+
+        # Apply transformer layers
+        for layer in self.transformer.layers:
+            hidden_states = layer(
+                hidden_states,
+                src_key_padding_mask=(
+                    ~attention_mask.bool() if attention_mask is not None else None
+                ),
+            )
+            if output_hidden_states:
+                all_hidden_states.append(hidden_states)
+
+        # Output logits
+        logits = self.output_head(hidden_states)
+
+        # Create output object
+        class ModelOutput:
+            def __init__(self):
+                self.hidden_states = all_hidden_states if output_hidden_states else None
+                self.logits = logits
+                self.attentions = None  # Not implemented in this example
+
+        return ModelOutput()
+
+
+def create_example_model_and_tokenizer(
+    save_path="./examples/custom_model/example_protein_model",
+):
+    """
+    Create and save an example protein model and tokenizer.
+    """
+    # Create model
+    model = ExampleProteinModel(
+        vocab_size=25,
+        hidden_size=384,  # Smaller for example
+        num_layers=6,
+        num_heads=12,
+        max_length=512,
+    )
+
+    # Initialize weights (in practice, you would train the model)
+    for param in model.parameters():
+        if param.dim() > 1:
+            nn.init.xavier_uniform_(param)
+
+    # Create save directory
+    os.makedirs(save_path, exist_ok=True)
+
+    # Save model
+    torch.save(
+        {
+            "model": model.state_dict(),
+            "config": {
+                "num_layers": 6,
+                "num_attention_heads": 12,
+                "hidden_size": 384,
+                "vocab_size": 25,
+                "max_position_embeddings": 512,
+                "model_type": "protein_transformer",
+            },
+        },
+        os.path.join(save_path, "pytorch_model.pt"),
+    )
+
+    # Save config
+    config = {
+        "num_layers": 6,
+        "num_attention_heads": 12,
+        "hidden_size": 384,
+        "vocab_size": 25,
+        "max_position_embeddings": 512,
+        "model_type": "protein_transformer",
+    }
+
+    with open(os.path.join(save_path, "config.json"), "w") as f:
+        json.dump(config, f, indent=2)
+
+    # Create tokenizer files
+    _create_tokenizer_files(save_path)
+
+    print(f"Example model saved to: {save_path}")
+    print(f"Model config: {config}")
+
+    return save_path
+
+
+def _create_tokenizer_files(save_path):
+    """
+    Create tokenizer files compatible with AutoTokenizer.from_pretrained().
+    """
+    # Define protein vocabulary
+    vocab = {
+        "<pad>": 0,
+        "<cls>": 1,
+        "<sep>": 2,
+        "<unk>": 3,
+        "A": 4,
+        "R": 5,
+        "N": 6,
+        "D": 7,
+        "C": 8,
+        "Q": 9,
+        "E": 10,
+        "G": 11,
+        "H": 12,
+        "I": 13,
+        "L": 14,
+        "K": 15,
+        "M": 16,
+        "F": 17,
+        "P": 18,
+        "S": 19,
+        "T": 20,
+        "W": 21,
+        "Y": 22,
+        "V": 23,
+        "X": 24,  # Unknown amino acid
+    }
+
+    # Create tokenizer_config.json
+    tokenizer_config = {
+        "tokenizer_class": "BertTokenizer",
+        "model_max_length": 512,
+        "do_lower_case": False,
+        "cls_token": "<cls>",
+        "sep_token": "<sep>",
+        "pad_token": "<pad>",
+        "unk_token": "<unk>",
+        "model_type": "protein_transformer",
+    }
+
+    with open(os.path.join(save_path, "tokenizer_config.json"), "w") as f:
+        json.dump(tokenizer_config, f, indent=2)
+
+    # Create vocab.json
+    with open(os.path.join(save_path, "vocab.json"), "w") as f:
+        json.dump(vocab, f, indent=2)
+
+    # Create vocab.txt (for BertTokenizer compatibility)
+    vocab_txt_content = []
+    for token, id in sorted(vocab.items(), key=lambda x: x[1]):
+        vocab_txt_content.append(token)
+
+    with open(os.path.join(save_path, "vocab.txt"), "w") as f:
+        f.write("\n".join(vocab_txt_content))
+
+    # Create special_tokens_map.json
+    special_tokens_map = {
+        "cls_token": "<cls>",
+        "sep_token": "<sep>",
+        "pad_token": "<pad>",
+        "unk_token": "<unk>",
+    }
+
+    with open(os.path.join(save_path, "special_tokens_map.json"), "w") as f:
+        json.dump(special_tokens_map, f, indent=2)
+
+    print(f"✓ Tokenizer files created in {save_path}")
+    print("  - tokenizer_config.json")
+    print("  - vocab.json")
+    print("  - vocab.txt")
+    print("  - special_tokens_map.json")
+
+
+def create_example_fasta(filename="./examples/custom_model/example_sequences.fasta"):
+    """
+    Create an example FASTA file with protein sequences.
+    """
+    sequences = [
+        (
+            "seq1",
+            "MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSLEVGQIA",
+        ),
+        (
+            "seq2",
+            "MRIILLGAPGAGKGTQAQFIMEKYGIPQISTGDMLRAAVKSGSELGKQAKDIMDAGKLVTDELVIALVKERIAQEDCRNGFLLDGFPRTIPQADAMKEAGINVDYVLEFDVPDELIVDRIVGRRVHAPSGRVYHVKFNPPKVEGKDDVTGEELTTRKDDQEETVRKRLVEYHQMTAPLIGYYSKEAEAGNTKYAKVDGTKPVAEVRADLEKILG",
+        ),
+        (
+            "seq3",
+            "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG",
+        ),
+        (
+            "seq4",
+            "MHHHHHHSSGVDLGTENLYFQSMKKWVTFISLLFLFSSAYSRGVFRRDAHKSEVAHRFKDLGEQHFKGLVLIAFSQYLQQCPFDEHVKLVNELTEFAKTCVADESHAGCEKSLHTLFGDELCKVASLRETARDLLHIQEKLLESYIFQGPRDTHYSQKFPQTLKLLDRKQLAGDLSSQFQEAFEQATRQRQAKVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQ",
+        ),
+        ("seq5", "ARNDCEQGHILKMFPSTWYV"),  # All amino acids
+    ]
+
+    with open(filename, "w") as f:
+        for seq_id, sequence in sequences:
+            f.write(f">{seq_id}\n{sequence}\n")
+
+    print(f"Example FASTA file created: {filename}")
+    return filename
+
+
+def create_example_substring(filename="./examples/custom_model/example_substring.csv"):
+    """
+    Create an example substring file.
+    """
+    substring_data = [
+        (
+            "seq1",
+            "AKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGDGTQDNLSGAEKAVQVKVKALPDAQFEVVHSLAKWKRQTLGQHDFSAGEGLYTHMKALRPDEDRLSLEVGQIA",
+        ),
+        (
+            "seq2",
+            "DGKLVTDELVIALVKERIAQEDCRNGFLLDGFPRTIPQADAMKEAGINVDYVLEFDVPDELIVDRIVGRRVHAPSGRVYHVKFNPPKVEGKDDVTGEELTTRKDDQEETVRKRLVEYHQMTAPLIGYYSKEAEAGNTKYAKVDGTKPVAEVRADLEKILG",
+        ),
+        ("seq3", "DKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"),
+        (
+            "seq4",
+            "TCVADESHAGCEKSLHTLFGDELCKVASLRETARDLLHIQEKLLESYIFQGPRDTHYSQKFPQTLKLLDRKQLAGDLSSQFQEAFEQATRQRQAKVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQVQWLKRQFRQAGNTDVQRQFQFSQIRLQAGRQAEFQWLQRQFQARAGQIEQTQRQAQ",
+        ),
+        ("seq5", "ARNDCEQGHILKMFPSTWYV"),
+    ]
+
+    with open(filename, "w") as f:
+        f.write("sequence_id,substring_aa\n")
+        for seq_id, substring in substring_data:
+            f.write(f"{seq_id},{substring}\n")
+
+    print(f"Example substring file created: {filename}")
+    return filename
+
+
+def main():
+    """
+    Main function to create example files and demonstrate usage.
+    """
+    print("Creating example custom protein model...")
+
+    # Create example model
+    model_path = create_example_model_and_tokenizer()
+
+    # Create example data
+    fasta_file = create_example_fasta()
+    substring_file = create_example_substring()
+
+    # Test loading the model
+    print("\nTesting model loading...")
+    try:
+        model_data = torch.load(
+            os.path.join(model_path, "pytorch_model.pt"), map_location="cpu"
+        )
+        print(f"✓ Model loaded successfully")
+        print(f"✓ Model config: {model_data.get('config', {})}")
+
+        # Test creating the model
+        model = ExampleProteinModel(
+            vocab_size=25, hidden_size=384, num_layers=6, num_heads=12, max_length=512
+        )
+        model.load_state_dict(model_data["model"])
+        print(f"✓ Model architecture created and weights loaded")
+
+        # Test tokenizer loading
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_path)
+            print(
+                f"✓ Tokenizer loaded successfully with AutoTokenizer.from_pretrained()"
+            )
+            print(f"✓ Tokenizer vocab size: {len(tokenizer.get_vocab())}")
+
+            # Test tokenization
+            test_sequence = "ARNDCEQGHILKMFPSTWYV"
+            tokens = tokenizer(test_sequence, return_tensors="pt")
+            print(f"✓ Tokenization test successful")
+            print(f"✓ Input sequence: {test_sequence}")
+            print(f"✓ Tokenized shape: {tokens['input_ids'].shape}")
+
+        except Exception as e:
+            print(f"✗ Error testing tokenizer: {e}")
+
+        # Test forward pass
+        dummy_input = torch.randint(0, 25, (1, 10))
+        dummy_mask = torch.ones(1, 10)
+
+        with torch.no_grad():
+            output = model(
+                dummy_input, attention_mask=dummy_mask, output_hidden_states=True
+            )
+
+        print(f"✓ Forward pass successful")
+        print(f"✓ Output logits shape: {output.logits.shape}")
+        print(
+            f"✓ Number of hidden states: {len(output.hidden_states) if output.hidden_states else 0}"
+        )
+
+    except Exception as e:
+        print(f"✗ Error testing model: {e}")
+        return False
+
+    print("\nAll tests passed! The custom model is ready to use with EmbedAIRR.")
+    return True
+
+
+if __name__ == "__main__":
+    main()

pepe_cli-1.0.0/examples/custom_model/example_protein_model/config.json

@@ -0,0 +1,8 @@
+{
+  "num_layers": 6,
+  "num_attention_heads": 12,
+  "hidden_size": 384,
+  "vocab_size": 25,
+  "max_position_embeddings": 512,
+  "model_type": "protein_transformer"
+}

pepe_cli-1.0.0/examples/custom_model/example_protein_model/special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "cls_token": "<cls>",
+  "sep_token": "<sep>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

pepe_cli-1.0.0/examples/custom_model/example_protein_model/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "tokenizer_class": "BertTokenizer",
+  "model_max_length": 512,
+  "do_lower_case": false,
+  "cls_token": "<cls>",
+  "sep_token": "<sep>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>",
+  "model_type": "protein_transformer"
+}