pepe-cli 1.4.0__tar.gz → 1.5.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {pepe_cli-1.4.0/src/pepe_cli.egg-info → pepe_cli-1.5.0}/PKG-INFO +51 -13
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/README.md +40 -9
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/create_example_custom_model.py +11 -11
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/config.json +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/special_tokens_map.json +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/tokenizer_config.json +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/vocab.json +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/vocab.txt +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/embedding_options.md +3 -3
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/model_selection.md +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/pyproject.toml +32 -3
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/requirements.txt +1 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/setup.py +16 -6
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/__init__.py +4 -2
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/__main__.py +1 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/api.py +18 -14
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/embedders/base_embedder.py +70 -58
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/embedders/custom_embedder.py +12 -14
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/embedders/esm_embedder.py +4 -2
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/embedders/huggingface_embedder.py +28 -23
- pepe_cli-1.5.0/src/pepe/embedders/metl_embedder.py +349 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/model_errors.py +8 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/model_selecter.py +47 -16
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/utils.py +149 -41
- {pepe_cli-1.4.0 → pepe_cli-1.5.0/src/pepe_cli.egg-info}/PKG-INFO +51 -13
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/SOURCES.txt +4 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/requires.txt +11 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/conftest.py +3 -3
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_api_unittest.py +20 -18
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_custom_embedder.py +1 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_device_logic.py +32 -24
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_esmc_modes.py +4 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_generic_hf_integration.py +2 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_generic_safeguards.py +17 -8
- pepe_cli-1.5.0/src/tests/test_metl_dispatch.py +100 -0
- pepe_cli-1.5.0/src/tests/test_metl_hooks.py +178 -0
- pepe_cli-1.5.0/src/tests/test_metl_integration.py +97 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_model_selection.py +21 -18
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_reconstruct_mean_pooled.py +1 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_splitting.py +72 -50
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_streaming_roundtrip.py +25 -9
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_sync_arguments.py +4 -1
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/LICENSE +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/MANIFEST.in +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/pytorch_model.pt +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_sequences.fasta +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/examples/custom_model/example_substring.csv +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/setup.cfg +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/embedders/__init__.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/parse_arguments.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe/py.typed +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/dependency_links.txt +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/entry_points.txt +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/top_level.txt +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/__init__.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_esm1_integration.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_esmc.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_load_layers_default.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_parse_arguments.py +0 -0
- {pepe_cli-1.4.0 → pepe_cli-1.5.0}/src/tests/test_t5_antiberta2_integration.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: pepe-cli
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.5.0
|
|
4
4
|
Summary: Pipeline for Easy Protein Embedding - Extract embeddings and attention matrices from protein sequences
|
|
5
5
|
Home-page: https://github.com/csi-greifflab/pepe-cli
|
|
6
6
|
Author: Jahn Zhong
|
|
@@ -16,13 +16,11 @@ Classifier: Intended Audience :: Science/Research
|
|
|
16
16
|
Classifier: License :: OSI Approved :: MIT License
|
|
17
17
|
Classifier: Operating System :: OS Independent
|
|
18
18
|
Classifier: Programming Language :: Python :: 3
|
|
19
|
-
Classifier: Programming Language :: Python :: 3.8
|
|
20
|
-
Classifier: Programming Language :: Python :: 3.9
|
|
21
19
|
Classifier: Programming Language :: Python :: 3.10
|
|
22
20
|
Classifier: Programming Language :: Python :: 3.11
|
|
23
21
|
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
|
|
24
22
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
25
|
-
Requires-Python: >=3.
|
|
23
|
+
Requires-Python: >=3.10
|
|
26
24
|
Description-Content-Type: text/markdown
|
|
27
25
|
License-File: LICENSE
|
|
28
26
|
Requires-Dist: torch>=1.9.0
|
|
@@ -32,6 +30,15 @@ Requires-Dist: numpy
|
|
|
32
30
|
Requires-Dist: protobuf
|
|
33
31
|
Requires-Dist: alive_progress
|
|
34
32
|
Requires-Dist: rjieba
|
|
33
|
+
Provides-Extra: esm
|
|
34
|
+
Requires-Dist: fair-esm; extra == "esm"
|
|
35
|
+
Provides-Extra: dev
|
|
36
|
+
Requires-Dist: pre-commit; extra == "dev"
|
|
37
|
+
Requires-Dist: ruff; extra == "dev"
|
|
38
|
+
Requires-Dist: mypy; extra == "dev"
|
|
39
|
+
Requires-Dist: pytest; extra == "dev"
|
|
40
|
+
Requires-Dist: pytest-cov; extra == "dev"
|
|
41
|
+
Requires-Dist: biopython; extra == "dev"
|
|
35
42
|
Dynamic: author
|
|
36
43
|
Dynamic: home-page
|
|
37
44
|
Dynamic: license-file
|
|
@@ -39,19 +46,19 @@ Dynamic: requires-python
|
|
|
39
46
|
|
|
40
47
|
# PEPE
|
|
41
48
|
|
|
42
|
-
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.
|
|
49
|
+
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.
|
|
43
50
|
|
|
44
51
|
### Citation
|
|
45
52
|
> **PEPE: Scalable extraction of multi-modal protein language model representations**
|
|
46
53
|
> Jahn Zhong, Niccolò Cardente, Geir Kjetil Sandve, Habib Bashour, Maria Francesca Abbate, Victor Greiff
|
|
47
|
-
> *bioRxiv* (2026)
|
|
54
|
+
> *bioRxiv* (2026)
|
|
48
55
|
> [DOI: 10.1101/2025.10.13.680902](https://doi.org/10.1101/2025.10.13.680902)
|
|
49
56
|
|
|
50
57
|
## Quick start
|
|
51
58
|
|
|
52
59
|
1. Install PEPE
|
|
53
60
|
|
|
54
|
-
From PyPI:
|
|
61
|
+
From PyPI:
|
|
55
62
|
```sh
|
|
56
63
|
pip install pepe-cli
|
|
57
64
|
```
|
|
@@ -59,7 +66,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
|
|
|
59
66
|
```sh
|
|
60
67
|
conda install -c jahn_zhong pepe-cli
|
|
61
68
|
```
|
|
62
|
-
Or install from the GitHub repository:
|
|
69
|
+
Or install from the GitHub repository:
|
|
63
70
|
```sh
|
|
64
71
|
git clone https://github.com/csi-greifflab/pepe-cli
|
|
65
72
|
cd pepe-cli
|
|
@@ -71,7 +78,35 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
|
|
|
71
78
|
pip install git+https://github.com/Biohub/transformers.git@main
|
|
72
79
|
```
|
|
73
80
|
|
|
74
|
-
3.
|
|
81
|
+
3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the backend directly from GitHub (it is not on PyPI):
|
|
82
|
+
|
|
83
|
+
```sh
|
|
84
|
+
pip install git+https://github.com/gitter-lab/metl-pretrained.git
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
CLI example:
|
|
88
|
+
|
|
89
|
+
```sh
|
|
90
|
+
pepe --model_name metl-g-20m-1d --fasta_path <file_path> --output_path <directory> --extract_embeddings mean_pooled
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Library example:
|
|
94
|
+
|
|
95
|
+
```python
|
|
96
|
+
import pepe
|
|
97
|
+
|
|
98
|
+
results = pepe.embed(
|
|
99
|
+
model_name="metl-g-20m-1d",
|
|
100
|
+
sequences={"prot1": "MADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"},
|
|
101
|
+
output_path="my_embeddings",
|
|
102
|
+
extract_embeddings=["mean_pooled"],
|
|
103
|
+
device="cpu",
|
|
104
|
+
)
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
**Limitations:** 1D METL models only (`per_token`, `mean_pooled`, `substring_pooled`). Logits and attention outputs are not supported. 3D METL identifiers (requiring structures) and the generic Hugging Face repo id `gitter-lab/METL` are rejected—use a `metl-*-1d` identifier from [metl-pretrained](https://github.com/gitter-lab/metl-pretrained).
|
|
108
|
+
|
|
109
|
+
4. Extract embeddings:\
|
|
75
110
|
Extract mean pooled embeddings from protein amino acid sequences in FASTA file:
|
|
76
111
|
```sh
|
|
77
112
|
pepe --experiment_name <optional_string> --fasta_path <file_path> --output_path <directory> --model_name <model_name>
|
|
@@ -155,7 +190,7 @@ results = pepe.embed(
|
|
|
155
190
|
|
|
156
191
|
# The embeddings are NOT loaded into RAM here.
|
|
157
192
|
# 'data' is a numpy.memmap object pointing to the file on disk.
|
|
158
|
-
data = results.mean_pooled["output_data"][-1]
|
|
193
|
+
data = results.mean_pooled["output_data"][-1]
|
|
159
194
|
|
|
160
195
|
# You can slice it like a normal array, which only loads those specific rows into RAM
|
|
161
196
|
first_100_embeddings = data[:100]
|
|
@@ -163,7 +198,7 @@ first_100_embeddings = data[:100]
|
|
|
163
198
|
# Optimizing RAM usage:
|
|
164
199
|
# If you are done with the model but want to keep working with the data,
|
|
165
200
|
# you can delete the embedder object to free up GPU/CPU memory while keeping the memmaps.
|
|
166
|
-
del results
|
|
201
|
+
del results
|
|
167
202
|
```
|
|
168
203
|
|
|
169
204
|
### Handling Long Sequences (Splitting & Reconstruction)
|
|
@@ -171,8 +206,8 @@ del results
|
|
|
171
206
|
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.
|
|
172
207
|
|
|
173
208
|
- **Automatic Detection**: When `--split_long_sequences` is enabled, PEPE automatically identifies sequences exceeding the model's capacity.
|
|
174
|
-
- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks.
|
|
175
|
-
- **Reconstruction**:
|
|
209
|
+
- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks.
|
|
210
|
+
- **Reconstruction**:
|
|
176
211
|
- In **Library mode**, sequences are reconstructed in memory automatically after `embed()`.
|
|
177
212
|
- 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.
|
|
178
213
|
|
|
@@ -238,6 +273,8 @@ results = pepe.embed(
|
|
|
238
273
|
- biohub/ESMC-300M
|
|
239
274
|
- biohub/ESMC-600M
|
|
240
275
|
- biohub/ESMC-6B
|
|
276
|
+
- METL 1D models (requires `metl-pretrained` from GitHub; see Quick start)
|
|
277
|
+
- `metl-g-20m-1d` and other `metl-*-1d` identifiers supported by [metl-pretrained](https://github.com/gitter-lab/metl-pretrained)
|
|
241
278
|
- Custom Hugging Face models
|
|
242
279
|
- Any compatible model from Hugging Face Hub: `username/model-name`
|
|
243
280
|
- Private models with authentication
|
|
@@ -253,6 +290,7 @@ results = pepe.embed(
|
|
|
253
290
|
- **`--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:
|
|
254
291
|
- ESM models: `esm2_t33_650M_UR50D`
|
|
255
292
|
- ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start)
|
|
293
|
+
- METL 1D models: `metl-g-20m-1d` (requires `metl-pretrained` from GitHub; see Quick start)
|
|
256
294
|
- Hugging Face models: `username/model-name`
|
|
257
295
|
- Custom PyTorch models: `/path/to/model.pt` or `/path/to/model_directory/`
|
|
258
296
|
- Local HF models: `/path/to/local_hf_directory/`
|
|
@@ -1,18 +1,18 @@
|
|
|
1
1
|
# PEPE
|
|
2
2
|
|
|
3
|
-
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.
|
|
3
|
+
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.
|
|
4
4
|
|
|
5
5
|
### Citation
|
|
6
6
|
> **PEPE: Scalable extraction of multi-modal protein language model representations**
|
|
7
7
|
> Jahn Zhong, Niccolò Cardente, Geir Kjetil Sandve, Habib Bashour, Maria Francesca Abbate, Victor Greiff
|
|
8
|
-
> *bioRxiv* (2026)
|
|
8
|
+
> *bioRxiv* (2026)
|
|
9
9
|
> [DOI: 10.1101/2025.10.13.680902](https://doi.org/10.1101/2025.10.13.680902)
|
|
10
10
|
|
|
11
11
|
## Quick start
|
|
12
12
|
|
|
13
13
|
1. Install PEPE
|
|
14
14
|
|
|
15
|
-
From PyPI:
|
|
15
|
+
From PyPI:
|
|
16
16
|
```sh
|
|
17
17
|
pip install pepe-cli
|
|
18
18
|
```
|
|
@@ -20,7 +20,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
|
|
|
20
20
|
```sh
|
|
21
21
|
conda install -c jahn_zhong pepe-cli
|
|
22
22
|
```
|
|
23
|
-
Or install from the GitHub repository:
|
|
23
|
+
Or install from the GitHub repository:
|
|
24
24
|
```sh
|
|
25
25
|
git clone https://github.com/csi-greifflab/pepe-cli
|
|
26
26
|
cd pepe-cli
|
|
@@ -32,7 +32,35 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a
|
|
|
32
32
|
pip install git+https://github.com/Biohub/transformers.git@main
|
|
33
33
|
```
|
|
34
34
|
|
|
35
|
-
3.
|
|
35
|
+
3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the backend directly from GitHub (it is not on PyPI):
|
|
36
|
+
|
|
37
|
+
```sh
|
|
38
|
+
pip install git+https://github.com/gitter-lab/metl-pretrained.git
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
CLI example:
|
|
42
|
+
|
|
43
|
+
```sh
|
|
44
|
+
pepe --model_name metl-g-20m-1d --fasta_path <file_path> --output_path <directory> --extract_embeddings mean_pooled
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Library example:
|
|
48
|
+
|
|
49
|
+
```python
|
|
50
|
+
import pepe
|
|
51
|
+
|
|
52
|
+
results = pepe.embed(
|
|
53
|
+
model_name="metl-g-20m-1d",
|
|
54
|
+
sequences={"prot1": "MADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"},
|
|
55
|
+
output_path="my_embeddings",
|
|
56
|
+
extract_embeddings=["mean_pooled"],
|
|
57
|
+
device="cpu",
|
|
58
|
+
)
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
**Limitations:** 1D METL models only (`per_token`, `mean_pooled`, `substring_pooled`). Logits and attention outputs are not supported. 3D METL identifiers (requiring structures) and the generic Hugging Face repo id `gitter-lab/METL` are rejected—use a `metl-*-1d` identifier from [metl-pretrained](https://github.com/gitter-lab/metl-pretrained).
|
|
62
|
+
|
|
63
|
+
4. Extract embeddings:\
|
|
36
64
|
Extract mean pooled embeddings from protein amino acid sequences in FASTA file:
|
|
37
65
|
```sh
|
|
38
66
|
pepe --experiment_name <optional_string> --fasta_path <file_path> --output_path <directory> --model_name <model_name>
|
|
@@ -116,7 +144,7 @@ results = pepe.embed(
|
|
|
116
144
|
|
|
117
145
|
# The embeddings are NOT loaded into RAM here.
|
|
118
146
|
# 'data' is a numpy.memmap object pointing to the file on disk.
|
|
119
|
-
data = results.mean_pooled["output_data"][-1]
|
|
147
|
+
data = results.mean_pooled["output_data"][-1]
|
|
120
148
|
|
|
121
149
|
# You can slice it like a normal array, which only loads those specific rows into RAM
|
|
122
150
|
first_100_embeddings = data[:100]
|
|
@@ -124,7 +152,7 @@ first_100_embeddings = data[:100]
|
|
|
124
152
|
# Optimizing RAM usage:
|
|
125
153
|
# If you are done with the model but want to keep working with the data,
|
|
126
154
|
# you can delete the embedder object to free up GPU/CPU memory while keeping the memmaps.
|
|
127
|
-
del results
|
|
155
|
+
del results
|
|
128
156
|
```
|
|
129
157
|
|
|
130
158
|
### Handling Long Sequences (Splitting & Reconstruction)
|
|
@@ -132,8 +160,8 @@ del results
|
|
|
132
160
|
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.
|
|
133
161
|
|
|
134
162
|
- **Automatic Detection**: When `--split_long_sequences` is enabled, PEPE automatically identifies sequences exceeding the model's capacity.
|
|
135
|
-
- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks.
|
|
136
|
-
- **Reconstruction**:
|
|
163
|
+
- **Overlapping Chunks**: Use `--split_overlap` to maintain context between chunks.
|
|
164
|
+
- **Reconstruction**:
|
|
137
165
|
- In **Library mode**, sequences are reconstructed in memory automatically after `embed()`.
|
|
138
166
|
- 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.
|
|
139
167
|
|
|
@@ -199,6 +227,8 @@ results = pepe.embed(
|
|
|
199
227
|
- biohub/ESMC-300M
|
|
200
228
|
- biohub/ESMC-600M
|
|
201
229
|
- biohub/ESMC-6B
|
|
230
|
+
- METL 1D models (requires `metl-pretrained` from GitHub; see Quick start)
|
|
231
|
+
- `metl-g-20m-1d` and other `metl-*-1d` identifiers supported by [metl-pretrained](https://github.com/gitter-lab/metl-pretrained)
|
|
202
232
|
- Custom Hugging Face models
|
|
203
233
|
- Any compatible model from Hugging Face Hub: `username/model-name`
|
|
204
234
|
- Private models with authentication
|
|
@@ -214,6 +244,7 @@ results = pepe.embed(
|
|
|
214
244
|
- **`--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:
|
|
215
245
|
- ESM models: `esm2_t33_650M_UR50D`
|
|
216
246
|
- ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start)
|
|
247
|
+
- METL 1D models: `metl-g-20m-1d` (requires `metl-pretrained` from GitHub; see Quick start)
|
|
217
248
|
- Hugging Face models: `username/model-name`
|
|
218
249
|
- Custom PyTorch models: `/path/to/model.pt` or `/path/to/model_directory/`
|
|
219
250
|
- Local HF models: `/path/to/local_hf_directory/`
|
|
@@ -3,12 +3,12 @@
|
|
|
3
3
|
Example script showing how to create and use a custom model with EmbedAIRR.
|
|
4
4
|
"""
|
|
5
5
|
|
|
6
|
+
import json
|
|
7
|
+
import os
|
|
8
|
+
|
|
6
9
|
import torch
|
|
7
10
|
import torch.nn as nn
|
|
8
|
-
import os
|
|
9
|
-
import json
|
|
10
11
|
from transformers import AutoTokenizer
|
|
11
|
-
import numpy as np
|
|
12
12
|
|
|
13
13
|
|
|
14
14
|
class ExampleProteinModel(nn.Module):
|
|
@@ -309,9 +309,9 @@ def main():
|
|
|
309
309
|
# Create example model
|
|
310
310
|
model_path = create_example_model_and_tokenizer()
|
|
311
311
|
|
|
312
|
-
# Create example data
|
|
313
|
-
|
|
314
|
-
|
|
312
|
+
# Create example data (written to disk for the demo; return paths unused here)
|
|
313
|
+
create_example_fasta()
|
|
314
|
+
create_example_substring()
|
|
315
315
|
|
|
316
316
|
# Test loading the model
|
|
317
317
|
print("\nTesting model loading...")
|
|
@@ -319,7 +319,7 @@ def main():
|
|
|
319
319
|
model_data = torch.load(
|
|
320
320
|
os.path.join(model_path, "pytorch_model.pt"), map_location="cpu"
|
|
321
321
|
)
|
|
322
|
-
print(
|
|
322
|
+
print("✓ Model loaded successfully")
|
|
323
323
|
print(f"✓ Model config: {model_data.get('config', {})}")
|
|
324
324
|
|
|
325
325
|
# Test creating the model
|
|
@@ -327,20 +327,20 @@ def main():
|
|
|
327
327
|
vocab_size=25, hidden_size=384, num_layers=6, num_heads=12, max_length=512
|
|
328
328
|
)
|
|
329
329
|
model.load_state_dict(model_data["model"])
|
|
330
|
-
print(
|
|
330
|
+
print("✓ Model architecture created and weights loaded")
|
|
331
331
|
|
|
332
332
|
# Test tokenizer loading
|
|
333
333
|
try:
|
|
334
334
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
335
335
|
print(
|
|
336
|
-
|
|
336
|
+
"✓ Tokenizer loaded successfully with AutoTokenizer.from_pretrained()"
|
|
337
337
|
)
|
|
338
338
|
print(f"✓ Tokenizer vocab size: {len(tokenizer.get_vocab())}")
|
|
339
339
|
|
|
340
340
|
# Test tokenization
|
|
341
341
|
test_sequence = "ARNDCEQGHILKMFPSTWYV"
|
|
342
342
|
tokens = tokenizer(test_sequence, return_tensors="pt")
|
|
343
|
-
print(
|
|
343
|
+
print("✓ Tokenization test successful")
|
|
344
344
|
print(f"✓ Input sequence: {test_sequence}")
|
|
345
345
|
print(f"✓ Tokenized shape: {tokens['input_ids'].shape}")
|
|
346
346
|
|
|
@@ -356,7 +356,7 @@ def main():
|
|
|
356
356
|
dummy_input, attention_mask=dummy_mask, output_hidden_states=True
|
|
357
357
|
)
|
|
358
358
|
|
|
359
|
-
print(
|
|
359
|
+
print("✓ Forward pass successful")
|
|
360
360
|
print(f"✓ Output logits shape: {output.logits.shape}")
|
|
361
361
|
print(
|
|
362
362
|
f"✓ Number of hidden states: {len(output.hidden_states) if output.hidden_states else 0}"
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
# Embedding options
|
|
2
|
-
PEPE can extract numerous different representations from the input sequences while embedding each only once.
|
|
2
|
+
PEPE can extract numerous different representations from the input sequences while embedding each only once.
|
|
3
3
|
## Layer selection
|
|
4
4
|
Protein representations can be extracted from any of a PLMs hidden layers using the ```--layers``` argument and passing a list of integers. Use negative integers to index layers from the last element. E.g. ```"1"``` is first layer, ```"-1"``` is the last layer, ```"-2"``` is second to last layer, etc. Use ```"all"``` to select all layers. Default option: ```"-1"```.
|
|
5
5
|
```sh
|
|
@@ -18,7 +18,7 @@ Multiple embedding modes can be selected at once using the ```--extract_embeddin
|
|
|
18
18
|
- ```"mean_pooled"``` (default option): Average of ```"per_token"``` embedding over all amino acid tokens of the protein sequence.
|
|
19
19
|
- ```"substring_pooled"```: Average of ```"per_token"``` embedding over a specified substring of the protein sequence. Additional arguments when selected:
|
|
20
20
|
- ```--substring_path``` (required for ```"substring_pooled"```): Path to a CSV file with two columns. The first column contains the ```sequence_id``` and the second column must contain a substring of the sequence provided in the FASTA input file.
|
|
21
|
-
- ```--context``` (optional): Specify the number of residues before and after the substring to include during pooling.
|
|
21
|
+
- ```--context``` (optional): Specify the number of residues before and after the substring to include during pooling.
|
|
22
22
|
- Attention weights:
|
|
23
23
|
- ```"attention_head"```: Asymmetrical pairwise attention weight matrices of input tokens from each self-attention head of the specified layer(s)
|
|
24
24
|
- ```"attention_layer"```: Average of ```"attention_head"``` per specified layer.
|
|
@@ -38,4 +38,4 @@ pepe \
|
|
|
38
38
|
The way PEPE outputs representations can be configured with the following arguments:
|
|
39
39
|
- ```--streaming_output```: ```True``` (default) or ```False```. PEPE preallocates the required disk space and writes each batch of outputs concurrently. Disable if encountering file system issues.
|
|
40
40
|
- ```--precision```: ```full``` (default) or ```half```. Specifies whether to save representations as ```float32``` (full precision) or ```float16``` (half precision) numerical values for smaller file sizes.
|
|
41
|
-
- ```--flatten```: ```True``` or ```False```(default). When enabled, two-dimensional embedding modes (e.g. ```"per_token"``` or ```"attention_layer"```) will be flattened to one-dimensional vector along the first axis.
|
|
41
|
+
- ```--flatten```: ```True``` or ```False```(default). When enabled, two-dimensional embedding modes (e.g. ```"per_token"``` or ```"attention_layer"```) will be flattened to one-dimensional vector along the first axis.
|
|
@@ -50,6 +50,6 @@ pepe \
|
|
|
50
50
|
--model_name "examples/custom_model/example_protein_model" \ # pass the directory path containing custom PyTorch model
|
|
51
51
|
--tokenizer_from "alchemab/antiberta2-cssp" \ # Uses the same tokenizer as AntiBERTa2-CSSP
|
|
52
52
|
--fasta_path "src/tests/test_files/test.fasta" \
|
|
53
|
-
--output_path "src/tests/test_files/test_output"
|
|
53
|
+
--output_path "src/tests/test_files/test_output"
|
|
54
54
|
```
|
|
55
55
|
For details, see the [example_protein_model folder](examples/custom_model/example_protein_model) and the [python script](examples/custom_model/create_example_custom_model.py) for generating the example_protein_model files.
|
|
@@ -11,15 +11,13 @@ authors = [
|
|
|
11
11
|
description = "Pipeline for Easy Protein Embedding - Extract embeddings and attention matrices from protein sequences"
|
|
12
12
|
readme = "README.md"
|
|
13
13
|
license = {text = "MIT"}
|
|
14
|
-
requires-python = ">=3.
|
|
14
|
+
requires-python = ">=3.10"
|
|
15
15
|
classifiers = [
|
|
16
16
|
"Development Status :: 3 - Alpha",
|
|
17
17
|
"Intended Audience :: Science/Research",
|
|
18
18
|
"License :: OSI Approved :: MIT License",
|
|
19
19
|
"Operating System :: OS Independent",
|
|
20
20
|
"Programming Language :: Python :: 3",
|
|
21
|
-
"Programming Language :: Python :: 3.8",
|
|
22
|
-
"Programming Language :: Python :: 3.9",
|
|
23
21
|
"Programming Language :: Python :: 3.10",
|
|
24
22
|
"Programming Language :: Python :: 3.11",
|
|
25
23
|
"Topic :: Scientific/Engineering :: Bio-Informatics",
|
|
@@ -36,6 +34,19 @@ dependencies = [
|
|
|
36
34
|
"rjieba",
|
|
37
35
|
]
|
|
38
36
|
|
|
37
|
+
[project.optional-dependencies]
|
|
38
|
+
esm = ["fair-esm"]
|
|
39
|
+
# Developer toolchain: `pip install -e .[dev]` bootstraps everything needed to
|
|
40
|
+
# run the hooks and the local test suite.
|
|
41
|
+
dev = [
|
|
42
|
+
"pre-commit",
|
|
43
|
+
"ruff",
|
|
44
|
+
"mypy",
|
|
45
|
+
"pytest",
|
|
46
|
+
"pytest-cov",
|
|
47
|
+
"biopython",
|
|
48
|
+
]
|
|
49
|
+
|
|
39
50
|
[project.urls]
|
|
40
51
|
Homepage = "https://github.com/csi-greifflab/pepe-cli"
|
|
41
52
|
"Bug Reports" = "https://github.com/csi-greifflab/pepe-cli/issues"
|
|
@@ -78,3 +89,21 @@ module = [
|
|
|
78
89
|
"pepe.utils",
|
|
79
90
|
]
|
|
80
91
|
check_untyped_defs = true
|
|
92
|
+
|
|
93
|
+
[tool.ruff]
|
|
94
|
+
line-length = 88
|
|
95
|
+
# Matches requires-python (>= 3.10) and the mypy python_version.
|
|
96
|
+
target-version = "py310"
|
|
97
|
+
src = ["src"]
|
|
98
|
+
# Notebooks follow their own conventions (imports mid-cell, etc.); don't lint them.
|
|
99
|
+
extend-exclude = ["notebooks", "*.ipynb"]
|
|
100
|
+
|
|
101
|
+
[tool.ruff.lint]
|
|
102
|
+
# Pyflakes (F), pycodestyle errors/warnings (E/W), import sorting (I).
|
|
103
|
+
# E501 (line length) is owned by the formatter, which does not wrap strings or
|
|
104
|
+
# comments; enforcing it as a lint rule would flag intentionally long lines.
|
|
105
|
+
select = ["E", "W", "F", "I"]
|
|
106
|
+
ignore = ["E501"]
|
|
107
|
+
|
|
108
|
+
[tool.ruff.lint.isort]
|
|
109
|
+
known-first-party = ["pepe"]
|
|
@@ -1,10 +1,19 @@
|
|
|
1
|
-
from setuptools import setup, find_packages
|
|
2
1
|
import os
|
|
3
2
|
import sys
|
|
4
3
|
|
|
4
|
+
from setuptools import find_packages, setup
|
|
5
|
+
|
|
5
6
|
# Add src to the Python path to import metadata
|
|
6
|
-
sys.path.insert(0, os.path.join(os.path.dirname(__file__),
|
|
7
|
-
from pepe import
|
|
7
|
+
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))
|
|
8
|
+
from pepe import (
|
|
9
|
+
__author__,
|
|
10
|
+
__author_email__,
|
|
11
|
+
__description__,
|
|
12
|
+
__homepage__,
|
|
13
|
+
__module_name__,
|
|
14
|
+
__package_name__,
|
|
15
|
+
__version__,
|
|
16
|
+
)
|
|
8
17
|
|
|
9
18
|
|
|
10
19
|
# Read the README file for long description
|
|
@@ -32,14 +41,12 @@ setup(
|
|
|
32
41
|
"Operating System :: Linux",
|
|
33
42
|
"Operating System :: macOS",
|
|
34
43
|
"Programming Language :: Python :: 3",
|
|
35
|
-
"Programming Language :: Python :: 3.8",
|
|
36
|
-
"Programming Language :: Python :: 3.9",
|
|
37
44
|
"Programming Language :: Python :: 3.10",
|
|
38
45
|
"Programming Language :: Python :: 3.11",
|
|
39
46
|
"Topic :: Scientific/Engineering :: Bio-Informatics",
|
|
40
47
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
41
48
|
],
|
|
42
|
-
python_requires=">=3.
|
|
49
|
+
python_requires=">=3.10",
|
|
43
50
|
install_requires=[
|
|
44
51
|
"torch>=1.9.0",
|
|
45
52
|
"transformers>=4.20.0",
|
|
@@ -49,6 +56,9 @@ setup(
|
|
|
49
56
|
"alive_progress",
|
|
50
57
|
"rjieba",
|
|
51
58
|
],
|
|
59
|
+
extras_require={
|
|
60
|
+
"esm": ["fair-esm"],
|
|
61
|
+
},
|
|
52
62
|
entry_points={
|
|
53
63
|
"console_scripts": [
|
|
54
64
|
f"pepe = {__module_name__}.__main__:main",
|
|
@@ -2,12 +2,14 @@ import logging
|
|
|
2
2
|
import sys
|
|
3
3
|
|
|
4
4
|
try:
|
|
5
|
-
from pepe.api import
|
|
5
|
+
from pepe.api import (
|
|
6
|
+
embed, # noqa: F401 (re-exported to top level; needs torch at runtime)
|
|
7
|
+
)
|
|
6
8
|
except ImportError:
|
|
7
9
|
pass
|
|
8
10
|
|
|
9
11
|
# Package metadata - single source of truth
|
|
10
|
-
__version__ = "1.
|
|
12
|
+
__version__ = "1.5.0"
|
|
11
13
|
__package_name__ = "pepe-cli"
|
|
12
14
|
__module_name__ = "pepe"
|
|
13
15
|
__author__ = "Jahn Zhong"
|
|
@@ -1,14 +1,14 @@
|
|
|
1
|
+
import logging
|
|
1
2
|
import os
|
|
2
3
|
import tempfile
|
|
3
|
-
import logging
|
|
4
4
|
from types import SimpleNamespace
|
|
5
|
-
from typing import Dict, List, Optional, Union
|
|
5
|
+
from typing import Any, Dict, List, Optional, Union
|
|
6
6
|
|
|
7
7
|
from pepe.model_selecter import select_model
|
|
8
|
-
import pepe.utils
|
|
9
8
|
|
|
10
9
|
logger = logging.getLogger("pepe.api")
|
|
11
10
|
|
|
11
|
+
|
|
12
12
|
def embed(
|
|
13
13
|
model_name: str,
|
|
14
14
|
sequences: Optional[Union[Dict[str, str], List[str]]] = None,
|
|
@@ -23,7 +23,7 @@ def embed(
|
|
|
23
23
|
discard_padding: bool = False,
|
|
24
24
|
max_input_length: str = "max_length",
|
|
25
25
|
experiment_name: Optional[str] = None,
|
|
26
|
-
**kwargs
|
|
26
|
+
**kwargs,
|
|
27
27
|
) -> Dict[str, Any]:
|
|
28
28
|
"""
|
|
29
29
|
High-level API for generating protein embeddings.
|
|
@@ -52,8 +52,10 @@ def embed(
|
|
|
52
52
|
if sequences is not None:
|
|
53
53
|
if isinstance(sequences, list):
|
|
54
54
|
sequences = {f"seq_{i}": seq for i, seq in enumerate(sequences)}
|
|
55
|
-
|
|
56
|
-
temp_fasta = tempfile.NamedTemporaryFile(
|
|
55
|
+
|
|
56
|
+
temp_fasta = tempfile.NamedTemporaryFile(
|
|
57
|
+
mode="w", suffix=".fasta", delete=False
|
|
58
|
+
)
|
|
57
59
|
for label, seq in sequences.items():
|
|
58
60
|
temp_fasta.write(f">{label}\n{seq}\n")
|
|
59
61
|
temp_fasta.close()
|
|
@@ -65,11 +67,13 @@ def embed(
|
|
|
65
67
|
return_results = False
|
|
66
68
|
if output_path is None:
|
|
67
69
|
if streaming_output:
|
|
68
|
-
logger.warning(
|
|
70
|
+
logger.warning(
|
|
71
|
+
"No output_path provided. Disabling streaming_output and returning in-memory results."
|
|
72
|
+
)
|
|
69
73
|
output_path = tempfile.mkdtemp()
|
|
70
74
|
streaming_output = False
|
|
71
75
|
return_results = True
|
|
72
|
-
|
|
76
|
+
|
|
73
77
|
# Create args object
|
|
74
78
|
args_dict = {
|
|
75
79
|
"model_name": model_name,
|
|
@@ -97,26 +101,26 @@ def embed(
|
|
|
97
101
|
"flush_batches_after": kwargs.get("flush_batches_after", 128),
|
|
98
102
|
"verbose": kwargs.get("verbose", False),
|
|
99
103
|
}
|
|
100
|
-
|
|
104
|
+
|
|
101
105
|
args = SimpleNamespace(**args_dict)
|
|
102
|
-
|
|
106
|
+
|
|
103
107
|
selected_model_class = select_model(
|
|
104
108
|
model_name, trust_remote_code=args.trust_remote_code
|
|
105
109
|
)
|
|
106
110
|
embedder = selected_model_class(args)
|
|
107
111
|
embedder.run()
|
|
108
|
-
|
|
112
|
+
|
|
109
113
|
results = {"output_path": output_path}
|
|
110
|
-
|
|
114
|
+
|
|
111
115
|
if return_results:
|
|
112
116
|
# Pick up in-memory results before they are lost
|
|
113
117
|
for output_type in extract_embeddings:
|
|
114
118
|
obj = getattr(embedder, output_type, None)
|
|
115
119
|
if obj and "output_data" in obj:
|
|
116
120
|
results[output_type] = obj["output_data"]
|
|
117
|
-
|
|
121
|
+
|
|
118
122
|
# Cleanup temp file
|
|
119
123
|
if temp_fasta:
|
|
120
124
|
os.unlink(temp_fasta.name)
|
|
121
|
-
|
|
125
|
+
|
|
122
126
|
return results
|