pepe-cli 1.3.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.
Files changed (66) hide show
  1. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/MANIFEST.in +1 -0
  2. {pepe_cli-1.3.0/src/pepe_cli.egg-info → pepe_cli-1.5.0}/PKG-INFO +51 -13
  3. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/README.md +40 -9
  4. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/create_example_custom_model.py +11 -11
  5. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/config.json +1 -1
  6. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/special_tokens_map.json +1 -1
  7. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/tokenizer_config.json +1 -1
  8. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/vocab.json +1 -1
  9. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/vocab.txt +1 -1
  10. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/embedding_options.md +3 -3
  11. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/model_selection.md +1 -1
  12. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/pyproject.toml +56 -3
  13. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/requirements.txt +1 -1
  14. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/setup.py +16 -6
  15. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/__init__.py +6 -4
  16. pepe_cli-1.5.0/src/pepe/__main__.py +49 -0
  17. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/api.py +23 -15
  18. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/embedders/base_embedder.py +386 -205
  19. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/embedders/custom_embedder.py +53 -35
  20. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/embedders/esm_embedder.py +42 -27
  21. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/embedders/huggingface_embedder.py +205 -114
  22. pepe_cli-1.5.0/src/pepe/embedders/metl_embedder.py +349 -0
  23. pepe_cli-1.5.0/src/pepe/model_errors.py +138 -0
  24. pepe_cli-1.5.0/src/pepe/model_selecter.py +277 -0
  25. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/parse_arguments.py +23 -5
  26. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/utils.py +280 -132
  27. {pepe_cli-1.3.0 → pepe_cli-1.5.0/src/pepe_cli.egg-info}/PKG-INFO +51 -13
  28. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/SOURCES.txt +15 -3
  29. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/requires.txt +11 -0
  30. pepe_cli-1.5.0/src/tests/__init__.py +0 -0
  31. pepe_cli-1.5.0/src/tests/conftest.py +137 -0
  32. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_api_unittest.py +27 -18
  33. pepe_cli-1.5.0/src/tests/test_custom_embedder.py +56 -0
  34. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_device_logic.py +32 -24
  35. pepe_cli-1.5.0/src/tests/test_esm1_integration.py +65 -0
  36. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_esmc_modes.py +4 -1
  37. pepe_cli-1.5.0/src/tests/test_generic_hf_integration.py +58 -0
  38. pepe_cli-1.5.0/src/tests/test_generic_safeguards.py +244 -0
  39. pepe_cli-1.5.0/src/tests/test_metl_dispatch.py +100 -0
  40. pepe_cli-1.5.0/src/tests/test_metl_hooks.py +178 -0
  41. pepe_cli-1.5.0/src/tests/test_metl_integration.py +97 -0
  42. pepe_cli-1.5.0/src/tests/test_model_selection.py +292 -0
  43. pepe_cli-1.5.0/src/tests/test_reconstruct_mean_pooled.py +88 -0
  44. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_splitting.py +72 -51
  45. pepe_cli-1.5.0/src/tests/test_streaming_roundtrip.py +159 -0
  46. pepe_cli-1.5.0/src/tests/test_sync_arguments.py +160 -0
  47. pepe_cli-1.5.0/src/tests/test_t5_antiberta2_integration.py +138 -0
  48. pepe_cli-1.3.0/src/pepe/__main__.py +0 -25
  49. pepe_cli-1.3.0/src/pepe/model_selecter.py +0 -169
  50. pepe_cli-1.3.0/src/tests/test_model_selection.py +0 -104
  51. pepe_cli-1.3.0/src/tests/test_run.py +0 -45
  52. pepe_cli-1.3.0/src/tests/verify_cross_tool_consistency.py +0 -132
  53. pepe_cli-1.3.0/src/tests/verify_readme.py +0 -65
  54. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/LICENSE +0 -0
  55. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_protein_model/pytorch_model.pt +0 -0
  56. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_sequences.fasta +0 -0
  57. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/examples/custom_model/example_substring.csv +0 -0
  58. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/setup.cfg +0 -0
  59. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe/embedders/__init__.py +0 -0
  60. /pepe_cli-1.3.0/src/tests/__init__.py → /pepe_cli-1.5.0/src/pepe/py.typed +0 -0
  61. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/dependency_links.txt +0 -0
  62. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/entry_points.txt +0 -0
  63. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/pepe_cli.egg-info/top_level.txt +0 -0
  64. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_esmc.py +0 -0
  65. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_load_layers_default.py +0 -0
  66. {pepe_cli-1.3.0 → pepe_cli-1.5.0}/src/tests/test_parse_arguments.py +0 -0
@@ -3,6 +3,7 @@ include LICENSE
3
3
  include requirements.txt
4
4
  include pyproject.toml
5
5
  recursive-include src *.py
6
+ include src/pepe/py.typed
6
7
  recursive-include examples *
7
8
  recursive-exclude * __pycache__
8
9
  recursive-exclude * *.py[co]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pepe-cli
