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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pepe-cli
-Version: 1.0.4
+Version: 1.1
 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
@@ -27,11 +27,9 @@ 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
@@ -49,6 +47,10 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
     ```sh
     pip install pepe-cli
     ```
+    From Conda:
+    ```sh
+    conda install -c jahn_zhong pepe-cli
+    ```
     Or install from the GitHub repository:    
     ```sh
     git clone https://github.com/csi-greifflab/pepe-cli
@@ -61,6 +63,135 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
     pepe --experiment_name <optional_string> --fasta_path <file_path> --output_path <directory> --model_name <model_name>
     ```
 
+## Quick start with Library
+
+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).
+2. Use the `pepe.embed()` function in your script:
+
+```python
+import pepe
+
+# Example: Embed sequences from a dictionary
+sequences = {
+    "prot1": "MADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
+    "prot2": "MERIKELRDLMSQSRTREILTKLAEAGIDVPRLFK"
+}
+
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="my_embeddings",
+    extract_embeddings=["mean_pooled"],
+    device="cpu"  # Use "cuda", "cuda:0", "cuda:1", etc. if available
+)
+
+# Or from a FASTA file
+pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    fasta_path="path/to/fasta",
+    output_path="my_embeddings",
+    extract_embeddings=["mean_pooled"],
+    device="cpu"  # Use "cuda" if available
+)
+```
+
+### Advanced Usage
+
+For more control, you can use the embedder classes directly:
+
+```python
+from pepe.model_selecter import select_model
+
+# Select the appropriate model class
+# This acts as a factory returning the correct subclass (ESMEmbedder, HuggingfaceEmbedder, etc.)
+ModelClass = select_model("esm2_t6_8M_UR50D")
+
+# Initialize the embedder
+embedder = ModelClass(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    fasta_path="path/to/sequences.fasta",
+    output_path="output_directory",
+    extract_embeddings=["mean_pooled", "attention_head"],
+    layers=[[-1], [6]]  # Layers are expected as a list of lists of ints
+)
+
+# Run the embedding pipeline
+embedder.run()
+```
+
+### Memory Management & Large Scale Processing
+
+PEPE is designed to handle protein datasets of any size by utilizing **streamed outputs** and **memory mapping**. This feature is enabled by default (`streaming_output=True`).
+
+- **CLI Usage**: No action needed. PEPE automatically streams batches to disk to avoid OOM errors.
+- **Library Usage**: When using `pepe.embed()` or the embedder classes, the returned object does **not** hold the full embeddings in RAM. Instead, it provides `numpy.memmap` handles to the data on disk.
+
+```python
+import pepe
+
+# Returns an embedder object
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="large_dataset_output"
+    # streaming_output=True  <-- Default
+)
+
+# The embeddings are NOT loaded into RAM here.
+# 'data' is a numpy.memmap object pointing to the file on disk.
+data = results.mean_pooled["output_data"][-1] 
+
+# You can slice it like a normal array, which only loads those specific rows into RAM
+first_100_embeddings = data[:100]
+
+# Optimizing RAM usage:
+# If you are done with the model but want to keep working with the data,
+# you can delete the embedder object to free up GPU/CPU memory while keeping the memmaps.
+del results 
+```
+
+### Handling Long Sequences (Splitting & Reconstruction)
+
+Some models have strict architectural limits on input length (e.g., 1024 for ESM-2, 256 for AntiBERTa2). PEPE can automatically detect sequences that exceed these limits and handle them through chunking and reconstruction.
+
+- **Automatic Detection**: When `--split_long_sequences` is enabled, PEPE automatically identifies sequences exceeding the model's capacity.
+- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks. 
+- **Reconstruction**: 
+    - In **Library mode**, sequences are reconstructed in memory automatically after `embed()`.
+    - In **CLI mode**, sequences are reconstructed if `streaming_output=False`. If `streaming_output=True`, chunks are exported individually to maximize efficiency and minimize RAM usage.
+
+```python
+# Library Example: Process a 2000 AA protein with ESM-2 (1024 limit)
+results = pepe.embed(
+    model_name="facebook/esm2_t33_650M_UR50D",
+    sequences={"long_prot": "M" * 2000},
+    split_long_sequences=True,
+    split_overlap=50
+)
+
+# 'results.per_token' will contain a single reconstructed tensor of length ~2002
+# (including special tokens) despite the model's 1024 limit.
+```
+
+
+#### Performance Optimization
+
+If you have sufficient RAM to hold the entire dataset in memory, you can disable streaming output for faster execution. This avoids the overhead of writing to disk during the embedding process.
+
+```python
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="output",
+    streaming_output=False  # Store everything in valid RAM for speed
+)
+
+# Now 'results.mean_pooled["output_data"]' is a standard Numpy object
+# accessible immediately in memory.
+```
+
 ## List of supported models:
 - ESM-family models
     - ESM1:
@@ -112,23 +243,35 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
 ### 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`.
+- **`--device`** (str, optional): Device to run the model on. Choose from `cuda`, `cpu`, or specific GPU indices like `cuda:0`, `cuda:1`. 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`.
+- **`--extract_embeddings`** (str, optional): Set the embedding return types. Choose one or more from:
+  - `per_token`: Extracts embeddings for each token (amino acid) in the sequence. Output shape: `(num_sequences, max_length, embedding_size)`.
+  - `mean_pooled`: Computes the average embedding across all tokens in a sequence, excluding special tokens (BOS, EOS, padding). Output shape: `(num_sequences, embedding_size)`.
+  - `substring_pooled`: Computes the average embedding for a specific substring within each sequence (e.g., a CDR3 region). Requires `--substring_path`. Output shape: `(num_sequences, embedding_size)`.
+  - `attention_head`: Extracts raw attention weights for every individual head in the specified layers. Output shape: `(num_sequences, max_length, max_length)` per head.
+  - `attention_layer`: Extracts the average attention weights across all heads within each specified layer. Output shape: `(num_sequences, max_length, max_length)` per layer.
+  - `attention_model`: Extracts the average attention weights across all heads and all specified layers. Output shape: `(num_sequences, max_length, max_length)`.
+  - `logits`: Extracts the raw language model output (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`.
+- **`--max_length`** (int, optional): Length to which sequences will be padded. Default is length of longest sequence in input file + special token(s). If shorter than longest sequence, will forcefully default to length of longest sequence + special token(s).
+- **`--split_long_sequences`** (bool, optional): When True, automatically detect sequences exceeding the model's maximum allowed length and split them into chunks for processing. Default is `False`.
+- **`--split_overlap`** (int, optional): Number of tokens to overlap when splitting long sequences. This helps maintain context across chunk boundaries. Default is `0`.
+- **`--force_split_length`** (int, optional): Explicitly force sequence splitting at a specific length, overriding the model's auto-detected limits. Default is `None`.
+- **`--discard_padding`** (bool, optional): Discard padding tokens from per_token embeddings output. **Note**: Setting this to `True` will automatically disable `--streaming_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`.
+When False, all outputs are stored in RAM and written to disk at once after computation has finished. **Note**: Automatically disabled if `--discard_padding` is `True`. 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`.
 

pepe_cli-1.1/README.md

@@ -0,0 +1,242 @@
+# 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
+    ```
+    From Conda:
+    ```sh
+    conda install -c jahn_zhong 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>
+    ```
+
+## Quick start with Library
+
+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).
+2. Use the `pepe.embed()` function in your script:
+
+```python
+import pepe
+
+# Example: Embed sequences from a dictionary
+sequences = {
+    "prot1": "MADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK",
+    "prot2": "MERIKELRDLMSQSRTREILTKLAEAGIDVPRLFK"
+}
+
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="my_embeddings",
+    extract_embeddings=["mean_pooled"],
+    device="cpu"  # Use "cuda", "cuda:0", "cuda:1", etc. if available
+)
+
+# Or from a FASTA file
+pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    fasta_path="path/to/fasta",
+    output_path="my_embeddings",
+    extract_embeddings=["mean_pooled"],
+    device="cpu"  # Use "cuda" if available
+)
+```
+
+### Advanced Usage
+
+For more control, you can use the embedder classes directly:
+
+```python
+from pepe.model_selecter import select_model
+
+# Select the appropriate model class
+# This acts as a factory returning the correct subclass (ESMEmbedder, HuggingfaceEmbedder, etc.)
+ModelClass = select_model("esm2_t6_8M_UR50D")
+
+# Initialize the embedder
+embedder = ModelClass(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    fasta_path="path/to/sequences.fasta",
+    output_path="output_directory",
+    extract_embeddings=["mean_pooled", "attention_head"],
+    layers=[[-1], [6]]  # Layers are expected as a list of lists of ints
+)
+
+# Run the embedding pipeline
+embedder.run()
+```
+
+### Memory Management & Large Scale Processing
+
+PEPE is designed to handle protein datasets of any size by utilizing **streamed outputs** and **memory mapping**. This feature is enabled by default (`streaming_output=True`).
+
+- **CLI Usage**: No action needed. PEPE automatically streams batches to disk to avoid OOM errors.
+- **Library Usage**: When using `pepe.embed()` or the embedder classes, the returned object does **not** hold the full embeddings in RAM. Instead, it provides `numpy.memmap` handles to the data on disk.
+
+```python
+import pepe
+
+# Returns an embedder object
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="large_dataset_output"
+    # streaming_output=True  <-- Default
+)
+
+# The embeddings are NOT loaded into RAM here.
+# 'data' is a numpy.memmap object pointing to the file on disk.
+data = results.mean_pooled["output_data"][-1] 
+
+# You can slice it like a normal array, which only loads those specific rows into RAM
+first_100_embeddings = data[:100]
+
+# Optimizing RAM usage:
+# If you are done with the model but want to keep working with the data,
+# you can delete the embedder object to free up GPU/CPU memory while keeping the memmaps.
+del results 
+```
+
+### Handling Long Sequences (Splitting & Reconstruction)
+
+Some models have strict architectural limits on input length (e.g., 1024 for ESM-2, 256 for AntiBERTa2). PEPE can automatically detect sequences that exceed these limits and handle them through chunking and reconstruction.
+
+- **Automatic Detection**: When `--split_long_sequences` is enabled, PEPE automatically identifies sequences exceeding the model's capacity.
+- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks. 
+- **Reconstruction**: 
+    - In **Library mode**, sequences are reconstructed in memory automatically after `embed()`.
+    - In **CLI mode**, sequences are reconstructed if `streaming_output=False`. If `streaming_output=True`, chunks are exported individually to maximize efficiency and minimize RAM usage.
+
+```python
+# Library Example: Process a 2000 AA protein with ESM-2 (1024 limit)
+results = pepe.embed(
+    model_name="facebook/esm2_t33_650M_UR50D",
+    sequences={"long_prot": "M" * 2000},
+    split_long_sequences=True,
+    split_overlap=50
+)
+
+# 'results.per_token' will contain a single reconstructed tensor of length ~2002
+# (including special tokens) despite the model's 1024 limit.
+```
+
+
+#### Performance Optimization
+
+If you have sufficient RAM to hold the entire dataset in memory, you can disable streaming output for faster execution. This avoids the overhead of writing to disk during the embedding process.
+
+```python
+results = pepe.embed(
+    model_name="facebook/esm2_t6_8M_UR50D",
+    sequences=sequences,
+    output_path="output",
+    streaming_output=False  # Store everything in valid RAM for speed
+)
+
+# Now 'results.mean_pooled["output_data"]' is a standard Numpy object
+# accessible immediately in memory.
+```
+
+## 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`, `cpu`, or specific GPU indices like `cuda:0`, `cuda:1`. 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`: Extracts embeddings for each token (amino acid) in the sequence. Output shape: `(num_sequences, max_length, embedding_size)`.
+  - `mean_pooled`: Computes the average embedding across all tokens in a sequence, excluding special tokens (BOS, EOS, padding). Output shape: `(num_sequences, embedding_size)`.
+  - `substring_pooled`: Computes the average embedding for a specific substring within each sequence (e.g., a CDR3 region). Requires `--substring_path`. Output shape: `(num_sequences, embedding_size)`.
+  - `attention_head`: Extracts raw attention weights for every individual head in the specified layers. Output shape: `(num_sequences, max_length, max_length)` per head.
+  - `attention_layer`: Extracts the average attention weights across all heads within each specified layer. Output shape: `(num_sequences, max_length, max_length)` per layer.
+  - `attention_model`: Extracts the average attention weights across all heads and all specified layers. Output shape: `(num_sequences, max_length, max_length)`.
+  - `logits`: Extracts the raw language model output (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 + special token(s). If shorter than longest sequence, will forcefully default to length of longest sequence + special token(s).
+- **`--split_long_sequences`** (bool, optional): When True, automatically detect sequences exceeding the model's maximum allowed length and split them into chunks for processing. Default is `False`.
+- **`--split_overlap`** (int, optional): Number of tokens to overlap when splitting long sequences. This helps maintain context across chunk boundaries. Default is `0`.
+- **`--force_split_length`** (int, optional): Explicitly force sequence splitting at a specific length, overriding the model's auto-detected limits. Default is `None`.
+- **`--discard_padding`** (bool, optional): Discard padding tokens from per_token embeddings output. **Note**: Setting this to `True` will automatically disable `--streaming_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. **Note**: Automatically disabled if `--discard_padding` is `True`. 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.4 → pepe_cli-1.1}/pyproject.toml

@@ -29,11 +29,9 @@ keywords = ["protein", "embeddings", "bioinformatics", "machine-learning", "nlp"
 dependencies = [
     "torch>=1.9.0",
     "transformers>=4.20.0",
-    "fair-esm",
     "sentencepiece",
     "numpy",
     "protobuf",
-    "rjieba",
     "alive_progress",
 ]
 

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

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

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

@@ -43,11 +43,9 @@ setup(
     install_requires=[
         "torch>=1.9.0",
         "transformers>=4.20.0",
-        "fair-esm",
         "sentencepiece",
         "numpy",
         "protobuf",
-        "rjieba",
         "alive_progress",
     ],
     entry_points={

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

@@ -1,8 +1,13 @@
 import logging
 import sys
 
+try:
+    from pepe.api import embed  # Exported to top-level; requires torch at runtime
+except ImportError:
+    pass
+
 # Package metadata - single source of truth
-__version__ = "1.0.4"
+__version__ = "1.1"
 __package_name__ = "pepe-cli"
 __module_name__ = "pepe"
 __author__ = "Jahn Zhong"

pepe_cli-1.1/src/pepe/api.py

@@ -0,0 +1,118 @@
+import os
+import tempfile
+import logging
+from types import SimpleNamespace
+from typing import Dict, List, Optional, Union, Any
+
+from pepe.model_selecter import select_model
+import pepe.utils
+
+logger = logging.getLogger("pepe.api")
+
+def embed(
+    model_name: str,
+    sequences: Optional[Union[Dict[str, str], List[str]]] = None,
+    fasta_path: Optional[str] = None,
+    output_path: Optional[str] = None,
+    extract_embeddings: List[str] = ["mean_pooled"],
+    layers: Union[List[int], List[List[int]]] = [[-1]],
+    batch_size: int = 1024,
+    device: str = "cuda",
+    precision: str = "32",
+    streaming_output: bool = True,
+    discard_padding: bool = False,
+    max_input_length: str = "max_length",
+    experiment_name: Optional[str] = None,
+    **kwargs
+) -> Dict[str, Any]:
+    """
+    High-level API for generating protein embeddings.
+
+    Args:
+        model_name: Model name or path to custom model.
+        sequences: Dictionary mapping labels to protein sequences, or a list of sequences.
+        fasta_path: Path to a FASTA file. Used if 'sequences' is not provided.
+        output_path: Directory for output files. If not provided, a temporary directory will be used.
+        extract_embeddings: List of embedding types to extract.
+        layers: Representation layers to extract. Default is the last layer ([-1]).
+        batch_size: Number of tokens per batch. Default is 1024.
+        device: Device to run the model on ('cuda' or 'cpu'). Default is 'cuda'.
+        precision: Output precision (e.g., '32', '16'). Default is '32'.
+        streaming_output: Whether to stream outputs to disk. Default is True.
+        discard_padding: Whether to discard padding tokens. Default is False.
+        max_input_length: Length to which sequences will be padded. Default is "max_length".
+        experiment_name: Optional prefix for output files.
+        **kwargs: Additional arguments supported by the embedders.
+
+    Returns:
+        A dictionary containing the results and/or the output path.
+    """
+    # Create a temporary fasta file if sequences are provided
+    temp_fasta = None
+    if sequences is not None:
+        if isinstance(sequences, list):
+            sequences = {f"seq_{i}": seq for i, seq in enumerate(sequences)}
+            
+        temp_fasta = tempfile.NamedTemporaryFile(mode='w', suffix='.fasta', delete=False)
+        for label, seq in sequences.items():
+            temp_fasta.write(f">{label}\n{seq}\n")
+        temp_fasta.close()
+        fasta_path = temp_fasta.name
+
+    if fasta_path is None:
+        raise ValueError("Either 'sequences' or 'fasta_path' must be provided.")
+
+    return_results = False
+    if output_path is None:
+        if streaming_output:
+            logger.warning("No output_path provided. Disabling streaming_output and returning in-memory results.")
+        output_path = tempfile.mkdtemp()
+        streaming_output = False
+        return_results = True
+    
+    # Create args object
+    args_dict = {
+        "model_name": model_name,
+        "fasta_path": fasta_path,
+        "output_path": output_path,
+        "extract_embeddings": extract_embeddings,
+        "layers": layers,
+        "batch_size": batch_size,
+        "device": device,
+        "precision": precision,
+        "streaming_output": streaming_output,
+        "discard_padding": discard_padding,
+        "max_input_length": max_input_length,
+        "experiment_name": experiment_name,
+        "tokenizer_from": kwargs.get("tokenizer_from"),
+        "substring_path": kwargs.get("substring_path"),
+        "context": kwargs.get("context", 0),
+        "split_long_sequences": kwargs.get("split_long_sequences", False),
+        "split_overlap": kwargs.get("split_overlap", 0),
+        "force_split_length": kwargs.get("force_split_length"),
+        "num_workers": kwargs.get("num_workers", 8),
+        "disable_special_tokens": kwargs.get("disable_special_tokens", False),
+        "flatten": kwargs.get("flatten", False),
+        "flush_batches_after": kwargs.get("flush_batches_after", 128),
+    }
+    
+    args = SimpleNamespace(**args_dict)
+    
+    selected_model_class = select_model(model_name)
+    embedder = selected_model_class(args)
+    embedder.run()
+    
+    results = {"output_path": output_path}
+    
+    if return_results:
+        # Pick up in-memory results before they are lost
+        for output_type in extract_embeddings:
+            obj = getattr(embedder, output_type, None)
+            if obj and "output_data" in obj:
+                results[output_type] = obj["output_data"]
+    
+    # Cleanup temp file
+    if temp_fasta:
+        os.unlink(temp_fasta.name)
+        
+    return results