3
- Version: 1.3.0
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.8
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. Extract embeddings:\
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. Extract embeddings:\
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
- fasta_file = create_example_fasta()
314
- substring_file = create_example_substring()
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(f"✓ Model loaded successfully")
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(f"✓ Model architecture created and weights loaded")
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
- f"✓ Tokenizer loaded successfully with AutoTokenizer.from_pretrained()"
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(f"✓ Tokenization test successful")
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(f"✓ Forward pass successful")
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}"
@@ -5,4 +5,4 @@
5
5
  "vocab_size": 25,
6
6
  "max_position_embeddings": 512,
7
7
  "model_type": "protein_transformer"
8
- }
8
+ }
@@ -3,4 +3,4 @@
3
3
  "sep_token": "<sep>",
4
4
  "pad_token": "<pad>",
5
5
  "unk_token": "<unk>"
6
- }
6
+ }
@@ -7,4 +7,4 @@
7
7
  "pad_token": "<pad>",
8
8
  "unk_token": "<unk>",
9
9
  "model_type": "protein_transformer"
10
- }
10
+ }
@@ -24,4 +24,4 @@
24
24
  "Y": 22,
25
25
  "V": 23,
26
26
  "X": 24
27
- }
27
+ }
@@ -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.8"
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"
@@ -52,5 +63,47 @@ package-dir = {"" = "src"}
52
63
  [tool.setuptools.packages.find]
53
64
  where = ["src"]
54
65
 
66
+ [tool.setuptools.package-data]
67
+ pepe = ["py.typed"]
68
+
55
69
  [tool.setuptools.dynamic]
56
70
  version = {attr = "pepe.__version__"}
71
+
72
+ [tool.mypy]
73
+ python_version = "3.10"
74
+ warn_return_any = false
75
+ warn_unused_ignores = true
76
+ ignore_missing_imports = true
77
+ check_untyped_defs = false
78
+ disallow_untyped_defs = false
79
+
80
+ [[tool.mypy.overrides]]
81
+ module = [
82
+ "pepe.api",
83
+ "pepe.model_selecter",
84
+ "pepe.model_errors",
85
+ "pepe.embedders.base_embedder",
86
+ "pepe.embedders.huggingface_embedder",
87
+ "pepe.embedders.custom_embedder",
88
+ "pepe.embedders.esm_embedder",
89
+ "pepe.utils",
90
+ ]
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"]
@@ -6,4 +6,4 @@ sentencepiece
6
6
  numpy
7
7
  protobuf
8
8
  alive_progress
9
- rjieba
9
+ rjieba
@@ -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__), 'src'))
7
- from pepe import __version__, __package_name__, __author__, __author_email__, __description__, __homepage__, __module_name__
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.8",
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 embed # Exported to top-level; requires torch at runtime
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.3.0"
12
+ __version__ = "1.5.0"
11
13
  __package_name__ = "pepe-cli"
12
14
  __module_name__ = "pepe"
13
15
  __author__ = "Jahn Zhong"
@@ -16,7 +18,7 @@ __description__ = "Pipeline for Easy Protein Embedding - Extract embeddings and
16
18
  __homepage__ = "https://github.com/csi-greifflab/pepe-cli"
17
19
 
18
20
 
19
- # Configure logging for the embedairr package
21
+ # Configure logging for the pepe package
20
22
  def setup_logging(level=logging.INFO):
21
23
  """Set up logging configuration for the PEPE package."""
22
24
  # Create formatter
@@ -30,7 +32,7 @@ def setup_logging(level=logging.INFO):
30
32
  console_handler.setFormatter(formatter)
31
33
 
32
34
  # Get root logger for the package
33
- logger = logging.getLogger("src")
35
+ logger = logging.getLogger("pepe")
34
36
  logger.setLevel(level)
35
37
 
36
38
  # Remove existing handlers to avoid duplication
@@ -0,0 +1,49 @@
1
+ import logging
2
+ import sys
3
+
4
+ from pepe.parse_arguments import parse_arguments
5
+
6
+ logger = logging.getLogger("pepe.__main__")
7
+
8
+
9
+ def main():
10
+ args = parse_arguments()
11
+
12
+ if args.check_model:
13
+ from pepe.model_selecter import report_model
14
+
15
+ report_model(args.check_model, trust_remote_code=args.trust_remote_code)
16
+ return
17
+
18
+ missing = []
19
+ if not args.model_name:
20
+ missing.append("--model_name")
21
+ if not args.fasta_path:
22
+ missing.append("--fasta_path")
23
+ if not args.output_path:
24
+ missing.append("--output_path")
25
+ if missing:
26
+ print(
27
+ f"Error: the following arguments are required: {', '.join(missing)}",
28
+ file=sys.stderr,
29
+ )
30
+ sys.exit(2)
31
+
32
+ # Lazy import to avoid loading heavy dependencies during --help
33
+ from pepe.model_selecter import select_model
34
+
35
+ selected_model = select_model(
36
+ args.model_name, trust_remote_code=args.trust_remote_code
37
+ )
38
+
39
+ embedder = selected_model(args)
40
+
41
+ logger.info("Embedder initialized")
42
+
43
+ embedder.run()
44
+
45
+ logger.info("All outputs saved.")
46
+
47
+
48
+ if __name__ == "__main__":
49
+ main()