From 42e27ea33d4ec846eac4b3a46d42e4ef4333469b Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 19:17:10 +0200 Subject: [PATCH 01/17] Add pre-commit hooks (Ruff + fixers) with CI enforcement Introduce a fast local hook layer that shifts cheap, deterministic checks left without replacing CI: - .pre-commit-config.yaml: hygiene fixers + Ruff lint(--fix)/format at commit; fast mocked unit tests at pre-push; opt-in local mypy hook (off by default). - pyproject.toml: [tool.ruff] (E,W,F,I; E501 owned by the formatter; py38) and a [project.optional-dependencies] dev extra so `pip install -e .[dev]` bootstraps the toolchain. - CI: new `lint` job runs `pre-commit run --all-files` so hooks are enforced even when skipped locally; model-download integration tests stay CI-only. - CONTRIBUTING.md: dev setup + how the commit/pre-push/CI layers relate. Co-Authored-By: Claude Fable 5 --- .github/workflows/test.yml | 24 +++++++++++ .pre-commit-config.yaml | 81 ++++++++++++++++++++++++++++++++++++++ CONTRIBUTING.md | 25 +++++++++++- pyproject.toml | 31 +++++++++++++++ 4 files changed, 159 insertions(+), 2 deletions(-) create mode 100644 .pre-commit-config.yaml diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4431eaa..6f0c336 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -15,6 +15,30 @@ concurrency: cancel-in-progress: true jobs: + lint: + name: Lint & format (pre-commit) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v5 + - uses: actions/setup-python@v6 + with: + python-version: "3.11" + cache: pip + - name: Install pre-commit + run: | + python -m pip install --upgrade pip + pip install pre-commit + # Cache the hook environments (ruff, etc.) keyed on the config. + - uses: actions/cache@v4 + with: + path: ~/.cache/pre-commit + key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }} + # Enforce the same fixers + Ruff lint/format contributors run locally. + # Pre-push hooks (unit tests) and the opt-in mypy hook do NOT run here: + # the `unit` job already runs the tests and the `mypy` job runs mypy. + - name: Run pre-commit on all files + run: pre-commit run --all-files --show-diff-on-failure + unit: name: Unit tests (mocked, no model downloads) runs-on: ubuntu-latest diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..32f97ea --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,81 @@ +# Pre-commit hooks — a fast, local first line of defense that shifts cheap, +# deterministic checks left so they don't waste a CI round-trip. These do NOT +# replace CI: CI remains the source of truth (full Python matrix + the heavy +# model-download integration tests) and re-runs these same hooks via +# `pre-commit run --all-files` so nothing slips through when a hook is skipped. +# +# Setup: +# pip install -e .[dev] +# pre-commit install # commit-time hooks +# pre-commit install --hook-type pre-push # pre-push hooks (unit tests) +# +# `git commit --no-verify` bypasses hooks locally, but CI will still enforce them. + +default_install_hook_types: [pre-commit, pre-push] + +repos: + # Generic hygiene fixers (run at commit time). + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + # The conda recipe is a Jinja2 template (`{% set %}` / `{{ }}`), not plain + # YAML, so the strict parser can't read it. + exclude: ^\.github/conda/meta\.yaml$ + - id: check-toml + - id: check-added-large-files + args: [--maxkb=1024] + - id: check-merge-conflict + - id: debug-statements + - id: mixed-line-ending + + # Ruff: lint (with autofix) + format. One tool in place of black/isort/flake8. + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.20 + hooks: + - id: ruff + args: [--fix] + - id: ruff-format + + # Fast mocked unit tests — run at PUSH time only (too slow for every commit). + # Mirrors the CI `unit` job's file list; model-download tests are excluded and + # stay CI-only. + - repo: local + hooks: + - id: pytest-unit + name: pytest (fast unit tests) + entry: python -m pytest -q + language: system + pass_filenames: false + stages: [pre-push] + args: + - src/tests/test_parse_arguments.py + - src/tests/test_sync_arguments.py + - src/tests/test_device_logic.py + - src/tests/test_load_layers_default.py + - src/tests/test_model_selection.py + - src/tests/test_generic_safeguards.py + - src/tests/test_custom_embedder.py + - src/tests/test_reconstruct_mean_pooled.py + + # OPT-IN (disabled by default): local mypy at push time, using the repo's + # venv and the existing [tool.mypy] config. CI already runs mypy + # (non-blocking); enable this if you want type errors to block locally. + # Remove the `stages: [manual]` line to activate at pre-push. + - id: mypy-local + name: mypy (scoped, opt-in) + entry: python -m mypy + language: system + pass_filenames: false + stages: [manual] + args: + - src/pepe/api.py + - src/pepe/model_selecter.py + - src/pepe/model_errors.py + - src/pepe/embedders/base_embedder.py + - src/pepe/embedders/huggingface_embedder.py + - src/pepe/embedders/custom_embedder.py + - src/pepe/embedders/esm_embedder.py + - src/pepe/utils.py diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 75adc27..d7ca543 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -9,8 +9,9 @@ first-time contributor or future-you coming back after six months. # Clone, then from the repo root: python -m venv .venv source .venv/bin/activate # Windows: .venv\Scripts\activate -pip install -e . # editable install: code changes take effect immediately -pip install pytest biopython # test-only dependencies +pip install -e ".[dev]" # editable install + dev toolchain (pytest, ruff, pre-commit, mypy) +pre-commit install # commit-time hooks (lint/format/fixers) +pre-commit install --hook-type pre-push # pre-push hooks (fast unit tests) # Optional backends, only if you work on those models: pip install fair-esm # ESM-1 models @@ -41,6 +42,26 @@ ESMC tests skip themselves automatically if the Biohub fork isn't installed. **Rule of thumb:** every time you fix a bug, add a test that would have caught it. That is how the suite becomes a map of everything that has ever gone wrong. +## Code style and pre-commit hooks + +Formatting and linting are handled by [Ruff](https://docs.astral.sh/ruff/) and run +automatically through pre-commit (config: `.pre-commit-config.yaml`). The checks are +layered by speed so nothing slows you down more than it has to: + +- **At commit** (instant): whitespace/EOF/YAML/TOML fixers + Ruff lint (`--fix`) and format. +- **At push** (seconds): the fast mocked unit tests, so breakage is caught before the CI + round-trip. An opt-in local `mypy` hook is available (disabled by default — CI runs mypy). +- **In CI** (minutes): the same hooks run again via `pre-commit run --all-files`, plus the + full Python matrix and the model-download integration tests. + +Pre-commit **does not replace CI** — it just shifts the cheap, deterministic checks left. +CI stays the source of truth: the model-download tests only run there, and the CI `lint` +job re-runs the hooks so a skipped hook can't sneak past. + +You can bypass hooks in a pinch with `git commit --no-verify` (or `git push --no-verify`), +but CI will still enforce them, so fix the issue before merging. To run everything on the +whole tree yourself: `pre-commit run --all-files`. + ## Branching and pull requests - `main` is the release branch (publishes to **PyPI**). diff --git a/pyproject.toml b/pyproject.toml index 924ad87..7bb99b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,18 @@ dependencies = [ "rjieba", ] +[project.optional-dependencies] +# Developer toolchain: `pip install -e .[dev]` bootstraps everything needed to +# run the hooks and the local test suite. +dev = [ + "pre-commit", + "ruff", + "mypy", + "pytest", + "pytest-cov", + "biopython", +] + [project.urls] Homepage = "https://github.com/csi-greifflab/pepe-cli" "Bug Reports" = "https://github.com/csi-greifflab/pepe-cli/issues" @@ -78,3 +90,22 @@ module = [ "pepe.utils", ] check_untyped_defs = true + +[tool.ruff] +line-length = 88 +# Match the declared floor (requires-python >= 3.8). NOTE: mypy targets 3.10 and +# the CI matrix starts at 3.9 — decide separately whether to drop 3.8 support. +target-version = "py38" +src = ["src"] +# Notebooks follow their own conventions (imports mid-cell, etc.); don't lint them. +extend-exclude = ["notebooks", "*.ipynb"] + +[tool.ruff.lint] +# Pyflakes (F), pycodestyle errors/warnings (E/W), import sorting (I). +# E501 (line length) is owned by the formatter, which does not wrap strings or +# comments; enforcing it as a lint rule would flag intentionally long lines. +select = ["E", "W", "F", "I"] +ignore = ["E501"] + +[tool.ruff.lint.isort] +known-first-party = ["pepe"] From efd7940ff67f9597bd65c87763683a244a2d66c9 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 19:21:53 +0200 Subject: [PATCH 02/17] Apply Ruff format + hooks normalization across the repo One-time mechanical pass to satisfy the newly added pre-commit hooks. No functional changes: - Ruff format + import sorting (isort) across src/, tests/, examples/. - Ruff autofixes: remove unused imports, collapse multi-imports; drop dead local assignments (conftest esm1 `device`, example script vars). - Mark the intentional `pepe.api.embed` re-export in __init__.py with noqa. - Generic fixers: trailing-whitespace + final-newline normalization on markdown/json/txt/yaml/workflow files. Co-Authored-By: Claude Fable 5 --- .github/QUICK_SETUP.md | 2 +- .github/bioconda_submission_guide.md | 4 +- .../workflows/publish-main-branch-trusted.yml | 42 +++--- .../workflows/publish-test-branch-trusted.yml | 32 ++--- README.md | 16 +-- .../create_example_custom_model.py | 22 ++-- .../example_protein_model/config.json | 2 +- .../special_tokens_map.json | 2 +- .../tokenizer_config.json | 2 +- .../example_protein_model/vocab.json | 2 +- .../example_protein_model/vocab.txt | 2 +- examples/embedding_options.md | 6 +- examples/model_selection.md | 2 +- notebooks/verify_embedder_consistency.ipynb | 87 ++++++++----- requirements.txt | 2 +- setup.py | 15 ++- src/pepe/__init__.py | 4 +- src/pepe/__main__.py | 1 + src/pepe/api.py | 32 +++-- src/pepe/embedders/base_embedder.py | 117 +++++++++-------- src/pepe/embedders/custom_embedder.py | 26 ++-- src/pepe/embedders/esm_embedder.py | 6 +- src/pepe/embedders/huggingface_embedder.py | 51 ++++---- src/pepe/model_selecter.py | 20 ++- src/pepe/utils.py | 119 +++++++++++------ src/tests/conftest.py | 6 +- src/tests/test_api_unittest.py | 38 +++--- src/tests/test_custom_embedder.py | 1 + src/tests/test_device_logic.py | 50 +++---- src/tests/test_esmc_modes.py | 5 +- src/tests/test_generic_hf_integration.py | 3 +- src/tests/test_generic_safeguards.py | 5 +- src/tests/test_model_selection.py | 23 ++-- src/tests/test_reconstruct_mean_pooled.py | 1 + src/tests/test_splitting.py | 122 +++++++++++------- src/tests/test_streaming_roundtrip.py | 34 +++-- src/tests/test_sync_arguments.py | 5 +- 37 files changed, 525 insertions(+), 384 deletions(-) diff --git a/.github/QUICK_SETUP.md b/.github/QUICK_SETUP.md index 583bba8..c67b164 100644 --- a/.github/QUICK_SETUP.md +++ b/.github/QUICK_SETUP.md @@ -6,7 +6,7 @@ ``` GitHub Repository → Settings → Environments → New environment Name: testpypi -Protection rules: +Protection rules: - Deployment branches: test - Save protection rules ``` diff --git a/.github/bioconda_submission_guide.md b/.github/bioconda_submission_guide.md index 888705d..9a9066b 100644 --- a/.github/bioconda_submission_guide.md +++ b/.github/bioconda_submission_guide.md @@ -59,7 +59,7 @@ about: license_file: LICENSE summary: Pipeline for Easy Protein Embedding description: | - PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting + PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings and attention matrices from protein sequences using pre-trained models. doc_url: https://github.com/csi-greifflab/pepe-cli#readme dev_url: https://github.com/csi-greifflab/pepe-cli @@ -71,7 +71,7 @@ extra: ## Step-by-Step Submission -1. **Fork Bioconda Recipes**: +1. **Fork Bioconda Recipes**: - Go to [bioconda/bioconda-recipes](https://github.com/bioconda/bioconda-recipes) and fork it. 2. **Clone your fork**: ```bash diff --git a/.github/workflows/publish-main-branch-trusted.yml b/.github/workflows/publish-main-branch-trusted.yml index 483185d..6a5cfbb 100644 --- a/.github/workflows/publish-main-branch-trusted.yml +++ b/.github/workflows/publish-main-branch-trusted.yml @@ -13,49 +13,49 @@ jobs: id-token: write # Required for trusted publishing contents: write # Required for creating releases pull-requests: read - + steps: - uses: actions/checkout@v5 - + - name: Set up Python uses: actions/setup-python@v6 with: python-version: '3.10' - + - name: Install build dependencies run: | python3 -m pip install --upgrade pip pip install build - + - name: Clean build artifacts run: | # Remove any cached files that might interfere with the build find . -name "__pycache__" -type d -exec rm -rf {} + || true find . -name "*.pyc" -type f -delete || true rm -rf dist/ build/ *.egg-info/ src/*.egg-info/ || true - + - name: Verify version for main branch run: | # Extract metadata directly from __init__.py (avoids importing package deps like torch) CURRENT_VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') PACKAGE_NAME=$(grep -m1 '__package_name__ = ' src/pepe/__init__.py | sed 's/__package_name__ = "\(.*\)"/\1/') MODULE_NAME=$(grep -m1 '__module_name__ = ' src/pepe/__init__.py | sed 's/__module_name__ = "\(.*\)"/\1/') - + # Check if version contains -dev or -test (should not be in main branch) if [[ "$CURRENT_VERSION" == *"-dev"* ]] || [[ "$CURRENT_VERSION" == *"-test"* ]]; then echo "❌ Error: Version $CURRENT_VERSION contains development/test suffix" echo "Main branch should have a clean version number (e.g., 1.0.0, 1.2.3)" exit 1 fi - + echo "✅ Version validation passed: $CURRENT_VERSION" echo "Package name: $PACKAGE_NAME" echo "Module name: $MODULE_NAME" grep "__version__" src/pepe/__init__.py - + - name: Build package run: python3 -m build - + - name: Verify build artifacts run: | echo "📦 Build artifacts created:" @@ -68,13 +68,13 @@ jobs: file "$file" echo "SHA256: $(sha256sum "$file" | cut -d' ' -f1)" done - + - name: Publish to PyPI (Trusted Publishing) uses: pypa/gh-action-pypi-publish@release/v1 with: verbose: true skip_existing: true - + - name: Output installation command run: | PACKAGE_NAME=$(grep -m1 '__package_name__ = ' src/pepe/__init__.py | sed 's/__package_name__ = "\(.*\)"/\1/') @@ -84,7 +84,7 @@ jobs: echo "pip install \"$PACKAGE_NAME==$CURRENT_VERSION\"" echo "Or for the latest version:" echo "pip install \"$PACKAGE_NAME\"" - + - name: Wait for package availability run: | echo "Waiting for package to be available on PyPI..." @@ -106,10 +106,10 @@ jobs: echo "Attempt $i failed, retrying in 30s..." sleep 30 done - + # Verify installation pip list | grep -i pepe - + # Test basic import python3 -c "import $MODULE_NAME; print('Package imported successfully')" @@ -132,13 +132,13 @@ jobs: echo "🎉 Package verification completed successfully!" echo "Package $PACKAGE_NAME version $CURRENT_VERSION is now available on PyPI!" - + - name: Get version for release id: get_version run: | VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') echo "version=$VERSION" >> $GITHUB_OUTPUT - + - name: Create GitHub Release uses: softprops/action-gh-release@v3 with: @@ -146,20 +146,20 @@ jobs: name: Release v${{ steps.get_version.outputs.version }} body: | Release of pepe-cli version ${{ steps.get_version.outputs.version }} - + ## Installation - + ### From PyPI (Recommended) ```bash pip install pepe-cli==${{ steps.get_version.outputs.version }} ``` - + ### From GitHub Release Download the wheel file and install: ```bash pip install pepe_cli-${{ steps.get_version.outputs.version }}-py3-none-any.whl ``` - + files: | dist/*.whl dist/*.tar.gz @@ -174,7 +174,7 @@ jobs: needs: build-and-publish-main steps: - uses: actions/checkout@v5 - + - name: Set up Conda uses: conda-incubator/setup-miniconda@v4 with: diff --git a/.github/workflows/publish-test-branch-trusted.yml b/.github/workflows/publish-test-branch-trusted.yml index d564ebe..9a19345 100644 --- a/.github/workflows/publish-test-branch-trusted.yml +++ b/.github/workflows/publish-test-branch-trusted.yml @@ -14,20 +14,20 @@ jobs: contents: read outputs: test_version: ${{ steps.version_step.outputs.test_version }} - + steps: - uses: actions/checkout@v5 - + - name: Set up Python uses: actions/setup-python@v6 with: python-version: '3.10' - + - name: Install build dependencies run: | python3 -m pip install --upgrade pip pip install build - + - name: Update version for test branch id: version_step run: | @@ -41,53 +41,53 @@ jobs: # then appending .dev{timestamp} (TestPyPI rejects re-uploading the same version). BASE_VERSION=$(echo "$CURRENT_VERSION" | sed -E 's/(-dev|\.dev[0-9]*)$//') TEST_VERSION="${BASE_VERSION}.dev${TIMESTAMP}" - + # Output version for other jobs echo "test_version=$TEST_VERSION" >> $GITHUB_OUTPUT - + # Update version in __init__.py sed -i "s/__version__ = \"$CURRENT_VERSION\"/__version__ = \"$TEST_VERSION\"/" src/pepe/__init__.py echo "Updated version from $CURRENT_VERSION to $TEST_VERSION" echo "Package name: $PACKAGE_NAME" echo "Module name: $MODULE_NAME" grep "__version__" src/pepe/__init__.py - + - name: Build package run: python3 -m build - + - name: Publish to TestPyPI (Trusted Publishing) uses: pypa/gh-action-pypi-publish@release/v1 with: repository-url: https://test.pypi.org/legacy/ verbose: true - + - name: Output installation command run: | PACKAGE_NAME=$(grep -m1 '__package_name__ = ' src/pepe/__init__.py | sed 's/__package_name__ = "\(.*\)"/\1/') echo "Package published to TestPyPI!" echo "To install the test version, run:" echo "pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple \"$PACKAGE_NAME\"" - + - name: Wait for package availability run: | echo "Waiting for package to be available on TestPyPI..." sleep 30 - + - name: Test installation and functionality run: | PACKAGE_NAME=$(grep -m1 '__package_name__ = ' src/pepe/__init__.py | sed 's/__package_name__ = "\(.*\)"/\1/') MODULE_NAME=$(grep -m1 '__module_name__ = ' src/pepe/__init__.py | sed 's/__module_name__ = "\(.*\)"/\1/') - + # Create a fresh virtual environment for testing python3 -m venv test_env source test_env/bin/activate - + # Install the package from TestPyPI pip install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple "$PACKAGE_NAME" - + # Verify installation pip list | grep -i pepe - + # Test basic import python3 -c "import $MODULE_NAME; print('Package imported successfully')" @@ -107,7 +107,7 @@ jobs: needs: build-and-publish-test steps: - uses: actions/checkout@v5 - + - name: Set up Conda uses: conda-incubator/setup-miniconda@v4 with: diff --git a/README.md b/README.md index f00744e..664c641 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,18 @@ # 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. +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. ### Citation > **PEPE: Scalable extraction of multi-modal protein language model representations** > Jahn Zhong, Niccolò Cardente, Geir Kjetil Sandve, Habib Bashour, Maria Francesca Abbate, Victor Greiff -> *bioRxiv* (2026) +> *bioRxiv* (2026) > [DOI: 10.1101/2025.10.13.680902](https://doi.org/10.1101/2025.10.13.680902) ## Quick start 1. Install PEPE - From PyPI: + From PyPI: ```sh pip install pepe-cli ``` @@ -20,7 +20,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a ```sh conda install -c jahn_zhong pepe-cli ``` - Or install from the GitHub repository: + Or install from the GitHub repository: ```sh git clone https://github.com/csi-greifflab/pepe-cli cd pepe-cli @@ -116,7 +116,7 @@ results = pepe.embed( # 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] +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] @@ -124,7 +124,7 @@ 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 +del results ``` ### Handling Long Sequences (Splitting & Reconstruction) @@ -132,8 +132,8 @@ del results 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**: +- **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. diff --git a/examples/custom_model/create_example_custom_model.py b/examples/custom_model/create_example_custom_model.py index 2845312..6edb0f2 100644 --- a/examples/custom_model/create_example_custom_model.py +++ b/examples/custom_model/create_example_custom_model.py @@ -3,12 +3,12 @@ Example script showing how to create and use a custom model with EmbedAIRR. """ +import json +import os + import torch import torch.nn as nn -import os -import json from transformers import AutoTokenizer -import numpy as np class ExampleProteinModel(nn.Module): @@ -309,9 +309,9 @@ def main(): # Create example model model_path = create_example_model_and_tokenizer() - # Create example data - fasta_file = create_example_fasta() - substring_file = create_example_substring() + # Create example data (written to disk for the demo; return paths unused here) + create_example_fasta() + create_example_substring() # Test loading the model print("\nTesting model loading...") @@ -319,7 +319,7 @@ def main(): model_data = torch.load( os.path.join(model_path, "pytorch_model.pt"), map_location="cpu" ) - print(f"✓ Model loaded successfully") + print("✓ Model loaded successfully") print(f"✓ Model config: {model_data.get('config', {})}") # Test creating the model @@ -327,20 +327,20 @@ def main(): vocab_size=25, hidden_size=384, num_layers=6, num_heads=12, max_length=512 ) model.load_state_dict(model_data["model"]) - print(f"✓ Model architecture created and weights loaded") + print("✓ Model architecture created and weights loaded") # Test tokenizer loading try: tokenizer = AutoTokenizer.from_pretrained(model_path) print( - f"✓ Tokenizer loaded successfully with AutoTokenizer.from_pretrained()" + "✓ Tokenizer loaded successfully with AutoTokenizer.from_pretrained()" ) print(f"✓ Tokenizer vocab size: {len(tokenizer.get_vocab())}") # Test tokenization test_sequence = "ARNDCEQGHILKMFPSTWYV" tokens = tokenizer(test_sequence, return_tensors="pt") - print(f"✓ Tokenization test successful") + print("✓ Tokenization test successful") print(f"✓ Input sequence: {test_sequence}") print(f"✓ Tokenized shape: {tokens['input_ids'].shape}") @@ -356,7 +356,7 @@ def main(): dummy_input, attention_mask=dummy_mask, output_hidden_states=True ) - print(f"✓ Forward pass successful") + print("✓ Forward pass successful") print(f"✓ Output logits shape: {output.logits.shape}") print( f"✓ Number of hidden states: {len(output.hidden_states) if output.hidden_states else 0}" diff --git a/examples/custom_model/example_protein_model/config.json b/examples/custom_model/example_protein_model/config.json index ad4ef93..b570c3c 100644 --- a/examples/custom_model/example_protein_model/config.json +++ b/examples/custom_model/example_protein_model/config.json @@ -5,4 +5,4 @@ "vocab_size": 25, "max_position_embeddings": 512, "model_type": "protein_transformer" -} \ No newline at end of file +} diff --git a/examples/custom_model/example_protein_model/special_tokens_map.json b/examples/custom_model/example_protein_model/special_tokens_map.json index 168df0d..afd9ba7 100644 --- a/examples/custom_model/example_protein_model/special_tokens_map.json +++ b/examples/custom_model/example_protein_model/special_tokens_map.json @@ -3,4 +3,4 @@ "sep_token": "", "pad_token": "", "unk_token": "" -} \ No newline at end of file +} diff --git a/examples/custom_model/example_protein_model/tokenizer_config.json b/examples/custom_model/example_protein_model/tokenizer_config.json index b64bfb2..a7abea3 100644 --- a/examples/custom_model/example_protein_model/tokenizer_config.json +++ b/examples/custom_model/example_protein_model/tokenizer_config.json @@ -7,4 +7,4 @@ "pad_token": "", "unk_token": "", "model_type": "protein_transformer" -} \ No newline at end of file +} diff --git a/examples/custom_model/example_protein_model/vocab.json b/examples/custom_model/example_protein_model/vocab.json index 0377446..935ca2c 100644 --- a/examples/custom_model/example_protein_model/vocab.json +++ b/examples/custom_model/example_protein_model/vocab.json @@ -24,4 +24,4 @@ "Y": 22, "V": 23, "X": 24 -} \ No newline at end of file +} diff --git a/examples/custom_model/example_protein_model/vocab.txt b/examples/custom_model/example_protein_model/vocab.txt index 6449278..ea8926b 100644 --- a/examples/custom_model/example_protein_model/vocab.txt +++ b/examples/custom_model/example_protein_model/vocab.txt @@ -22,4 +22,4 @@ T W Y V -X \ No newline at end of file +X diff --git a/examples/embedding_options.md b/examples/embedding_options.md index b82e4cf..a960f40 100644 --- a/examples/embedding_options.md +++ b/examples/embedding_options.md @@ -1,5 +1,5 @@ # Embedding options -PEPE can extract numerous different representations from the input sequences while embedding each only once. +PEPE can extract numerous different representations from the input sequences while embedding each only once. ## Layer selection 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"```. ```sh @@ -18,7 +18,7 @@ Multiple embedding modes can be selected at once using the ```--extract_embeddin - ```"mean_pooled"``` (default option): Average of ```"per_token"``` embedding over all amino acid tokens of the protein sequence. - ```"substring_pooled"```: Average of ```"per_token"``` embedding over a specified substring of the protein sequence. Additional arguments when selected: - ```--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. - - ```--context``` (optional): Specify the number of residues before and after the substring to include during pooling. + - ```--context``` (optional): Specify the number of residues before and after the substring to include during pooling. - Attention weights: - ```"attention_head"```: Asymmetrical pairwise attention weight matrices of input tokens from each self-attention head of the specified layer(s) - ```"attention_layer"```: Average of ```"attention_head"``` per specified layer. @@ -38,4 +38,4 @@ pepe \ The way PEPE outputs representations can be configured with the following arguments: - ```--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. - ```--precision```: ```full``` (default) or ```half```. Specifies whether to save representations as ```float32``` (full precision) or ```float16``` (half precision) numerical values for smaller file sizes. -- ```--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. \ No newline at end of file +- ```--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. diff --git a/examples/model_selection.md b/examples/model_selection.md index ceffd36..d40dfff 100644 --- a/examples/model_selection.md +++ b/examples/model_selection.md @@ -50,6 +50,6 @@ pepe \ --model_name "examples/custom_model/example_protein_model" \ # pass the directory path containing custom PyTorch model --tokenizer_from "alchemab/antiberta2-cssp" \ # Uses the same tokenizer as AntiBERTa2-CSSP --fasta_path "src/tests/test_files/test.fasta" \ - --output_path "src/tests/test_files/test_output" + --output_path "src/tests/test_files/test_output" ``` 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. diff --git a/notebooks/verify_embedder_consistency.ipynb b/notebooks/verify_embedder_consistency.ipynb index 0f82d9b..b646b9f 100644 --- a/notebooks/verify_embedder_consistency.ipynb +++ b/notebooks/verify_embedder_consistency.ipynb @@ -55,11 +55,12 @@ ], "source": [ "import os\n", - "import sys\n", + "import shutil\n", "import subprocess\n", - "import torch\n", + "import sys\n", + "\n", "import numpy as np\n", - "import shutil\n", + "import torch\n", "\n", "# Define base directory and paths\n", "base_dir = \"/doctorai/userdata/pepe-cli\"\n", @@ -99,7 +100,7 @@ " print(f\"Running PEPE (precision={precision})...\")\n", " if os.path.exists(output_dir):\n", " shutil.rmtree(output_dir)\n", - " \n", + "\n", " pepe.embed(\n", " model_name=\"facebook/esm2_t6_8M_UR50D\",\n", " fasta_path=fasta_path,\n", @@ -108,12 +109,17 @@ " extract_embeddings=[\"mean_pooled\"],\n", " streaming_output=False,\n", " device=\"cpu\",\n", - " precision=precision\n", + " precision=precision,\n", " )\n", - " \n", + "\n", " model_name = \"esm2_t6_8M_UR50D\"\n", - " mean_pooled_file = os.path.join(output_dir, model_name, \"mean_pooled\", f\"pepe_output_{model_name}_mean_pooled_layer_6.npy\")\n", - " \n", + " mean_pooled_file = os.path.join(\n", + " output_dir,\n", + " model_name,\n", + " \"mean_pooled\",\n", + " f\"pepe_output_{model_name}_mean_pooled_layer_6.npy\",\n", + " )\n", + "\n", " return np.load(mean_pooled_file)" ] }, @@ -132,36 +138,45 @@ "source": [ "def run_plmfit(plmfit_repo_path, venv_path, output_dir):\n", " print(\"Running PLMFit...\")\n", - " \n", + "\n", " cmd = [\n", " os.path.join(venv_path, \"bin\", \"python3\"),\n", - " \"-m\", \"plmfit\",\n", - " \"--function\", \"extract_embeddings\",\n", - " \"--data_type\", \"verify\",\n", - " \"--plm\", \"esm2_t6_8M_UR50D\",\n", - " \"--output_dir\", output_dir,\n", - " \"--experiment_dir\", \"verify_exp\",\n", - " \"--experiment_name\", \"verify\",\n", - " \"--layer\", \"6\",\n", - " \"--reduction\", \"mean\"\n", + " \"-m\",\n", + " \"plmfit\",\n", + " \"--function\",\n", + " \"extract_embeddings\",\n", + " \"--data_type\",\n", + " \"verify\",\n", + " \"--plm\",\n", + " \"esm2_t6_8M_UR50D\",\n", + " \"--output_dir\",\n", + " output_dir,\n", + " \"--experiment_dir\",\n", + " \"verify_exp\",\n", + " \"--experiment_name\",\n", + " \"verify\",\n", + " \"--layer\",\n", + " \"6\",\n", + " \"--reduction\",\n", + " \"mean\",\n", " ]\n", - " \n", + "\n", " env = os.environ.copy()\n", " env[\"DATA_DIR\"] = os.path.join(plmfit_repo_path, \"data\")\n", " env[\"CONFIG_DIR\"] = os.path.join(plmfit_repo_path, \"config\")\n", " env[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n", - " \n", + "\n", " subprocess.run(cmd, check=True, cwd=plmfit_repo_path, env=env)\n", - " \n", + "\n", " output_pt = os.path.join(plmfit_repo_path, \"verify_exp\", \"verify.pt\")\n", - " \n", + "\n", " if not os.path.exists(output_pt):\n", - " for root, dirs, files in os.walk(os.path.join(plmfit_repo_path, 'verify_exp')):\n", + " for root, dirs, files in os.walk(os.path.join(plmfit_repo_path, \"verify_exp\")):\n", " for f in files:\n", " if f.endswith(\"verify.pt\"):\n", " output_pt = os.path.join(root, f)\n", " break\n", - " \n", + "\n", " return torch.load(output_pt, map_location=\"cpu\")" ] }, @@ -184,10 +199,10 @@ " os.path.join(venv_path, \"bin\", \"python3\"),\n", " os.path.join(base_dir, \"src/tests/run_official_esm.py\"),\n", " fasta_path,\n", - " output_path\n", + " output_path,\n", " ]\n", " subprocess.run(cmd, check=True)\n", - " \n", + "\n", " return torch.load(output_path, map_location=\"cpu\")" ] }, @@ -2075,7 +2090,9 @@ "\n", "# Run Official ESM\n", "esm_out_path = os.path.join(base_dir, \"test_verify_esm_official_nb.pt\")\n", - "esm_results = run_official_esm(fasta_path, os.path.join(base_dir, \"venv_esm_official\"), esm_out_path)\n", + "esm_results = run_official_esm(\n", + " fasta_path, os.path.join(base_dir, \"venv_esm_official\"), esm_out_path\n", + ")\n", "mean_pooled_esm_residues = esm_results[\"mean_pooled_residues\"]\n", "mean_pooled_esm_all = esm_results[\"mean_pooled_all\"]\n", "\n", @@ -2091,12 +2108,18 @@ "mean_pooled_esm_all = mean_pooled_esm_all.float()\n", "\n", "# 1. PEPE vs Official ESM (Residue-only pooling)\n", - "match_mp_res = torch.allclose(mean_pooled_pepe_torch, mean_pooled_esm_residues, atol=1e-5)\n", - "print(f\"PEPE vs Official ESM (Mean Pooled Residues): {'MATCH' if match_mp_res else 'FAIL'}\")\n", + "match_mp_res = torch.allclose(\n", + " mean_pooled_pepe_torch, mean_pooled_esm_residues, atol=1e-5\n", + ")\n", + "print(\n", + " f\"PEPE vs Official ESM (Mean Pooled Residues): {'MATCH' if match_mp_res else 'FAIL'}\"\n", + ")\n", "\n", "# 2. PLMFit vs Official ESM (All-token pooling)\n", "match_plmfit = torch.allclose(mean_pooled_plmfit, mean_pooled_esm_all, atol=1e-5)\n", - "print(f\"PLMFit vs Official ESM (Mean Pooled All Tokens): {'MATCH' if match_plmfit else 'FAIL'}\")" + "print(\n", + " f\"PLMFit vs Official ESM (Mean Pooled All Tokens): {'MATCH' if match_plmfit else 'FAIL'}\"\n", + ")" ] }, { @@ -3954,7 +3977,9 @@ "# Compare PEPE FP16 vs Official ESM (cast to FP16)\n", "mean_pooled_esm_16 = mean_pooled_esm_residues.half().float()\n", "match_16 = torch.allclose(mean_pooled_pepe_16_torch, mean_pooled_esm_16, atol=1e-3)\n", - "print(f\"PEPE FP16 vs Official ESM (casted, atol=1e-3): {'MATCH' if match_16 else 'FAIL'}\")" + "print(\n", + " f\"PEPE FP16 vs Official ESM (casted, atol=1e-3): {'MATCH' if match_16 else 'FAIL'}\"\n", + ")" ] }, { diff --git a/requirements.txt b/requirements.txt index e446af4..5c1b66a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,4 @@ sentencepiece numpy protobuf alive_progress -rjieba \ No newline at end of file +rjieba diff --git a/setup.py b/setup.py index bfc6e58..2e1e099 100644 --- a/setup.py +++ b/setup.py @@ -1,10 +1,19 @@ -from setuptools import setup, find_packages import os import sys +from setuptools import find_packages, setup + # Add src to the Python path to import metadata -sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src')) -from pepe import __version__, __package_name__, __author__, __author_email__, __description__, __homepage__, __module_name__ +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) +from pepe import ( + __author__, + __author_email__, + __description__, + __homepage__, + __module_name__, + __package_name__, + __version__, +) # Read the README file for long description diff --git a/src/pepe/__init__.py b/src/pepe/__init__.py index 33f8e77..fa7ef4f 100644 --- a/src/pepe/__init__.py +++ b/src/pepe/__init__.py @@ -2,7 +2,9 @@ import sys try: - from pepe.api import embed # Exported to top-level; requires torch at runtime + from pepe.api import ( + embed, # noqa: F401 (re-exported to top level; needs torch at runtime) + ) except ImportError: pass diff --git a/src/pepe/__main__.py b/src/pepe/__main__.py index ac755b2..852c5b1 100644 --- a/src/pepe/__main__.py +++ b/src/pepe/__main__.py @@ -1,5 +1,6 @@ import logging import sys + from pepe.parse_arguments import parse_arguments logger = logging.getLogger("pepe.__main__") diff --git a/src/pepe/api.py b/src/pepe/api.py index 621acb6..b64a10e 100644 --- a/src/pepe/api.py +++ b/src/pepe/api.py @@ -1,14 +1,14 @@ +import logging import os import tempfile -import logging from types import SimpleNamespace -from typing import Dict, List, Optional, Union, Any +from typing import Any, Dict, List, Optional, Union 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, @@ -23,7 +23,7 @@ def embed( discard_padding: bool = False, max_input_length: str = "max_length", experiment_name: Optional[str] = None, - **kwargs + **kwargs, ) -> Dict[str, Any]: """ High-level API for generating protein embeddings. @@ -52,8 +52,10 @@ def embed( 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) + + 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() @@ -65,11 +67,13 @@ def embed( 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.") + 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, @@ -97,26 +101,26 @@ def embed( "flush_batches_after": kwargs.get("flush_batches_after", 128), "verbose": kwargs.get("verbose", False), } - + args = SimpleNamespace(**args_dict) - + selected_model_class = select_model( model_name, trust_remote_code=args.trust_remote_code ) 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 diff --git a/src/pepe/embedders/base_embedder.py b/src/pepe/embedders/base_embedder.py index 8c93cc3..d12a6ed 100644 --- a/src/pepe/embedders/base_embedder.py +++ b/src/pepe/embedders/base_embedder.py @@ -1,16 +1,18 @@ -import os import csv -import torch +import inspect +import logging +import os import re +import time +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple, cast + import numpy as np +import torch +from alive_progress import alive_bar from numpy.lib.format import open_memmap -import inspect -from typing import Any, Dict, Iterable, List, Optional, Tuple, cast + from pepe.utils import MultiIODispatcher, check_disk_free_space -from alive_progress import alive_bar -import time -from pathlib import Path -import logging logger = logging.getLogger("pepe.embedders.base_embedder") @@ -364,8 +366,7 @@ def _safe_compute( mask_chunks = (None, None) outs = [ - self._safe_compute(tc, mc) - for tc, mc in zip(toks_chunks, mask_chunks) + self._safe_compute(tc, mc) for tc, mc in zip(toks_chunks, mask_chunks) ] # outs is list of (logits, reps, attn) logits = ( @@ -451,7 +452,7 @@ def embed(self) -> None: toks, attention_mask ) self.total_gpu_time += time.time() - t0_gpu - + output_bundle = { "logits": logits, "attention_matrices": attention_matrices, @@ -475,7 +476,7 @@ def embed(self) -> None: time.sleep(0.05) if backpressure_triggered: self.total_backpressure_time += time.time() - t0_bp - + t0_io = time.time() self._extract_batch(output_bundle) self.total_io_enqueue_time += time.time() - t0_io @@ -500,9 +501,7 @@ def embed(self) -> None: logger.info( f"Total Backpressure wait time: {self.total_backpressure_time:.2f}s" ) - logger.info( - f"Total IO Enqueue time: {self.total_io_enqueue_time:.2f}s" - ) + logger.info(f"Total IO Enqueue time: {self.total_io_enqueue_time:.2f}s") if self.total_gpu_time > 0: overhead = ( self.total_backpressure_time + self.total_io_enqueue_time @@ -580,13 +579,8 @@ def get_substring_positions( substring = substring.replace("-", "") # get position of substring in sequence - start = max(full_sequence.find(substring) - context, 0) + int( - special_tokens - ) - end = ( - min(start + len(substring) + context, len(full_sequence)) - + special_tokens - ) + start = max(full_sequence.find(substring) - context, 0) + int(special_tokens) + end = min(start + len(substring) + context, len(full_sequence)) + special_tokens return start, end @@ -755,9 +749,7 @@ def _extract_attention_head( ), # Ensure it's on CPU and NumPy ) else: - self.attention_head["output_data"][layer][ - head - ].extend(tensor) + self.attention_head["output_data"][layer][head].extend(tensor) def _extract_attention_layer( self, @@ -788,9 +780,7 @@ def _extract_attention_layer( array=self._to_numpy(tensor), # Ensure it's on CPU and NumPy ) else: - self.attention_layer["output_data"][layer].extend( - tensor - ) + self.attention_layer["output_data"][layer].extend(tensor) def _extract_attention_model( self, @@ -855,7 +845,7 @@ def _prepare_tensor(self, data_list: Any, flatten: bool) -> Any: if self.discard_padding: # Handle variable-length sequences by returning an object array of numpy arrays return np.array([t.numpy() for t in data_list], dtype=object) - + tensor = torch.stack(data_list, dim=0) if flatten: tensor = tensor.flatten(start_dim=1) @@ -943,6 +933,7 @@ def run(self) -> None: self._cleanup_checkpoint() logger.info("Pipeline completed successfully!") + def _check_max_input_length(self) -> None: """Check if max_input_length exceeds the model's allowed maximum length and handle splitting.""" max_allowed = self._get_model_max_allowed() @@ -1030,12 +1021,14 @@ def _handle_sequence_splitting(self, max_allowed: int) -> None: """Split sequences that exceed max_allowed into chunks.""" new_sequences = {} self.chunks_mapping = {} - special_tokens_count = 2 # cls + eos (conservative default) + special_tokens_count = 2 # cls + eos (conservative default) chunk_size = max_allowed - special_tokens_count overlap = self.split_overlap if chunk_size <= overlap: - logger.error(f"chunk_size ({chunk_size}) must be greater than overlap ({overlap}). Disabling splitting.") + logger.error( + f"chunk_size ({chunk_size}) must be greater than overlap ({overlap}). Disabling splitting." + ) return for label, sequence in self.sequences.items(): @@ -1051,23 +1044,23 @@ def _handle_sequence_splitting(self, max_allowed: int) -> None: end = min(start + chunk_size, len(sequence)) chunk_payload = sequence[start:end] chunk_label = f"{label}_chunk_{chunk_idx}" - + new_sequences[chunk_label] = chunk_payload self.chunk_payload_lengths[chunk_label] = len(chunk_payload) chunks.append(chunk_label) - + if end == len(sequence): break start = end - overlap chunk_idx += 1 - + self.chunks_mapping[label] = chunks - + self.sequences = new_sequences # Update num_sequences - if hasattr(self, 'num_sequences'): + if hasattr(self, "num_sequences"): self.num_sequences = len(self.sequences) - + # Update max_input_length to chunk size self.max_input_length = chunk_size @@ -1078,9 +1071,9 @@ def _reconstruct_chunks(self) -> None: assert self.layers is not None logger.info("Reconstructing original sequences from chunks...") - + label_to_idx = {label: i for i, label in enumerate(self.sequence_labels)} - + # For each output type, we need to rebuild the data rebuild_logits = "logits" in self.output_types per_token_requested = "per_token" in self.output_types @@ -1128,24 +1121,26 @@ def _reconstruct_chunks(self) -> None: orig_label = parent is_chunk = True break - + if orig_label in labels_processed: continue - + labels_processed.add(orig_label) new_sequence_labels.append(orig_label) - + if not is_chunk: # Just copy the existing data idx = label_to_idx[label] for output_type, obj in output_type_map: for layer in self.layers: - reconstructed_data[output_type][layer].append(obj["output_data"][layer][idx]) + reconstructed_data[output_type][layer].append( + obj["output_data"][layer][idx] + ) continue - + # Reconstruct from chunks chunk_labels = self.chunks_mapping[orig_label] - + # 1. Concatenate per-token and logits if build_per_token or rebuild_logits: for output_type, obj, flag in [ @@ -1159,28 +1154,40 @@ def _reconstruct_chunks(self) -> None: idx = label_to_idx[cl] full_tensor = obj["output_data"][layer][idx] payload_len = self.chunk_payload_lengths[cl] - + # Identify indices for extraction start_idx = 1 if i > 0: start_idx += self.split_overlap - + end_idx = 1 + payload_len meat = full_tensor[start_idx:end_idx] - if i == 0: meat = torch.cat([full_tensor[0:1], meat], dim=0) if i == len(chunk_labels) - 1: expected_unpadded_len = 1 + payload_len - # Safe bet: if there's no EOS, it's either padding or out of bounds. + # Safe bet: if there's no EOS, it's either padding or out of bounds. # To be robust during reconstruction of standard models, we should append if `add_special_tokens` was True and the tokenizer adds EOS. # The simplest heuristic: the original tokenizer encoded "" into >1 token or it has an EOS token - eos_count = len(self.tokenizer.encode("", add_special_tokens=True)) - 1 if hasattr(self, "tokenizer") and hasattr(self.tokenizer, "encode") else 0 - + eos_count = ( + len( + self.tokenizer.encode( + "", add_special_tokens=True + ) + ) + - 1 + if hasattr(self, "tokenizer") + and hasattr(self.tokenizer, "encode") + else 0 + ) + if eos_count > 0: - appended = full_tensor[expected_unpadded_len : expected_unpadded_len + 1] + appended = full_tensor[ + expected_unpadded_len : expected_unpadded_len + + 1 + ] meat = torch.cat([meat, appended], dim=0) parts.append(meat) @@ -1205,5 +1212,7 @@ def _reconstruct_chunks(self) -> None: continue for layer in self.layers: obj["output_data"][layer] = reconstructed_data[output_type][layer] - - logger.info(f"Reconstruction complete. Final sequence count: {self.num_sequences}") + + logger.info( + f"Reconstruction complete. Final sequence count: {self.num_sequences}" + ) diff --git a/src/pepe/embedders/custom_embedder.py b/src/pepe/embedders/custom_embedder.py index 261791c..e182fe6 100644 --- a/src/pepe/embedders/custom_embedder.py +++ b/src/pepe/embedders/custom_embedder.py @@ -1,13 +1,13 @@ +import json import logging -import torch -import torch.nn as nn import os -import json -import sys from typing import Any, Dict, List, Optional, Set, Tuple + +import torch +from transformers import AutoTokenizer + import pepe.utils from pepe.embedders.base_embedder import BaseEmbedder -from transformers import AutoTokenizer logger = logging.getLogger("pepe.embedders.custom_embedder") @@ -121,9 +121,7 @@ def _initialize_model( # for the config/metadata dict format) and therefore executes arbitrary # code from the file — acceptable here because a custom .pt path is the # user's own local model, not a remote download. - model_data = torch.load( - model_file_path, map_location="cpu", weights_only=False - ) + model_data = torch.load(model_file_path, map_location="cpu", weights_only=False) # Handle different model saving formats if isinstance(model_data, dict): @@ -229,9 +227,9 @@ def _infer_num_heads(self, state_dict: Any) -> int: def _infer_embedding_size(self, state_dict: Any) -> int: """Infer embedding size from state dict.""" - assert ( - state_dict is not None - ), "State dict cannot be None for embedding size inference" + assert state_dict is not None, ( + "State dict cannot be None for embedding size inference" + ) # Look for embedding layers emb_keys = [ @@ -307,9 +305,9 @@ def _load_layers(self, layers: Optional[List[int]] = None) -> List[int]: layers = [-1] # Validate layer indices - assert all( - -(self.num_layers + 1) <= i <= self.num_layers for i in layers - ), f"Layer indices must be in range [{-(self.num_layers + 1)}, {self.num_layers}]" + assert all(-(self.num_layers + 1) <= i <= self.num_layers for i in layers), ( + f"Layer indices must be in range [{-(self.num_layers + 1)}, {self.num_layers}]" + ) # Convert negative indices to positive layers = [(i + self.num_layers + 1) % (self.num_layers + 1) for i in layers] diff --git a/src/pepe/embedders/esm_embedder.py b/src/pepe/embedders/esm_embedder.py index 08f76fb..5aac638 100644 --- a/src/pepe/embedders/esm_embedder.py +++ b/src/pepe/embedders/esm_embedder.py @@ -1,8 +1,10 @@ import logging -import torch from typing import Any, Dict, List, Optional, Tuple -from pepe.embedders.base_embedder import BaseEmbedder + +import torch + import pepe.utils +from pepe.embedders.base_embedder import BaseEmbedder # Lazy imports to avoid loading heavy dependencies at import time diff --git a/src/pepe/embedders/huggingface_embedder.py b/src/pepe/embedders/huggingface_embedder.py index e23d824..f4dd35b 100644 --- a/src/pepe/embedders/huggingface_embedder.py +++ b/src/pepe/embedders/huggingface_embedder.py @@ -1,9 +1,11 @@ import logging import os +from typing import Any, Dict, List, Optional, Tuple + import torch + import pepe.utils from pepe.embedders.base_embedder import BaseEmbedder -from typing import Any, Dict, List, Optional, Tuple # Lazy imports to avoid loading heavy dependencies at import time @@ -13,6 +15,7 @@ def _import_transformers(): try: from transformers import T5EncoderModel + try: _t5_fast = importlib.import_module( "transformers.models.t5.tokenization_t5_fast" @@ -22,12 +25,17 @@ def _import_transformers(): import transformers T5TokenizerFast = getattr(transformers, "T5TokenizerFast") - from transformers import RoFormerTokenizer, RoFormerModel + from transformers import ( + AutoModel, + AutoModelForCausalLM, + AutoModelForMaskedLM, + AutoTokenizer, + RoFormerModel, + RoFormerTokenizer, + ) from transformers.models.roformer.modeling_roformer import ( RoFormerSinusoidalPositionalEmbedding, ) - from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM - from transformers import AutoModelForMaskedLM return ( T5EncoderModel, @@ -134,13 +142,11 @@ def _compute_outputs( hidden_states = outputs.hidden_states if isinstance(hidden_states, torch.Tensor): representations = { - layer: hidden_states[layer].to(dtype).cpu() - for layer in self.layers + layer: hidden_states[layer].to(dtype).cpu() for layer in self.layers } else: representations = { - layer: hidden_states[layer].to(dtype).cpu() - for layer in self.layers + layer: hidden_states[layer].to(dtype).cpu() for layer in self.layers } torch.cuda.empty_cache() else: @@ -179,7 +185,9 @@ def __init__(self, args: Any) -> None: ) self._set_output_objects() if not self.split_long_sequences: - assert self.max_input_length <= 256, "AntiBERTa2 only supports max_length <= 256. Use --split_long_sequences to process longer sequences." + assert self.max_input_length <= 256, ( + "AntiBERTa2 only supports max_length <= 256. Use --split_long_sequences to process longer sequences." + ) def _initialize_model( self, @@ -297,9 +305,7 @@ def _initialize_model( ) import importlib - _t5_slow = importlib.import_module( - "transformers.models.t5.tokenization_t5" - ) + _t5_slow = importlib.import_module("transformers.models.t5.tokenization_t5") tokenizer = _t5_slow.T5Tokenizer.from_pretrained( model_link, legacy=True, trust_remote_code=self.trust_remote_code ) @@ -418,9 +424,9 @@ def _initialize_model( model_kwargs["attn_implementation"] = "eager" if self.return_logits: - model = AutoModelForMaskedLM.from_pretrained( - model_link, **model_kwargs - ).to(device) + model = AutoModelForMaskedLM.from_pretrained(model_link, **model_kwargs).to( + device + ) else: model = AutoModel.from_pretrained(model_link, **model_kwargs).to(device) model.eval() @@ -449,8 +455,9 @@ def _compute_outputs( ) if return_logits: logits = ( - outputs.logits - .to(dtype=self._precision_to_dtype(self.precision, "torch")) + outputs.logits.to( + dtype=self._precision_to_dtype(self.precision, "torch") + ) .permute(2, 0, 1) .cpu() ) @@ -585,9 +592,9 @@ def _initialize_model( model_kwargs["attn_implementation"] = "eager" if self.return_logits: - model = AutoModelForMaskedLM.from_pretrained( - model_link, **model_kwargs - ).to(device) + model = AutoModelForMaskedLM.from_pretrained(model_link, **model_kwargs).to( + device + ) else: model = AutoModel.from_pretrained(model_link, **model_kwargs).to(device) model.eval() @@ -599,9 +606,7 @@ def _initialize_model( num_layers = _get_config_attr( config, "num_hidden_layers", "n_layers", "num_layers" ) - embedding_size = _get_config_attr( - config, "hidden_size", "d_model", "embed_dim" - ) + embedding_size = _get_config_attr(config, "hidden_size", "d_model", "embed_dim") return model, tokenizer, num_heads, num_layers, embedding_size diff --git a/src/pepe/model_selecter.py b/src/pepe/model_selecter.py index 06739e9..7d22065 100644 --- a/src/pepe/model_selecter.py +++ b/src/pepe/model_selecter.py @@ -1,10 +1,10 @@ +import os +from typing import Tuple, Type + from pepe.embedders.base_embedder import BaseEmbedder from pepe.embedders.custom_embedder import CustomEmbedder from pepe.model_errors import translate_hf_config_error -import os -from typing import Any, Tuple, Type - def _get_esm_embedder() -> Type[BaseEmbedder]: """Lazy import of ESM embedder to avoid loading heavy dependencies.""" @@ -29,7 +29,7 @@ def _get_esmc_embedder() -> Type[BaseEmbedder]: def _get_huggingface_embedders() -> Tuple[Type[BaseEmbedder], Type[BaseEmbedder]]: """Lazy import of HuggingFace embedders to avoid loading heavy dependencies.""" - from pepe.embedders.huggingface_embedder import T5Embedder, Antiberta2Embedder + from pepe.embedders.huggingface_embedder import Antiberta2Embedder, T5Embedder return T5Embedder, Antiberta2Embedder @@ -97,9 +97,7 @@ def _select_hf_model( model_name, trust_remote_code=trust_remote_code ) except Exception as e: - translate_hf_config_error( - model_name, e, trust_remote_code=trust_remote_code - ) + translate_hf_config_error(model_name, e, trust_remote_code=trust_remote_code) model_type = (getattr(config, "model_type", "") or "").lower() @@ -153,10 +151,10 @@ def _format_max_length(config, tokenizer): def report_model(model_name, trust_remote_code=False): """Load config + tokenizer only and print a compatibility summary.""" - import pepe.utils - from transformers import AutoConfig, AutoTokenizer + import pepe.utils + try: config = AutoConfig.from_pretrained( model_name, trust_remote_code=trust_remote_code @@ -165,9 +163,7 @@ def report_model(model_name, trust_remote_code=False): model_name, trust_remote_code=trust_remote_code ) except Exception as e: - translate_hf_config_error( - model_name, e, trust_remote_code=trust_remote_code - ) + translate_hf_config_error(model_name, e, trust_remote_code=trust_remote_code) embedder_cls = select_model(model_name, trust_remote_code=trust_remote_code) embedder_name = embedder_cls.__name__ diff --git a/src/pepe/utils.py b/src/pepe/utils.py index a8c1adf..3f95441 100644 --- a/src/pepe/utils.py +++ b/src/pepe/utils.py @@ -1,16 +1,17 @@ +import gc +import json import logging -from typing import Any, Dict, List, Optional, Sequence, Tuple, cast -from torch.utils.data import Dataset +import os +import queue import re -import torch +import shutil +import threading import time -import gc -import os -import json -from transformers import RoFormerTokenizer, T5Tokenizer -import threading, queue +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import torch from alive_progress import alive_bar -import shutil +from torch.utils.data import Dataset logger = logging.getLogger("pepe.utils") @@ -29,7 +30,7 @@ def __init__(self, dataset: Any, token_budget: int) -> None: # dataset[idx] -> (label, seq_str, toks, ...) sample_seq_len = len(dataset[0][2]) - + if sample_seq_len > token_budget: logger.warning( f"A sequence has length {sample_seq_len}, which exceeds the specified token budget (batch_size) of {token_budget}. " @@ -55,8 +56,6 @@ def __len__(self) -> int: return len(self.batches) - - class SequenceDictDataset(Dataset): substring_dict: Optional[Dict[str, str]] encoded_data: List[Any] @@ -193,7 +192,9 @@ def __init__( ) # (label, seq, toks, attention_mask) self.substring_masks = self._get_substring_masks() - def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gapped_sequences = True): + def _encode_sequences( + self, data, tokenizer, max_length, add_special_tokens, gapped_sequences=True + ): labels, raw_strs = zip(*data) strs: List[str] if gapped_sequences: @@ -205,12 +206,18 @@ def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gap max_token_length = max(len(seq) for seq in strs) # Account for special tokens (if any) - special_tokens_count = len(tokenizer.encode("", add_special_tokens=add_special_tokens)) if hasattr(tokenizer, "encode") else 0 + special_tokens_count = ( + len(tokenizer.encode("", add_special_tokens=add_special_tokens)) + if hasattr(tokenizer, "encode") + else 0 + ) max_token_length += special_tokens_count if max_length == "max_length": max_length = max_token_length - logger.info(f"Setting max_length to {max_length} (including {special_tokens_count} special tokens).") + logger.info( + f"Setting max_length to {max_length} (including {special_tokens_count} special tokens)." + ) else: max_length = int(max_length) if max_length < max_token_length: @@ -362,6 +369,7 @@ def __getitem__(self, idx): else: return labels, seqs, toks, None, None + class CustomDataset(SequenceDictDataset): """ Dataset class for custom embedder. @@ -375,7 +383,7 @@ def __init__( tokenizer, max_length, add_special_tokens=True, - gapped_sequences = True, + gapped_sequences=True, ): super().__init__(sequences, substring_dict, context) self.tokenizer = tokenizer @@ -400,7 +408,9 @@ def __init__( ) self.substring_masks = self._get_substring_masks() - def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gapped_sequences = True): + def _encode_sequences( + self, data, tokenizer, max_length, add_special_tokens, gapped_sequences=True + ): """Encode sequences using the tokenizer.""" labels, raw_strs = zip(*data) strs: List[str] @@ -415,7 +425,9 @@ def _encode_sequences(self, data, tokenizer, max_length, add_special_tokens, gap # Convert max_length to int if needed if max_length == "max_length": - max_token_length = max(len(s) for s in strs) + (2 if add_special_tokens else 0) + max_token_length = max(len(s) for s in strs) + ( + 2 if add_special_tokens else 0 + ) logger.info(f"Tokenizing {len(strs)} sequences...") out = tokenizer( @@ -464,6 +476,7 @@ def __getitem__(self, idx): else: return labels, seqs, toks, attn_mask + def check_input_tokens( valid_tokens: Any, sequences: Dict[str, str], @@ -484,10 +497,12 @@ def __str_to_list(sequence, bracket_type="square"): ) as bar: for label, sequence in sequences.items(): sequence = __str_to_list(sequence, bracket_type) - if "antiberta" in model_name and not split_long_sequences: # check for longest sequence - assert ( - len(sequence) <= 256 - ), f"Antiberta2 does not support sequences longer than 256 tokens. Found {len(sequence)} tokens in sequence {label}. Use --split_long_sequences to process longer sequences." + if ( + "antiberta" in model_name and not split_long_sequences + ): # check for longest sequence + assert len(sequence) <= 256, ( + f"Antiberta2 does not support sequences longer than 256 tokens. Found {len(sequence)} tokens in sequence {label}. Use --split_long_sequences to process longer sequences." + ) if not set(sequence).issubset(valid_tokens): raise ValueError( @@ -502,13 +517,13 @@ def __str_to_list(sequence, bracket_type="square"): def is_character_tokenizer(tokenizer) -> bool: """Return True when one amino acid maps to one token (no subword merging).""" - encoded = tokenizer.encode( - _CHARACTER_TOKENIZER_PROBE, add_special_tokens=False - ) + encoded = tokenizer.encode(_CHARACTER_TOKENIZER_PROBE, add_special_tokens=False) return len(encoded) == len(_CHARACTER_TOKENIZER_PROBE) -def warn_if_non_character_tokenizer(tokenizer: Any, model_name: Optional[str] = None) -> None: +def warn_if_non_character_tokenizer( + tokenizer: Any, model_name: Optional[str] = None +) -> None: """Log a prominent warning when the tokenizer is not character-level.""" if is_character_tokenizer(tokenizer): return @@ -544,6 +559,8 @@ def flush(): flush() return seq_dict + + def get_bracket_type(tokenizer: Any) -> str: """ Attempt to determine if special tokens use [ ] or < > brackets. @@ -551,17 +568,22 @@ def get_bracket_type(tokenizer: Any) -> str: """ try: special_tokens = [] - if hasattr(tokenizer, "special_tokens_map") and "additional_special_tokens" in tokenizer.special_tokens_map: + if ( + hasattr(tokenizer, "special_tokens_map") + and "additional_special_tokens" in tokenizer.special_tokens_map + ): tokens = tokenizer.special_tokens_map["additional_special_tokens"] if tokens and len(tokens) > 0: special_tokens.append(tokens[0]) - + if not special_tokens and hasattr(tokenizer, "all_special_tokens"): if tokenizer.all_special_tokens: special_tokens.append(tokenizer.all_special_tokens[0]) - + if not special_tokens: - logger.warning("Could not find any special tokens. Defaulting to 'square' brackets.") + logger.warning( + "Could not find any special tokens. Defaulting to 'square' brackets." + ) return "square" first_char = special_tokens[0][0] @@ -570,11 +592,17 @@ def get_bracket_type(tokenizer: Any) -> str: elif first_char == "[": return "square" else: - logger.warning(f"Unrecognized special token format (starts with '{first_char}'). Defaulting to 'square' brackets.") + logger.warning( + f"Unrecognized special token format (starts with '{first_char}'). Defaulting to 'square' brackets." + ) return "square" except Exception as e: - logger.warning(f"Error determining bracket type: {e}. Defaulting to 'square' brackets.") + logger.warning( + f"Error determining bracket type: {e}. Defaulting to 'square' brackets." + ) return "square" + + def gap_sequence(sequences: Sequence[str], bracket_type: str = "square") -> List[str]: if bracket_type == "square": seqs = [" ".join(re.findall(r"\[.*?\]|.", sequence)) for sequence in sequences] @@ -584,6 +612,7 @@ def gap_sequence(sequences: Sequence[str], bracket_type: str = "square") -> List raise ValueError(f"Invalid bracket type: {bracket_type}") return seqs + def flush_memmaps(obj: Any) -> None: """Recursively flush memory maps.""" if hasattr(obj, "flush") and callable(obj.flush): @@ -735,8 +764,10 @@ def flush_key(self, key): self.mark_range_completed(output_type, layer, head, offset, len(arr)) mmap_handle.flush() duration = time.time() - t0 - if duration > 1.0: # Only log slow flushes - logger.info(f"[IOFlushWorker] Slow flush for {key}: {duration:.2f}s for {batch_count} batches ({total_elements} elements)") + if duration > 1.0: # Only log slow flushes + logger.info( + f"[IOFlushWorker] Slow flush for {key}: {duration:.2f}s for {batch_count} batches ({total_elements} elements)" + ) except Exception as e: logger.error(f"[IOFlushWorker] Exception during flush: {e}") raise e @@ -749,7 +780,7 @@ def enqueue(self, output_type, layer, head, offset, array): # Check if this range was already completed (crash recovery) if self.is_range_completed(output_type, layer, head, offset, len(array)): logger.debug( - f"[IOFlushWorker] Skipping already completed range: {output_type}, {layer}, {head}, {offset}-{offset+len(array)}" + f"[IOFlushWorker] Skipping already completed range: {output_type}, {layer}, {head}, {offset}-{offset + len(array)}" ) return @@ -785,7 +816,7 @@ def stop(self, max_wait_time=60, force_shutdown=True): max_wait_time: Maximum time to wait for pending operations (seconds) force_shutdown: If True, force shutdown after max_wait_time even if work remains """ - logger.info(f"[IOFlushWorker] Initiating shutdown...") + logger.info("[IOFlushWorker] Initiating shutdown...") # Signal that we're shutting down self.shutdown_flag.set() @@ -903,9 +934,11 @@ def __init__( self.num_heavy_workers = 0 self.num_light_workers = num_workers elif num_workers == 1: - logger.info("[MultiIODispatcher] Only 1 worker available. Assigning all keys to this worker.") + logger.info( + "[MultiIODispatcher] Only 1 worker available. Assigning all keys to this worker." + ) self.num_heavy_workers = 1 - self.num_light_workers = 1 # Both point to the same worker + self.num_light_workers = 1 # Both point to the same worker else: self.num_heavy_workers = max(1, int(num_workers * heavy_proportion)) self.num_light_workers = num_workers - self.num_heavy_workers @@ -927,7 +960,9 @@ def __init__( shard_id = hash(key) % self.num_heavy_workers else: # Assign light keys to light workers - shard_id = self.num_heavy_workers + (hash(key) % self.num_light_workers) + shard_id = self.num_heavy_workers + ( + hash(key) % self.num_light_workers + ) sharded_registries[shard_id][key] = mmap for i, reg in enumerate(sharded_registries): @@ -947,7 +982,9 @@ def __init__( self.light_workers = self.workers else: self.heavy_workers = ( - self.workers[: self.num_heavy_workers] if self.num_heavy_workers > 0 else [] + self.workers[: self.num_heavy_workers] + if self.num_heavy_workers > 0 + else [] ) self.light_workers = self.workers[self.num_heavy_workers :] @@ -1049,7 +1086,7 @@ def _load_global_checkpoint(self): for key, ranges in self.global_completed_ranges.items(): total_bytes = sum(end - start for start, end in ranges) logger.info( - f" {key}: {len(ranges)} ranges, {total_bytes / (1024*1024):.1f}MB" + f" {key}: {len(ranges)} ranges, {total_bytes / (1024 * 1024):.1f}MB" ) except Exception as e: diff --git a/src/tests/conftest.py b/src/tests/conftest.py index 05faf2a..bc45955 100644 --- a/src/tests/conftest.py +++ b/src/tests/conftest.py @@ -1,4 +1,5 @@ """Shared pytest fixtures for PEPE tests.""" + import os import sys @@ -95,11 +96,10 @@ def esm1_model_cache(): "ESM-1 integration tests require the esm package" ) - from pepe.embedders.esm_embedder import ESMEmbedder - from esm import pretrained - device = torch.device("cpu") + from pepe.embedders.esm_embedder import ESMEmbedder + model, alphabet = pretrained.load_model_and_alphabet_hub(ESM1_MODEL_NAME) model.eval() model.prepend_bos = True diff --git a/src/tests/test_api_unittest.py b/src/tests/test_api_unittest.py index 94e82ba..4621ad9 100644 --- a/src/tests/test_api_unittest.py +++ b/src/tests/test_api_unittest.py @@ -1,10 +1,8 @@ import os import sys -import unittest -import shutil import tempfile -import numpy as np -import torch +import unittest + import pytest # Add src to sys.path @@ -18,12 +16,13 @@ pytest.mark.usefixtures("esm2_model_cache"), ] + class TestPepeAPI(unittest.TestCase): @classmethod def setUpClass(cls): cls.test_sequences = { "seq1": "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALTNAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR", - "seq2": "MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNCNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYK" + "seq2": "MNIFEMLRIDEGLRLKIYKDTEGYYTIGIGHLLTKSPSLNAAKSELDKAIGRNCNGVITKDEAEKLFNQDVDAAVRGILRNAKLKPVYDSLDAVRRAALINMVFQMGETGVAGFTNSLRMLQQKRWDEAAVNLAKSRWYNQTPNRAKRVITTFRTGTWDAYK", } cls.model_name = "esm2_t6_8M_UR50D" @@ -33,12 +32,12 @@ def test_embed_sequences_dict_in_memory(self): model_name=self.model_name, sequences=self.test_sequences, extract_embeddings=["mean_pooled"], - device="cpu" + device="cpu", ) - + self.assertIn("mean_pooled", results) self.assertIn("output_path", results) - + # ESM2-8M has 6 layers, default is last layer (-1 -> 6) mean_pooled = results["mean_pooled"] self.assertIn(6, mean_pooled) @@ -52,9 +51,9 @@ def test_embed_sequences_list_in_memory(self): model_name=self.model_name, sequences=seq_list, extract_embeddings=["mean_pooled"], - device="cpu" + device="cpu", ) - + self.assertIn("mean_pooled", results) mean_pooled = results["mean_pooled"] self.assertEqual(len(mean_pooled[6]), 2) @@ -68,16 +67,16 @@ def test_embed_to_disk(self): output_path=tmp_out, extract_embeddings=["mean_pooled"], device="cpu", - streaming_output=False + streaming_output=False, ) - + self.assertEqual(results["output_path"], tmp_out) # Verify file exists # Output structured as: {output_path}/{model_name_basename}/mean_pooled/... model_basename = os.path.basename(self.model_name) expected_dir = os.path.join(tmp_out, model_basename, "mean_pooled") self.assertTrue(os.path.exists(expected_dir)) - + files = os.listdir(expected_dir) self.assertTrue(any(f.endswith(".npy") for f in files)) @@ -89,12 +88,12 @@ def test_embed_discard_padding(self): extract_embeddings=["per_token"], device="cpu", discard_padding=True, - streaming_output=False + streaming_output=False, ) - + self.assertIn("per_token", results) per_token = results["per_token"][6] - + # Lengths should be different because padding is discarded len1 = per_token[0].shape[0] len2 = per_token[1].shape[0] @@ -113,12 +112,15 @@ def test_embed_streaming_no_path(self): sequences=self.test_sequences, extract_embeddings=["mean_pooled"], device="cpu", - streaming_output=True + streaming_output=True, ) # Ensure warning was logged - self.assertTrue(any("No output_path provided" in line for line in cm.output)) + self.assertTrue( + any("No output_path provided" in line for line in cm.output) + ) # Ensure results were returned (streaming was disabled) self.assertIn("mean_pooled", results) + if __name__ == "__main__": unittest.main() diff --git a/src/tests/test_custom_embedder.py b/src/tests/test_custom_embedder.py index 2184d35..27eede9 100644 --- a/src/tests/test_custom_embedder.py +++ b/src/tests/test_custom_embedder.py @@ -1,4 +1,5 @@ """Integration coverage for the CustomEmbedder (.pt) path.""" + import json import os import sys diff --git a/src/tests/test_device_logic.py b/src/tests/test_device_logic.py index d598201..21d7073 100644 --- a/src/tests/test_device_logic.py +++ b/src/tests/test_device_logic.py @@ -1,15 +1,15 @@ -import unittest -from unittest.mock import MagicMock, patch -import torch import os import sys +import unittest from types import SimpleNamespace +from unittest.mock import MagicMock, patch # Add src to sys.path sys.path.insert(0, os.path.abspath("src")) from pepe.embedders.base_embedder import BaseEmbedder + class TestDeviceLogic(unittest.TestCase): def setUp(self): self.args = SimpleNamespace( @@ -30,21 +30,20 @@ def setUp(self): experiment_name=None, context=0, flatten=False, - flush_batches_after=128 + flush_batches_after=128, ) @patch("torch.cuda.is_available", return_value=True) def test_base_device_logic_gpu_available(self, mock_cuda): # Mocking file system calls and substring loading - with patch("os.path.exists", return_value=True), \ - patch("os.makedirs"), \ - patch.object(BaseEmbedder, "_load_substrings", return_value=None): - + with patch("os.path.exists", return_value=True), patch( + "os.makedirs" + ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): # 1. GPU requested and available self.args.device = "cuda" base = BaseEmbedder(self.args) self.assertEqual(base.device.type, "cuda") - + # 2. CPU specifically requested even if GPU available self.args.device = "cpu" base = BaseEmbedder(self.args) @@ -52,10 +51,9 @@ def test_base_device_logic_gpu_available(self, mock_cuda): @patch("torch.cuda.is_available", return_value=True) def test_base_device_logic_gpu_specific(self, mock_cuda): - with patch("os.path.exists", return_value=True), \ - patch("os.makedirs"), \ - patch.object(BaseEmbedder, "_load_substrings", return_value=None): - + with patch("os.path.exists", return_value=True), patch( + "os.makedirs" + ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): # GPU cuda:1 requested and available self.args.device = "cuda:1" base = BaseEmbedder(self.args) @@ -64,10 +62,9 @@ def test_base_device_logic_gpu_specific(self, mock_cuda): @patch("torch.cuda.is_available", return_value=False) def test_base_device_logic_gpu_not_available(self, mock_cuda): - with patch("os.path.exists", return_value=True), \ - patch("os.makedirs"), \ - patch.object(BaseEmbedder, "_load_substrings", return_value=None): - + with patch("os.path.exists", return_value=True), patch( + "os.makedirs" + ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): # GPU requested but NOT available -> should fallback to CPU self.args.device = "cuda" base = BaseEmbedder(self.args) @@ -75,14 +72,14 @@ def test_base_device_logic_gpu_not_available(self, mock_cuda): def test_get_bracket_type_robustness(self): from pepe.utils import get_bracket_type - + # Test unknown bracket (should default to square) mock_tokenizer = MagicMock() mock_tokenizer.all_special_tokens = ["__special__"] del mock_tokenizer.special_tokens_map - + self.assertEqual(get_bracket_type(mock_tokenizer), "square") - + # Test no special tokens (should default to square) mock_tokenizer.all_special_tokens = [] self.assertEqual(get_bracket_type(mock_tokenizer), "square") @@ -94,18 +91,23 @@ def test_get_bracket_type_robustness(self): def test_subclass_device_type_usage(self): # This test verifies that we are using .type == "cuda" in subclasses # by checking the source code of the relevant files - from pepe.embedders.esm_embedder import ESMEmbedder import inspect + + from pepe.embedders.esm_embedder import ESMEmbedder + source = inspect.getsource(ESMEmbedder._initialize_model) self.assertIn('self.device.type == "cuda"', source) - - from pepe.embedders.huggingface_embedder import Antiberta2Embedder, ESM2Embedder + import inspect + + from pepe.embedders.huggingface_embedder import Antiberta2Embedder, ESM2Embedder + source = inspect.getsource(ESM2Embedder._initialize_model) self.assertIn('self.device.type == "cuda"', source) - + source = inspect.getsource(Antiberta2Embedder._initialize_model) self.assertIn('self.device.type == "cuda"', source) + if __name__ == "__main__": unittest.main() diff --git a/src/tests/test_esmc_modes.py b/src/tests/test_esmc_modes.py index 977be45..9476169 100644 --- a/src/tests/test_esmc_modes.py +++ b/src/tests/test_esmc_modes.py @@ -1,4 +1,5 @@ """Smoke tests for ESMC embedding modes. Run with ESMC_TEST=1.""" + import os import sys import unittest @@ -29,7 +30,9 @@ def _embed(self, **kwargs): return pepe.embed(**defaults) def test_per_token(self): - sequences = {"seq1": "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALTNAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR"} + sequences = { + "seq1": "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALTNAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR" + } results = self._embed( sequences=sequences, extract_embeddings=["per_token"], diff --git a/src/tests/test_generic_hf_integration.py b/src/tests/test_generic_hf_integration.py index 3698470..2d5ca40 100644 --- a/src/tests/test_generic_hf_integration.py +++ b/src/tests/test_generic_hf_integration.py @@ -26,9 +26,10 @@ class TestGenericHFIntegration(unittest.TestCase): MODEL = "hf-internal-testing/tiny-random-BertModel" def test_tiny_bert_mean_pooled_end_to_end(self): - import pepe from transformers import AutoConfig + import pepe + config = AutoConfig.from_pretrained(self.MODEL) expected_hidden = config.hidden_size diff --git a/src/tests/test_generic_safeguards.py b/src/tests/test_generic_safeguards.py index fb8f700..4e0b250 100644 --- a/src/tests/test_generic_safeguards.py +++ b/src/tests/test_generic_safeguards.py @@ -107,7 +107,10 @@ def test_warns_when_split_long_sequences_disabled(self): with self.assertLogs("pepe.embedders.base_embedder", level="WARNING") as logs: self._build_embedder(split_long_sequences=False) self.assertTrue( - any("exceed the model's maximum allowed length" in msg for msg in logs.output) + any( + "exceed the model's maximum allowed length" in msg + for msg in logs.output + ) ) diff --git a/src/tests/test_model_selection.py b/src/tests/test_model_selection.py index 04f2284..8e29ac4 100644 --- a/src/tests/test_model_selection.py +++ b/src/tests/test_model_selection.py @@ -6,7 +6,6 @@ sys.path.insert(0, os.path.abspath("src")) -from pepe.model_selecter import select_model, report_model from pepe.model_errors import ( ESMCForkRequiredError, GatedModelError, @@ -16,6 +15,7 @@ RemoteCodeRequiredError, UnsupportedArchitectureError, ) +from pepe.model_selecter import report_model, select_model def _config(model_type): @@ -161,9 +161,7 @@ def test_gated_esmc_raises_fork_error(self): def test_trust_remote_code_skips_remote_code_error_when_flag_set(self): with patch( "transformers.AutoConfig.from_pretrained", - side_effect=ValueError( - "Pass `trust_remote_code=True` to load this model." - ), + side_effect=ValueError("Pass `trust_remote_code=True` to load this model."), ): with self.assertRaises(ModelSelectionError): select_model(self.MODEL, trust_remote_code=True) @@ -178,7 +176,9 @@ def test_bare_esm2_name(self): # Must not touch the network: AutoConfig would raise if called. with patch( "transformers.AutoConfig.from_pretrained", - side_effect=AssertionError("AutoConfig should not be called for bare names"), + side_effect=AssertionError( + "AutoConfig should not be called for bare names" + ), ): self.assertIs(select_model("esm2_t6_8M_UR50D"), ESM2Embedder) @@ -187,7 +187,9 @@ def test_bare_esm1_name(self): with patch( "transformers.AutoConfig.from_pretrained", - side_effect=AssertionError("AutoConfig should not be called for bare names"), + side_effect=AssertionError( + "AutoConfig should not be called for bare names" + ), ): self.assertIs(select_model("esm1b_t33_650M_UR50S"), ESMEmbedder) @@ -201,12 +203,11 @@ def test_select_model_forwards_trust_remote_code_to_autoconfig(self): with patch("transformers.AutoConfig.from_pretrained") as mock_config: mock_config.return_value = _config("bert") select_model("someuser/model-x", trust_remote_code=True) - mock_config.assert_called_once_with( - "someuser/model-x", trust_remote_code=True - ) + mock_config.assert_called_once_with("someuser/model-x", trust_remote_code=True) def test_generic_embedder_passes_trust_remote_code_to_from_pretrained(self): import torch + from pepe.embedders.huggingface_embedder import GenericHuggingFaceEmbedder recorded = {"tokenizer": None, "model": None} @@ -270,9 +271,7 @@ def test_report_names_embedder_and_flags_subword_tokenizer(self): ), patch( "pepe.model_selecter.select_model", return_value=MagicMock(__name__="GenericHuggingFaceEmbedder"), - ), patch( - "pepe.utils.is_character_tokenizer", return_value=False - ): + ), patch("pepe.utils.is_character_tokenizer", return_value=False): buf = StringIO() with patch("sys.stdout", buf): report_model("someuser/protbert", trust_remote_code=True) diff --git a/src/tests/test_reconstruct_mean_pooled.py b/src/tests/test_reconstruct_mean_pooled.py index 055477f..942c5a0 100644 --- a/src/tests/test_reconstruct_mean_pooled.py +++ b/src/tests/test_reconstruct_mean_pooled.py @@ -12,6 +12,7 @@ asserts structure (count/shape/non-zero) and, above all, that reconstruction does not raise. """ + import json import os import sys diff --git a/src/tests/test_splitting.py b/src/tests/test_splitting.py index ef74382..1090a41 100644 --- a/src/tests/test_splitting.py +++ b/src/tests/test_splitting.py @@ -1,16 +1,17 @@ import os +import shutil +import subprocess import sys import unittest -import torch + import numpy as np -import subprocess -import shutil # Add src to sys.path sys.path.append(os.path.abspath("src")) from pepe.model_selecter import select_model + class TestSplittingIntegrated(unittest.TestCase): @classmethod def setUpClass(cls): @@ -19,12 +20,12 @@ def setUpClass(cls): os.makedirs(cls.test_dir, exist_ok=True) cls.fasta_path = os.path.join(cls.test_dir, "long.fasta") cls.short_fasta_path = os.path.join(cls.test_dir, "short.fasta") - + # 300 AAs -> 300 + 2 special tokens = 302 tokens ( > 256) long_seq = "M" * 300 with open(cls.fasta_path, "w") as f: f.write(f">long_prot\n{long_seq}\n") - + short_seq = "M" * 50 with open(cls.short_fasta_path, "w") as f: f.write(f">short_prot\n{short_seq}\n") @@ -36,8 +37,7 @@ def tearDownClass(cls): def test_library_reconstruction(self): """Test that the library correctly reconstructs embeddings in memory.""" - from pepe.model_selecter import select_model - + class DummyArgs: def __init__(self, fasta, split=True): self.fasta_path = fasta @@ -67,27 +67,26 @@ def __init__(self, fasta, split=True): args = DummyArgs(self.fasta_path, split=True) EmbedderClass = select_model(args.model_name) embedder = EmbedderClass(args) - + self.assertTrue(len(embedder.chunks_mapping) > 0) embedder.embed() - + # Verify reconstruction self.assertEqual(len(embedder.sequence_labels), 1) self.assertEqual(embedder.sequence_labels[0], "long_prot") - + per_token = embedder.per_token["output_data"][embedder.layers[0]][0] - + # Calculate expected length: sequence length + special tokens count - special_tokens_count = len(embedder.tokenizer.encode("", add_special_tokens=True)) + special_tokens_count = len( + embedder.tokenizer.encode("", add_special_tokens=True) + ) expected_len = 300 + special_tokens_count self.assertEqual(per_token.shape[0], expected_len) - - def test_force_split_length(self): """Test that the library overrides model defaults when force_split_length is supplied.""" - from pepe.model_selecter import select_model - + class DummyArgs: def __init__(self, fasta, split=True): self.fasta_path = fasta @@ -116,90 +115,113 @@ def __init__(self, fasta, split=True): args = DummyArgs(self.fasta_path, split=True) EmbedderClass = select_model(args.model_name) embedder = EmbedderClass(args) - + self.assertTrue(len(embedder.chunks_mapping) > 0) - + # Since force_split_length is 150, and our protein is 300 characters, it should be chunked into more than 2 pieces! - self.assertTrue(len(embedder.chunks_mapping['long_prot']) > 2) + self.assertTrue(len(embedder.chunks_mapping["long_prot"]) > 2) embedder.embed() - + # Verify reconstruction self.assertEqual(len(embedder.sequence_labels), 1) self.assertEqual(embedder.sequence_labels[0], "long_prot") - + per_token = embedder.per_token["output_data"][embedder.layers[0]][0] - + # Calculate expected length: sequence length + special tokens count - special_tokens_count = len(embedder.tokenizer.encode("", add_special_tokens=True)) + special_tokens_count = len( + embedder.tokenizer.encode("", add_special_tokens=True) + ) expected_len = 300 + special_tokens_count self.assertEqual(per_token.shape[0], expected_len) - - def test_cli_reconstruction_no_streaming(self): """Test that CLI with streaming_output=False reconstructs files on disk.""" out_dir = os.path.join(self.test_dir, "cli_no_streaming") cmd = [ - sys.executable, "-m", "pepe", - "--model_name", "alchemab/antiberta2-cssp", - "--fasta_path", self.fasta_path, - "--output_path", out_dir, + sys.executable, + "-m", + "pepe", + "--model_name", + "alchemab/antiberta2-cssp", + "--fasta_path", + self.fasta_path, + "--output_path", + out_dir, "--split_long_sequences", - "--split_overlap", "50", - "--device", "cpu", - "--streaming_output", "False", - "--extract_embeddings", "per_token" + "--split_overlap", + "50", + "--device", + "cpu", + "--streaming_output", + "False", + "--extract_embeddings", + "per_token", ] - + env = os.environ.copy() env["PYTHONPATH"] = "src" subprocess.run(cmd, env=env, check=True) - + # Index should show original labels idx_path = os.path.join(out_dir, "antiberta2-cssp", "long_idx.csv") with open(idx_path, "r") as f: lines = f.readlines() - + # header + 1 sequence self.assertEqual(len(lines), 2) self.assertIn("long_prot", lines[1]) - + # Check tensor shape on disk (non-streaming uses .npy) - npy_path = os.path.join(out_dir, "antiberta2-cssp", "per_token", "long_antiberta2-cssp_per_token_layer_16.npy") + npy_path = os.path.join( + out_dir, + "antiberta2-cssp", + "per_token", + "long_antiberta2-cssp_per_token_layer_16.npy", + ) data = np.load(npy_path) self.assertEqual(data.shape[0], 1) # Check the sequence length (reconstructed) # For AntiBERTa2, 300 AA + [CLS] + [SEP] = 302 self.assertEqual(data[0].shape[0], 302) - def test_cli_streaming_remains_chunked(self): """Test that CLI with streaming_output=True (default) exports chunks.""" out_dir = os.path.join(self.test_dir, "cli_streaming") cmd = [ - sys.executable, "-m", "pepe", - "--model_name", "alchemab/antiberta2-cssp", - "--fasta_path", self.fasta_path, - "--output_path", out_dir, + sys.executable, + "-m", + "pepe", + "--model_name", + "alchemab/antiberta2-cssp", + "--fasta_path", + self.fasta_path, + "--output_path", + out_dir, "--split_long_sequences", - "--split_overlap", "50", - "--device", "cpu", - "--streaming_output", "True", - "--extract_embeddings", "per_token" + "--split_overlap", + "50", + "--device", + "cpu", + "--streaming_output", + "True", + "--extract_embeddings", + "per_token", ] - + env = os.environ.copy() env["PYTHONPATH"] = "src" subprocess.run(cmd, env=env, check=True) - + # Index should show chunks idx_path = os.path.join(out_dir, "antiberta2-cssp", "long_idx.csv") with open(idx_path, "r") as f: content = f.read() - + self.assertIn("long_prot_chunk_0", content) self.assertIn("long_prot_chunk_1", content) + if __name__ == "__main__": unittest.main() diff --git a/src/tests/test_streaming_roundtrip.py b/src/tests/test_streaming_roundtrip.py index 49cc5e2..ad10981 100644 --- a/src/tests/test_streaming_roundtrip.py +++ b/src/tests/test_streaming_roundtrip.py @@ -1,4 +1,5 @@ """Regression tests: streaming disk output must match in-memory results.""" + import glob import os import sys @@ -53,8 +54,7 @@ def in_memory_to_numpy(output_data, output_type, layer=LAYER, num_heads=NUM_HEAD return _stack_tensors(output_data) if output_type == "attention_head": return { - head: _stack_tensors(output_data[layer][head]) - for head in range(num_heads) + head: _stack_tensors(output_data[layer][head]) for head in range(num_heads) } return _stack_tensors(output_data[layer]) @@ -66,23 +66,37 @@ def _output_dir(output_path, output_type): def load_streaming_array(output_path, output_type, layer=LAYER, head=None): output_dir = _output_dir(output_path, output_type) if output_type == "attention_model": - pattern = os.path.join(output_dir, f"{EXPERIMENT_NAME}_{MODEL_NAME}_attention_model.npy") - files = [pattern] if os.path.exists(pattern) else glob.glob(os.path.join(output_dir, "*.npy")) + pattern = os.path.join( + output_dir, f"{EXPERIMENT_NAME}_{MODEL_NAME}_attention_model.npy" + ) + files = ( + [pattern] + if os.path.exists(pattern) + else glob.glob(os.path.join(output_dir, "*.npy")) + ) elif output_type == "attention_head": pattern = os.path.join( output_dir, f"{EXPERIMENT_NAME}_{MODEL_NAME}_attention_head_layer_{layer}_head_{head + 1}.npy", ) - files = [pattern] if os.path.exists(pattern) else glob.glob( - os.path.join(output_dir, f"*_layer_{layer}_head_{head + 1}.npy") + files = ( + [pattern] + if os.path.exists(pattern) + else glob.glob( + os.path.join(output_dir, f"*_layer_{layer}_head_{head + 1}.npy") + ) ) else: pattern = os.path.join( output_dir, f"{EXPERIMENT_NAME}_{MODEL_NAME}_{output_type}_layer_{layer}.npy", ) - files = [pattern] if os.path.exists(pattern) else glob.glob( - os.path.join(output_dir, f"*_{output_type}_layer_{layer}.npy") + files = ( + [pattern] + if os.path.exists(pattern) + else glob.glob( + os.path.join(output_dir, f"*_{output_type}_layer_{layer}.npy") + ) ) assert len(files) == 1, f"Expected one file for {output_type}, found {files}" return np.load(files[0]) @@ -103,7 +117,9 @@ def _assert_allclose(reference, streaming, output_type): np.testing.assert_allclose( reference[head], streaming[head], rtol=1e-4, atol=1e-4 ) - assert np.any(streaming[head] != 0), f"{output_type} head {head} is all zeros" + assert np.any(streaming[head] != 0), ( + f"{output_type} head {head} is all zeros" + ) else: np.testing.assert_allclose(reference, streaming, rtol=1e-4, atol=1e-4) assert np.any(streaming != 0), f"{output_type} is all zeros" diff --git a/src/tests/test_sync_arguments.py b/src/tests/test_sync_arguments.py index fc07b0e..4d32aba 100644 --- a/src/tests/test_sync_arguments.py +++ b/src/tests/test_sync_arguments.py @@ -1,4 +1,5 @@ """Guard against drift between CLI flags, embed() API, and embedder args.""" + import argparse import ast import inspect @@ -40,7 +41,9 @@ def capture(self, args=None, namespace=None): def _embed_params(): - return {name for name in inspect.signature(api.embed).parameters if name != "kwargs"} + return { + name for name in inspect.signature(api.embed).parameters if name != "kwargs" + } def _args_dict_keys(): From 29a3c7247355b115ceeabedf9654f78b2c16bbab Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 19:25:56 +0200 Subject: [PATCH 03/17] Drop Python 3.8 and 3.9 support (minimum is now 3.10) Raise the floor to Python 3.10 so packaging metadata, CI, and the type/lint tooling all agree (mypy and Ruff already targeted 3.10): - pyproject.toml / setup.py: requires-python >=3.10, drop 3.8/3.9 classifiers. - .github/workflows/test.yml: unit matrix now [3.10, 3.11]. - .github/conda/meta.yaml: python >=3.10 (host + run). - pyproject [tool.ruff]: target-version = "py310". - CHANGELOG: note the removal under [Unreleased]. Bumping the Ruff target to py310 lets the formatter use parenthesized `with` statements (3.10+ syntax); reformats four files accordingly. No behavior change. Co-Authored-By: Claude Fable 5 --- .github/conda/meta.yaml | 4 ++-- .github/workflows/test.yml | 2 +- CHANGELOG.md | 5 +++++ pyproject.toml | 9 +++------ setup.py | 4 +--- src/pepe/embedders/base_embedder.py | 11 +++++++---- src/tests/test_device_logic.py | 24 +++++++++++++++--------- src/tests/test_generic_safeguards.py | 20 +++++++++++++------- src/tests/test_model_selection.py | 20 ++++++++++++-------- 9 files changed, 59 insertions(+), 40 deletions(-) diff --git a/.github/conda/meta.yaml b/.github/conda/meta.yaml index b3b8d1b..1be1fe0 100644 --- a/.github/conda/meta.yaml +++ b/.github/conda/meta.yaml @@ -12,12 +12,12 @@ build: requirements: host: - - python >=3.8 + - python >=3.10 - pip - setuptools - wheel run: - - python >=3.8 + - python >=3.10 - pytorch >=1.9.0 - transformers >=4.20.0 - sentencepiece diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6f0c336..aa397d0 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -45,7 +45,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.9", "3.10", "3.11"] + python-version: ["3.10", "3.11"] steps: - uses: actions/checkout@v5 - uses: actions/setup-python@v6 diff --git a/CHANGELOG.md b/CHANGELOG.md index 400e847..f80750a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,11 @@ truth that drives publishing). ## [Unreleased] +### Removed +- Dropped support for Python 3.8 and 3.9. The minimum supported version is now + **3.10** (`requires-python >=3.10`), matching the mypy/Ruff target. The CI unit + matrix, packaging classifiers, and the conda recipe were updated accordingly. + ## [1.4.0] - 2026-07-07 ### Added diff --git a/pyproject.toml b/pyproject.toml index 7bb99b0..4eb4f34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,15 +11,13 @@ authors = [ description = "Pipeline for Easy Protein Embedding - Extract embeddings and attention matrices from protein sequences" readme = "README.md" license = {text = "MIT"} -requires-python = ">=3.8" +requires-python = ">=3.10" classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Science/Research", "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Topic :: Scientific/Engineering :: Bio-Informatics", @@ -93,9 +91,8 @@ check_untyped_defs = true [tool.ruff] line-length = 88 -# Match the declared floor (requires-python >= 3.8). NOTE: mypy targets 3.10 and -# the CI matrix starts at 3.9 — decide separately whether to drop 3.8 support. -target-version = "py38" +# Matches requires-python (>= 3.10) and the mypy python_version. +target-version = "py310" src = ["src"] # Notebooks follow their own conventions (imports mid-cell, etc.); don't lint them. extend-exclude = ["notebooks", "*.ipynb"] diff --git a/setup.py b/setup.py index 2e1e099..9e6f2b4 100644 --- a/setup.py +++ b/setup.py @@ -41,14 +41,12 @@ def read_readme(): "Operating System :: Linux", "Operating System :: macOS", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Topic :: Scientific/Engineering :: Bio-Informatics", "Topic :: Scientific/Engineering :: Artificial Intelligence", ], - python_requires=">=3.8", + python_requires=">=3.10", install_requires=[ "torch>=1.9.0", "transformers>=4.20.0", diff --git a/src/pepe/embedders/base_embedder.py b/src/pepe/embedders/base_embedder.py index d12a6ed..3599945 100644 --- a/src/pepe/embedders/base_embedder.py +++ b/src/pepe/embedders/base_embedder.py @@ -429,10 +429,13 @@ def embed(self) -> None: if resume_info: logger.info(f"Resuming from checkpoint: {resume_info}") - with alive_bar( - len(self.sequences), - title=f"{self.model_name}: Generating embeddings ...", - ) as bar, torch.no_grad(): + with ( + alive_bar( + len(self.sequences), + title=f"{self.model_name}: Generating embeddings ...", + ) as bar, + torch.no_grad(), + ): offset = 0 for ( labels, diff --git a/src/tests/test_device_logic.py b/src/tests/test_device_logic.py index 21d7073..8ec33e7 100644 --- a/src/tests/test_device_logic.py +++ b/src/tests/test_device_logic.py @@ -36,9 +36,11 @@ def setUp(self): @patch("torch.cuda.is_available", return_value=True) def test_base_device_logic_gpu_available(self, mock_cuda): # Mocking file system calls and substring loading - with patch("os.path.exists", return_value=True), patch( - "os.makedirs" - ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): + with ( + patch("os.path.exists", return_value=True), + patch("os.makedirs"), + patch.object(BaseEmbedder, "_load_substrings", return_value=None), + ): # 1. GPU requested and available self.args.device = "cuda" base = BaseEmbedder(self.args) @@ -51,9 +53,11 @@ def test_base_device_logic_gpu_available(self, mock_cuda): @patch("torch.cuda.is_available", return_value=True) def test_base_device_logic_gpu_specific(self, mock_cuda): - with patch("os.path.exists", return_value=True), patch( - "os.makedirs" - ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): + with ( + patch("os.path.exists", return_value=True), + patch("os.makedirs"), + patch.object(BaseEmbedder, "_load_substrings", return_value=None), + ): # GPU cuda:1 requested and available self.args.device = "cuda:1" base = BaseEmbedder(self.args) @@ -62,9 +66,11 @@ def test_base_device_logic_gpu_specific(self, mock_cuda): @patch("torch.cuda.is_available", return_value=False) def test_base_device_logic_gpu_not_available(self, mock_cuda): - with patch("os.path.exists", return_value=True), patch( - "os.makedirs" - ), patch.object(BaseEmbedder, "_load_substrings", return_value=None): + with ( + patch("os.path.exists", return_value=True), + patch("os.makedirs"), + patch.object(BaseEmbedder, "_load_substrings", return_value=None), + ): # GPU requested but NOT available -> should fallback to CPU self.args.device = "cuda" base = BaseEmbedder(self.args) diff --git a/src/tests/test_generic_safeguards.py b/src/tests/test_generic_safeguards.py index 4e0b250..eec7f2f 100644 --- a/src/tests/test_generic_safeguards.py +++ b/src/tests/test_generic_safeguards.py @@ -89,13 +89,19 @@ def tearDownClass(cls): def _build_embedder(self, split_long_sequences): args = _make_args(self.fasta_path, split_long_sequences=split_long_sequences) model, tokenizer = _mock_model_tokenizer(max_position_embeddings=512) - with patch.object( - GenericHuggingFaceEmbedder, - "_initialize_model", - return_value=(model, tokenizer, 8, 2, 64), - ), patch.object( - GenericHuggingFaceEmbedder, "_load_data", return_value=(MagicMock(), 512) - ), patch.object(GenericHuggingFaceEmbedder, "_set_output_objects"): + with ( + patch.object( + GenericHuggingFaceEmbedder, + "_initialize_model", + return_value=(model, tokenizer, 8, 2, 64), + ), + patch.object( + GenericHuggingFaceEmbedder, + "_load_data", + return_value=(MagicMock(), 512), + ), + patch.object(GenericHuggingFaceEmbedder, "_set_output_objects"), + ): return GenericHuggingFaceEmbedder(args) def test_splits_when_split_long_sequences_enabled(self): diff --git a/src/tests/test_model_selection.py b/src/tests/test_model_selection.py index 8e29ac4..fe98ac2 100644 --- a/src/tests/test_model_selection.py +++ b/src/tests/test_model_selection.py @@ -264,14 +264,18 @@ def test_report_names_embedder_and_flags_subword_tokenizer(self): mock_tokenizer.model_max_length = 512 mock_tokenizer.encode.return_value = [1, 2, 3] - with patch( - "transformers.AutoConfig.from_pretrained", return_value=mock_config - ), patch( - "transformers.AutoTokenizer.from_pretrained", return_value=mock_tokenizer - ), patch( - "pepe.model_selecter.select_model", - return_value=MagicMock(__name__="GenericHuggingFaceEmbedder"), - ), patch("pepe.utils.is_character_tokenizer", return_value=False): + with ( + patch("transformers.AutoConfig.from_pretrained", return_value=mock_config), + patch( + "transformers.AutoTokenizer.from_pretrained", + return_value=mock_tokenizer, + ), + patch( + "pepe.model_selecter.select_model", + return_value=MagicMock(__name__="GenericHuggingFaceEmbedder"), + ), + patch("pepe.utils.is_character_tokenizer", return_value=False), + ): buf = StringIO() with patch("sys.stdout", buf): report_model("someuser/protbert", trust_remote_code=True) From c31fbd031380773898baa7a2a835cb7e35d3227d Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 20:15:02 +0200 Subject: [PATCH 04/17] Add GHCR Docker publishing to test and main release workflows. Ship a CPU Docker image alongside PyPI/Conda releases so users can run PEPE without a local Python setup. Co-authored-by: Cursor --- .dockerignore | 32 +++++++++++ .../workflows/publish-main-branch-trusted.yml | 56 +++++++++++++++++++ .../workflows/publish-test-branch-trusted.yml | 48 ++++++++++++++++ Dockerfile | 21 +++++++ 4 files changed, 157 insertions(+) create mode 100644 .dockerignore create mode 100644 Dockerfile diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..dc1ac14 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,32 @@ +.git +.github +.claude +.venv +.venv-* +venv +env +__pycache__ +*.py[cod] +*.egg-info +dist +build +.pytest_cache +.mypy_cache +.coverage +htmlcov +*.pt +*.pth +data +ignore +ignore.py +test_library_out* +test_verify_* +readme_advanced_results +readme_example_results +src/tests/test_files/test_output +*.tsv +*.code-workspace +.vscode +.env +plan.md +ROADMAP.md diff --git a/.github/workflows/publish-main-branch-trusted.yml b/.github/workflows/publish-main-branch-trusted.yml index 483185d..b1fb348 100644 --- a/.github/workflows/publish-main-branch-trusted.yml +++ b/.github/workflows/publish-main-branch-trusted.yml @@ -159,6 +159,13 @@ jobs: ```bash pip install pepe_cli-${{ steps.get_version.outputs.version }}-py3-none-any.whl ``` + + ### From GHCR (Docker) + ```bash + docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} + docker run --rm -v "$PWD:/data" ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} \ + --model_name esm2_t6_8M_UR50D --fasta_path /data/sequences.fasta --output_path /data/out --extract_embeddings mean_pooled + ``` files: | dist/*.whl @@ -207,3 +214,52 @@ jobs: fi echo "Found package: $PACKAGE_PATH" anaconda -t $ANACONDA_TOKEN upload --label main $PACKAGE_PATH + + publish-docker: + runs-on: ubuntu-latest + needs: build-and-publish-main + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v5 + + - name: Get version + id: version + run: | + VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') + echo "version=$VERSION" >> $GITHUB_OUTPUT + echo "image=ghcr.io/csi-greifflab/pepe-cli" >> $GITHUB_OUTPUT + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to GHCR + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and push Docker image + uses: docker/build-push-action@v6 + with: + context: . + push: true + tags: | + ghcr.io/csi-greifflab/pepe-cli:${{ steps.version.outputs.version }} + ghcr.io/csi-greifflab/pepe-cli:latest + labels: | + org.opencontainers.image.source=${{ github.repository }} + org.opencontainers.image.version=${{ steps.version.outputs.version }} + org.opencontainers.image.title=pepe-cli + + - name: Verify container CLI + run: | + docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.version.outputs.version }} + docker run --rm ghcr.io/csi-greifflab/pepe-cli:${{ steps.version.outputs.version }} --help + + - name: Output Docker pull command + run: | + echo "Docker image published to GHCR:" + echo "docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.version.outputs.version }}" diff --git a/.github/workflows/publish-test-branch-trusted.yml b/.github/workflows/publish-test-branch-trusted.yml index d564ebe..8d14742 100644 --- a/.github/workflows/publish-test-branch-trusted.yml +++ b/.github/workflows/publish-test-branch-trusted.yml @@ -143,3 +143,51 @@ jobs: fi echo "Found package: $PACKAGE_PATH" anaconda -t $ANACONDA_TOKEN upload --label test $PACKAGE_PATH + + publish-docker-test: + runs-on: ubuntu-latest + needs: build-and-publish-test + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v5 + + - name: Sync version for Docker build + run: | + CURRENT_VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') + sed -i "s/__version__ = \"$CURRENT_VERSION\"/__version__ = \"${{ needs.build-and-publish-test.outputs.test_version }}\"/" src/pepe/__init__.py + grep "__version__" src/pepe/__init__.py + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to GHCR + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and push Docker image + uses: docker/build-push-action@v6 + with: + context: . + push: true + tags: | + ghcr.io/csi-greifflab/pepe-cli:${{ needs.build-and-publish-test.outputs.test_version }} + ghcr.io/csi-greifflab/pepe-cli:test + labels: | + org.opencontainers.image.source=${{ github.repository }} + org.opencontainers.image.version=${{ needs.build-and-publish-test.outputs.test_version }} + org.opencontainers.image.title=pepe-cli + + - name: Verify container CLI + run: | + docker pull ghcr.io/csi-greifflab/pepe-cli:${{ needs.build-and-publish-test.outputs.test_version }} + docker run --rm ghcr.io/csi-greifflab/pepe-cli:${{ needs.build-and-publish-test.outputs.test_version }} --help + + - name: Output Docker pull command + run: | + echo "Test Docker image published to GHCR:" + echo "docker pull ghcr.io/csi-greifflab/pepe-cli:${{ needs.build-and-publish-test.outputs.test_version }}" diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..1cf786f --- /dev/null +++ b/Dockerfile @@ -0,0 +1,21 @@ +FROM python:3.10-slim + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PIP_NO_CACHE_DIR=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 + +WORKDIR /app + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git \ + libgomp1 \ + && rm -rf /var/lib/apt/lists/* + +COPY pyproject.toml setup.py MANIFEST.in README.md LICENSE requirements.txt ./ +COPY src/ ./src/ + +RUN pip install . + +WORKDIR /data +ENTRYPOINT ["pepe"] From 6325e62e8f40fd703a4f2f516ca51dcf39904ae5 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:30:30 +0000 Subject: [PATCH 05/17] Address review: remove git dep, exclude tests, add --user flag, defer release to after Docker publish --- .dockerignore | 1 + .../workflows/publish-main-branch-trusted.yml | 104 +++++++++++------- Dockerfile | 1 - 3 files changed, 64 insertions(+), 42 deletions(-) diff --git a/.dockerignore b/.dockerignore index dc1ac14..fd5f5ea 100644 --- a/.dockerignore +++ b/.dockerignore @@ -30,3 +30,4 @@ src/tests/test_files/test_output .env plan.md ROADMAP.md +src/tests diff --git a/.github/workflows/publish-main-branch-trusted.yml b/.github/workflows/publish-main-branch-trusted.yml index b1fb348..2cd83e2 100644 --- a/.github/workflows/publish-main-branch-trusted.yml +++ b/.github/workflows/publish-main-branch-trusted.yml @@ -133,48 +133,12 @@ jobs: echo "🎉 Package verification completed successfully!" echo "Package $PACKAGE_NAME version $CURRENT_VERSION is now available on PyPI!" - - name: Get version for release - id: get_version - run: | - VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') - echo "version=$VERSION" >> $GITHUB_OUTPUT - - - name: Create GitHub Release - uses: softprops/action-gh-release@v3 + - name: Upload dist artifacts + uses: actions/upload-artifact@v4 with: - tag_name: v${{ steps.get_version.outputs.version }} - name: Release v${{ steps.get_version.outputs.version }} - body: | - Release of pepe-cli version ${{ steps.get_version.outputs.version }} - - ## Installation - - ### From PyPI (Recommended) - ```bash - pip install pepe-cli==${{ steps.get_version.outputs.version }} - ``` - - ### From GitHub Release - Download the wheel file and install: - ```bash - pip install pepe_cli-${{ steps.get_version.outputs.version }}-py3-none-any.whl - ``` - - ### From GHCR (Docker) - ```bash - docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} - docker run --rm -v "$PWD:/data" ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} \ - --model_name esm2_t6_8M_UR50D --fasta_path /data/sequences.fasta --output_path /data/out --extract_embeddings mean_pooled - ``` - - files: | - dist/*.whl - dist/*.tar.gz - draft: false - prerelease: false - generate_release_notes: true - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + name: dist-packages + path: dist/ + retention-days: 1 publish-conda: runs-on: ubuntu-latest @@ -263,3 +227,61 @@ jobs: run: | echo "Docker image published to GHCR:" echo "docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.version.outputs.version }}" + + create-release: + runs-on: ubuntu-latest + needs: [build-and-publish-main, publish-docker] + permissions: + contents: write + steps: + - uses: actions/checkout@v5 + + - name: Get version for release + id: get_version + run: | + VERSION=$(grep -m1 '__version__ = ' src/pepe/__init__.py | sed 's/__version__ = "\(.*\)"/\1/') + echo "version=$VERSION" >> $GITHUB_OUTPUT + + - name: Download dist artifacts + uses: actions/download-artifact@v4 + with: + name: dist-packages + path: dist/ + + - name: Create GitHub Release + uses: softprops/action-gh-release@v3 + with: + tag_name: v${{ steps.get_version.outputs.version }} + name: Release v${{ steps.get_version.outputs.version }} + body: | + Release of pepe-cli version ${{ steps.get_version.outputs.version }} + + ## Installation + + ### From PyPI (Recommended) + ```bash + pip install pepe-cli==${{ steps.get_version.outputs.version }} + ``` + + ### From GitHub Release + Download the wheel file and install: + ```bash + pip install pepe_cli-${{ steps.get_version.outputs.version }}-py3-none-any.whl + ``` + + ### From GHCR (Docker) + ```bash + docker pull ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} + docker run --rm -v "$PWD:/data" --user "$(id -u):$(id -g)" \ + ghcr.io/csi-greifflab/pepe-cli:${{ steps.get_version.outputs.version }} \ + --model_name esm2_t6_8M_UR50D --fasta_path /data/sequences.fasta --output_path /data/out --extract_embeddings mean_pooled + ``` + + files: | + dist/*.whl + dist/*.tar.gz + draft: false + prerelease: false + generate_release_notes: true + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/Dockerfile b/Dockerfile index 1cf786f..5f47dd4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,7 +8,6 @@ ENV PYTHONDONTWRITEBYTECODE=1 \ WORKDIR /app RUN apt-get update && apt-get install -y --no-install-recommends \ - git \ libgomp1 \ && rm -rf /var/lib/apt/lists/* From e84f3a018872aaa7311d4cdff710cd9333d53b9c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 18:34:58 +0000 Subject: [PATCH 06/17] Fix trailing whitespace in publish-main-branch-trusted.yml --- .github/workflows/publish-main-branch-trusted.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/publish-main-branch-trusted.yml b/.github/workflows/publish-main-branch-trusted.yml index f3fee54..1ae86ce 100644 --- a/.github/workflows/publish-main-branch-trusted.yml +++ b/.github/workflows/publish-main-branch-trusted.yml @@ -132,7 +132,7 @@ jobs: echo "🎉 Package verification completed successfully!" echo "Package $PACKAGE_NAME version $CURRENT_VERSION is now available on PyPI!" - + - name: Upload dist artifacts uses: actions/upload-artifact@v4 with: From 8e1ab5e41517288ebeaf67afa84c14299baaa373 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 21:00:32 +0200 Subject: [PATCH 07/17] Add METL 1D embeddings via optional metl-pretrained extra. Enables metl-*-1d model dispatch, METLEmbedder/METLDataset, typed errors for missing package and 3D models, CI integration test behind METL_TEST, and docs in CHANGELOG (1D-only; no logits/attention). Co-authored-by: Cursor --- .github/workflows/test.yml | 6 + CHANGELOG.md | 17 ++ pyproject.toml | 4 + setup.py | 4 + src/pepe/embedders/metl_embedder.py | 253 ++++++++++++++++++++++++++++ src/pepe/model_errors.py | 8 + src/pepe/model_selecter.py | 30 +++- src/pepe/utils.py | 75 +++++++++ src/tests/test_metl_dispatch.py | 78 +++++++++ src/tests/test_metl_integration.py | 41 +++++ 10 files changed, 515 insertions(+), 1 deletion(-) create mode 100644 src/pepe/embedders/metl_embedder.py create mode 100644 src/tests/test_metl_dispatch.py create mode 100644 src/tests/test_metl_integration.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 583b011..4a13947 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -40,6 +40,7 @@ jobs: src/tests/test_device_logic.py \ src/tests/test_load_layers_default.py \ src/tests/test_model_selection.py \ + src/tests/test_metl_dispatch.py \ src/tests/test_generic_safeguards.py integration: @@ -69,3 +70,8 @@ jobs: run: | GENERIC_HF_TEST=1 python -m pytest -v \ src/tests/test_generic_hf_integration.py + - name: Run METL integration test + run: | + pip install "git+https://github.com/gitter-lab/metl-pretrained.git" + METL_TEST=1 python -m pytest -v \ + src/tests/test_metl_integration.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c7f5cd2..f83ccbd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,23 @@ truth that drives publishing). ## [Unreleased] +### Added +- METL 1D protein embeddings via optional `metl-pretrained` backend: install with + `pip install pepe-cli[metl]` and use model identifiers such as `metl-g-20m-1d` + (and other `metl-*-1d` names). Dispatch lives in `model_selecter.py`; embedding + is handled by `METLEmbedder` with `METLDataset` tokenization. +- `[metl]` and `[esm]` optional dependency extras in `pyproject.toml` and + `setup.py`. +- Typed errors `METLPackageRequiredError` and `METL3DNotSupportedError` when + METL is requested without the extra or when a 3D METL model id is used. +- Unit tests for METL dispatch and an optional integration test (`METL_TEST=1`) + wired in CI (`.github/workflows/test.yml`). + +### Limitations +- METL support is **1D embeddings only** (`per_token`, `mean_pooled`, + `substring_pooled`). Logits and attention outputs are not available. 3D METL + model identifiers are rejected with a clear error. + ## [1.3.0] - 2026-07-07 ### Added diff --git a/pyproject.toml b/pyproject.toml index 282daaa..5f672e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,10 @@ dependencies = [ "rjieba", ] +[project.optional-dependencies] +metl = ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git"] +esm = ["fair-esm"] + [project.urls] Homepage = "https://github.com/csi-greifflab/pepe-cli" "Bug Reports" = "https://github.com/csi-greifflab/pepe-cli/issues" diff --git a/setup.py b/setup.py index bfc6e58..d748314 100644 --- a/setup.py +++ b/setup.py @@ -49,6 +49,10 @@ def read_readme(): "alive_progress", "rjieba", ], + extras_require={ + "metl": ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git"], + "esm": ["fair-esm"], + }, entry_points={ "console_scripts": [ f"pepe = {__module_name__}.__main__:main", diff --git a/src/pepe/embedders/metl_embedder.py b/src/pepe/embedders/metl_embedder.py new file mode 100644 index 0000000..adfe8d3 --- /dev/null +++ b/src/pepe/embedders/metl_embedder.py @@ -0,0 +1,253 @@ +import logging +import torch +import pepe.utils +from pepe.embedders.base_embedder import BaseEmbedder +from pepe.model_errors import ModelSelectionError, METLPackageRequiredError + + +def _import_metl(): + """Lazy import of metl-pretrained to avoid loading issues.""" + try: + import metl + + return metl + except ImportError as e: + raise METLPackageRequiredError( + "METL models require metl-pretrained. Install with: " + "pip install 'pepe-cli[metl]' or " + "pip install git+https://github.com/gitter-lab/metl-pretrained.git" + ) from e + + +def resolve_metl_ident(model_name): + """Map a PEPE model name to a metl-pretrained identifier.""" + normalized = model_name.strip() + if normalized.lower() in ("gitter-lab/metl", "gitter-lab/metl-pretrained"): + raise ModelSelectionError( + "gitter-lab/METL is the HuggingFace wrapper and is not supported by PEPE. " + "Use a metl-pretrained identifier instead (e.g. metl-g-20m-1d)." + ) + return normalized.lower() + + +logger = logging.getLogger("src.embedders.metl_embedder") + + +class METLEmbedder(BaseEmbedder): + def __init__(self, args): + super().__init__(args) + if self.return_logits: + logger.warning( + "Warning: Logits are not supported for METL models. Setting to False." + ) + self.return_logits = False + if "logits" in self.output_types: + self.output_types.remove("logits") + if self.return_contacts: + logger.warning( + "Warning: Attention matrices are not supported for METL models. Setting to False." + ) + self.return_contacts = False + for output_type in ("attention_head", "attention_layer", "attention_model"): + if output_type in self.output_types: + self.output_types.remove(output_type) + + self.sequences = pepe.utils.fasta_to_dict(args.fasta_path) + self.num_sequences = len(self.sequences) + ( + self.model, + self.data_encoder, + self.num_heads, + self.num_layers, + self.embedding_size, + ) = self._initialize_model(self.model_link) + self.valid_tokens = set(self.data_encoder.chars[1:]) + self._check_max_input_length() + pepe.utils.check_input_tokens( + self.valid_tokens, + self.sequences, + self.model_name, + split_long_sequences=self.split_long_sequences, + ) + self.special_tokens = torch.tensor([0], device=self.device, dtype=torch.int8) + self.layers = self._load_layers(self.layers) + self._repr_outputs = {} + self._hook_handles = [] + self._register_repr_hooks() + self.data_loader, self.max_input_length = self._load_data( + self.sequences, self.substring_dict + ) + self._set_output_objects() + + def _initialize_model(self, model_link): + logger.info("Loading METL model...") + metl = _import_metl() + ident = resolve_metl_ident(model_link) + model, data_encoder = metl.get_from_ident(ident) + model.eval() + + tr_encoder = model.model.tr_encoder + num_tr_layers = len(tr_encoder.layers) + num_layers = num_tr_layers + embedding_size = getattr(model, "embedding_len", None) or model.model.embedding_len + num_heads = 1 + + if torch.cuda.is_available() and self.device.type == "cuda": + model = model.cuda() + logger.info("Transferred model to GPU") + else: + logger.info("No GPU available, using CPU") + + return model, data_encoder, num_heads, num_layers, embedding_size + + def _register_repr_hooks(self): + tr_encoder = self.model.model.tr_encoder + num_tr_layers = len(tr_encoder.layers) + + for layer in self.layers: + if layer == num_tr_layers: + + def make_hook(captured_layer): + def hook(_module, _input, output): + self._repr_outputs[captured_layer] = output + + return hook + + handle = tr_encoder.norm.register_forward_hook(make_hook(layer)) + else: + + def make_hook(captured_layer): + def hook(_module, _input, output): + self._repr_outputs[captured_layer] = output + + return hook + + handle = tr_encoder.layers[layer - 1].register_forward_hook( + make_hook(layer) + ) + self._hook_handles.append(handle) + + def _load_layers(self, layers): + if layers is None: + return [self.num_layers] + if not layers: + layers = [-1] + assert all(-(self.num_layers + 1) <= i <= self.num_layers for i in layers) + layers = [ + (i + self.num_layers + 1) % (self.num_layers + 1) for i in layers + ] + return layers + + def _load_data(self, sequences, substring_dict=None): + dataset = pepe.utils.METLDataset( + sequences, + substring_dict, + self.context, + self.data_encoder, + self.max_input_length, + ) + logger.info("Generating batches...") + batches = pepe.utils.TokenBudgetBatchSampler(dataset, self.batch_size) + data_loader = torch.utils.data.DataLoader( + dataset, batch_sampler=batches, collate_fn=dataset.safe_collate + ) + logger.info("Data loaded") + max_length = dataset.get_max_encoded_length() + return data_loader, max_length + + def _compute_outputs( + self, + model, + toks, + attention_mask, + return_embeddings, + return_contacts, + return_logits, + ): + self._repr_outputs.clear() + model(toks) + + if return_embeddings: + dtype = self._precision_to_dtype(self.precision, "torch") + representations = { + layer: self._repr_outputs[layer].to(dtype=dtype).cpu() + for layer in self.layers + } + torch.cuda.empty_cache() + else: + representations = None + + return None, representations, None + + def _get_model_max_allowed(self): + if hasattr(self, "force_split_length") and self.force_split_length is not None: + return self.force_split_length + return self.model.aa_seq_len + + def _check_max_input_length(self): + """Check max length without BOS/EOS adjustment (METL has no special tokens).""" + max_allowed = self._get_model_max_allowed() + if max_allowed is None: + return + + sequences_too_long = any( + len(s) > max_allowed for s in self.sequences.values() + ) + needs_splitting = ( + isinstance(self.max_input_length, int) + and self.max_input_length > max_allowed + ) or sequences_too_long + + if needs_splitting: + if self.split_long_sequences: + logger.info( + f"Splitting sequences because they exceed model limit ({max_allowed})." + ) + self._handle_sequence_splitting(max_allowed) + else: + logger.warning( + f"Warning: Sequences exceed the model's maximum allowed length ({max_allowed})." + ) + + def _handle_sequence_splitting(self, max_allowed): + """Split sequences without reserving space for BOS/EOS tokens.""" + new_sequences = {} + self.chunks_mapping = {} + chunk_size = max_allowed + overlap = self.split_overlap + + if chunk_size <= overlap: + logger.error( + f"chunk_size ({chunk_size}) must be greater than overlap ({overlap}). Disabling splitting." + ) + return + + for label, sequence in self.sequences.items(): + if len(sequence) <= chunk_size: + new_sequences[label] = sequence + continue + + self.original_sequences[label] = sequence + chunks = [] + start = 0 + chunk_idx = 0 + while start < len(sequence): + end = min(start + chunk_size, len(sequence)) + chunk_payload = sequence[start:end] + chunk_label = f"{label}_chunk_{chunk_idx}" + + new_sequences[chunk_label] = chunk_payload + self.chunk_payload_lengths[chunk_label] = len(chunk_payload) + chunks.append(chunk_label) + + if end == len(sequence): + break + start = end - overlap + chunk_idx += 1 + + self.chunks_mapping[label] = chunks + + self.sequences = new_sequences + if hasattr(self, "num_sequences"): + self.num_sequences = len(self.sequences) + self.max_input_length = chunk_size diff --git a/src/pepe/model_errors.py b/src/pepe/model_errors.py index 8921f67..b56b86b 100644 --- a/src/pepe/model_errors.py +++ b/src/pepe/model_errors.py @@ -29,6 +29,14 @@ class ESMCForkRequiredError(ModelSelectionError): """ESMC models require Biohub's transformers fork.""" +class METL3DNotSupportedError(ModelSelectionError): + """METL 3D models require PDB structures and are not supported by PEPE.""" + + +class METLPackageRequiredError(ModelSelectionError): + """METL models require the metl-pretrained package.""" + + def _import_hf_hub_errors(): """Lazy import of huggingface_hub exception types (version-tolerant).""" try: diff --git a/src/pepe/model_selecter.py b/src/pepe/model_selecter.py index f3d2741..98df270 100644 --- a/src/pepe/model_selecter.py +++ b/src/pepe/model_selecter.py @@ -1,7 +1,8 @@ from pepe.embedders.custom_embedder import CustomEmbedder -from pepe.model_errors import translate_hf_config_error +from pepe.model_errors import METL3DNotSupportedError, translate_hf_config_error import os +import re def _get_esm_embedder(): @@ -39,6 +40,29 @@ def _get_generic_hf_embedder(): return GenericHuggingFaceEmbedder +def _get_metl_embedder(): + """Lazy import of METL embedder (metl-pretrained backend).""" + from pepe.embedders.metl_embedder import METLEmbedder + + return METLEmbedder + + +def _is_metl_model(model_name): + if re.match(r"^metl[-_]", model_name, re.I): + return True + if model_name.lower() in ("gitter-lab/metl",): + return True + return False + + +def _validate_metl_model_name(model_name): + if re.search(r"[-_]3d(?:[-_]|$)", model_name, re.I): + raise METL3DNotSupportedError( + "METL 3D models (requiring PDB structures) are not supported by PEPE. " + "Use a 1D model identifier (e.g. metl-g-20m-1d)." + ) + + def select_model(model_name, trust_remote_code=False): # 1. Local checkpoints / directories take precedence over any name heuristic, # so a local file or folder is never mistaken for a HuggingFace repo id or a @@ -54,6 +78,10 @@ def select_model(model_name, trust_remote_code=False): ): return CustomEmbedder + if _is_metl_model(model_name): + _validate_metl_model_name(model_name) + return _get_metl_embedder() + # 2. Anything shaped like a HuggingFace repo id (username/model-name) is # dispatched by inspecting its config, not its name. This is the primary # signal: a fine-tune whose slug says "esm2" but is really a BERT no longer diff --git a/src/pepe/utils.py b/src/pepe/utils.py index ca8b15d..80dd2b3 100644 --- a/src/pepe/utils.py +++ b/src/pepe/utils.py @@ -342,6 +342,81 @@ def __getitem__(self, idx): else: return labels, seqs, toks, None, None + +class METLDataset(SequenceDictDataset): + def __init__( + self, + sequences, + substring_dict, + context, + data_encoder, + max_length, + ): + super().__init__(sequences, substring_dict, context) + self.data_encoder = data_encoder + self.encoded_data = self._encode_sequences(self.data, max_length) + self.pad_token_id = 0 + if self.substring_dict: + logger.info("Tokenizing substrings...") + self.encoded_substring_data = self._encode_sequences( + self.filtered_substring_data, + "max_length", + ) + logger.info("Matching substrings to full sequences...") + self.substring_masks = self._get_substring_masks() + + def _encode_sequences(self, data, max_length): + labels, strs = zip(*data) + encoded = [] + with alive_bar(len(strs), title="Tokenizing sequences...") as bar: + for s in strs: + seq_encoded = self.data_encoder.encode_sequences([s])[0] + encoded.append(seq_encoded) + bar() + + max_encoded_length = max(len(seq_encoded) for seq_encoded in encoded) + if max_length == "max_length": + max_length = max_encoded_length + else: + max_length = int(max_length) + if max_length < max_encoded_length: + logger.warning( + f"max_length {max_length} is less than the length of the longest sequence: {max_encoded_length}. Setting max_length to {max_encoded_length}." + ) + max_length = max_encoded_length + tokens = torch.empty((len(encoded), max_length), dtype=torch.int64) + tokens.fill_(0) + for i, seq_encoded in enumerate(encoded): + seq = torch.tensor(seq_encoded, dtype=torch.int64) + tokens[i, : len(seq_encoded)] = seq + return list(zip(labels, strs, list(tokens))) + + def get_max_encoded_length(self): + return max(len(toks) for _, _, toks in self.encoded_data) + + def safe_collate(self, batch): + if self.substring_dict: + labels, seqs, toks, _, substring_masks = zip(*batch) + return ( + list(labels), + list(seqs), + torch.stack(toks), + None, + torch.stack(substring_masks), + ) + else: + labels, seqs, toks, _, _ = zip(*batch) + return list(labels), list(seqs), torch.stack(toks), None, None + + def __getitem__(self, idx): + labels, seqs, toks = self.encoded_data[idx] + if self.substring_dict: + substring_masks = self.substring_masks[idx] + return labels, seqs, toks, None, substring_masks + else: + return labels, seqs, toks, None, None + + class CustomDataset(SequenceDictDataset): """ Dataset class for custom embedder. diff --git a/src/tests/test_metl_dispatch.py b/src/tests/test_metl_dispatch.py new file mode 100644 index 0000000..51a14a0 --- /dev/null +++ b/src/tests/test_metl_dispatch.py @@ -0,0 +1,78 @@ +import builtins +import os +import sys +import unittest +from unittest.mock import patch + +sys.path.insert(0, os.path.abspath("src")) + +from pepe.model_selecter import select_model +from pepe.model_errors import ( + METL3DNotSupportedError, + METLPackageRequiredError, + ModelSelectionError, +) + + +class TestMETLDispatch(unittest.TestCase): + """Mocked unit tests for METL model dispatch (no downloads).""" + + def test_select_metl_1d_returns_metl_embedder(self): + with patch( + "transformers.AutoConfig.from_pretrained", + side_effect=AssertionError("AutoConfig should not be called for METL idents"), + ): + embedder_cls = select_model("metl-g-20m-1d") + + from pepe.embedders.metl_embedder import METLEmbedder + + self.assertIs(embedder_cls, METLEmbedder) + + def test_select_metl_3d_raises_typed_error(self): + with self.assertRaises(METL3DNotSupportedError) as ctx: + select_model("metl-l-2m-3d-gb1") + + msg = str(ctx.exception).lower() + self.assertIn("3d", msg) + self.assertIn("metl-g-20m-1d", msg) + + def test_select_gitter_lab_metl_routes_to_metl_embedder(self): + with patch("transformers.AutoConfig.from_pretrained") as mock_config: + embedder_cls = select_model("gitter-lab/METL") + + from pepe.embedders.metl_embedder import METLEmbedder + + self.assertIs(embedder_cls, METLEmbedder) + mock_config.assert_not_called() + + def test_import_metl_missing_package_has_install_hint(self): + real_import = builtins.__import__ + + def mock_import(name, *args, **kwargs): + if name == "metl": + raise ImportError("No module named 'metl'") + return real_import(name, *args, **kwargs) + + with patch("builtins.__import__", side_effect=mock_import): + from pepe.embedders.metl_embedder import _import_metl + + with self.assertRaises(METLPackageRequiredError) as ctx: + _import_metl() + + msg = str(ctx.exception) + self.assertIn("metl-pretrained", msg) + self.assertIn("pip install", msg) + + def test_gitter_lab_metl_init_requires_specific_ident(self): + from pepe.embedders.metl_embedder import resolve_metl_ident + + with self.assertRaises(ModelSelectionError) as ctx: + resolve_metl_ident("gitter-lab/METL") + + msg = str(ctx.exception).lower() + self.assertIn("gitter-lab/metl", msg) + self.assertIn("metl-g-20m-1d", msg) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/tests/test_metl_integration.py b/src/tests/test_metl_integration.py new file mode 100644 index 0000000..6384f37 --- /dev/null +++ b/src/tests/test_metl_integration.py @@ -0,0 +1,41 @@ +import os +import sys +import unittest + +import numpy as np + +sys.path.insert(0, os.path.abspath("src")) + + +@unittest.skipUnless( + os.environ.get("METL_TEST") == "1", + "Set METL_TEST=1 to run METL integration test (requires metl-pretrained)", +) +class TestMETLIntegration(unittest.TestCase): + def test_metl_mean_pooled(self): + import pepe + + sequences = { + "seq1": "MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKTYFPHFDLSHGSAQVKGHGKKVADALTNAVAHVDDMPNALSALSDLHAHKLRVDPVNFKLLSHCLLVTLAAHLPAEFTPAVHASLDKFLASVSTVLTSKYR", + } + results = pepe.embed( + model_name="metl-g-20m-1d", + sequences=sequences, + extract_embeddings=["mean_pooled"], + device="cpu", + ) + + self.assertIn("mean_pooled", results) + layer_outputs = results["mean_pooled"] + layer_key = next(iter(layer_outputs)) + embedding = layer_outputs[layer_key][0] + self.assertEqual(embedding.shape, (512,)) + if hasattr(embedding, "detach"): + values = embedding.detach().cpu().numpy() + else: + values = np.asarray(embedding) + self.assertTrue(np.any(values != 0)) + + +if __name__ == "__main__": + unittest.main() From a096e4e8417809935c321b97c3a6dee4825ab992 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 21:01:08 +0200 Subject: [PATCH 08/17] docs: document METL 1D optional backend in README. Quick start install from GitHub [metl] extra, CLI example, supported models, and limitations aligned with CHANGELOG. Co-authored-by: Cursor --- README.md | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index f00744e..c8d94b4 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,29 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a pip install git+https://github.com/Biohub/transformers.git@main ``` -3. Extract embeddings:\ +3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the `metl-pretrained` backend via PEPE's `[metl]` extra. This extra is not yet on PyPI; install from the repository: + + ```sh + git clone https://github.com/csi-greifflab/pepe-cli + cd pepe-cli + pip install -e ".[metl]" + ``` + + Or install directly from GitHub: + + ```sh + pip install "pepe-cli[metl] @ git+https://github.com/csi-greifflab/pepe-cli.git" + ``` + + Example: + + ```sh + pepe --model_name metl-g-20m-1d --fasta_path --output_path --extract_embeddings mean_pooled + ``` + + **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). + +4. Extract embeddings:\ Extract mean pooled embeddings from protein amino acid sequences in FASTA file: ```sh pepe --experiment_name --fasta_path --output_path --model_name @@ -199,6 +221,8 @@ results = pepe.embed( - biohub/ESMC-300M - biohub/ESMC-600M - biohub/ESMC-6B + - METL 1D models (requires `[metl]` extra; see Quick start—install from GitHub until published on PyPI) + - `metl-g-20m-1d` and other `metl-*-1d` identifiers supported by [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) - Custom Hugging Face models - Any compatible model from Hugging Face Hub: `username/model-name` - Private models with authentication @@ -214,6 +238,7 @@ results = pepe.embed( - **`--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` - ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start) + - METL 1D models: `metl-g-20m-1d` (requires `[metl]` extra; see Quick start) - 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/` From b8c862050a319f2829316283f503593e0d991f52 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 21:02:13 +0200 Subject: [PATCH 09/17] docs: refine METL README with install and library examples. Document pepe-cli[metl] and direct metl-pretrained install, plus CLI/library usage for metl-g-20m-1d. Co-authored-by: Cursor --- README.md | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index c8d94b4..f4412f3 100644 --- a/README.md +++ b/README.md @@ -32,26 +32,38 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a pip install git+https://github.com/Biohub/transformers.git@main ``` -3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the `metl-pretrained` backend via PEPE's `[metl]` extra. This extra is not yet on PyPI; install from the repository: +3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the optional backend: ```sh - git clone https://github.com/csi-greifflab/pepe-cli - cd pepe-cli - pip install -e ".[metl]" + pip install pepe-cli[metl] ``` - Or install directly from GitHub: + The underlying `metl-pretrained` package is not on PyPI. PEPE's `[metl]` extra installs it from GitHub; you can also install it directly: ```sh - pip install "pepe-cli[metl] @ git+https://github.com/csi-greifflab/pepe-cli.git" + pip install git+https://github.com/gitter-lab/metl-pretrained.git ``` - Example: + CLI example: ```sh pepe --model_name metl-g-20m-1d --fasta_path --output_path --extract_embeddings mean_pooled ``` + Library example: + + ```python + import pepe + + results = pepe.embed( + model_name="metl-g-20m-1d", + sequences={"prot1": "MADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK"}, + output_path="my_embeddings", + extract_embeddings=["mean_pooled"], + device="cpu", + ) + ``` + **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). 4. Extract embeddings:\ @@ -221,7 +233,7 @@ results = pepe.embed( - biohub/ESMC-300M - biohub/ESMC-600M - biohub/ESMC-6B - - METL 1D models (requires `[metl]` extra; see Quick start—install from GitHub until published on PyPI) + - METL 1D models (requires `[metl]` extra; see Quick start) - `metl-g-20m-1d` and other `metl-*-1d` identifiers supported by [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) - Custom Hugging Face models - Any compatible model from Hugging Face Hub: `username/model-name` From f5649199549590f539a0129e52a076e04499f807 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 21:21:28 +0200 Subject: [PATCH 10/17] Fix METL layer-hook mapping and default layer selection Address review findings on the METL 1D embedder: - _register_repr_hooks: layer 0 now captures the encoder input embeddings (via a forward_pre_hook) instead of aliasing to layers[-1]; the final-layer norm hook is guarded so norm-less METL models fall back to the last block instead of crashing on None.register_forward_hook. Collapses the duplicated make_hook closure into one definition. - _load_layers: default (layers is None) now returns all transformer layers, matching the HuggingFace embedders, instead of only the final layer. The batch-dimension review item was verified against metl-pretrained source (batch_first=True throughout) and needs no change. Adds test_metl_hooks.py (deterministic, no metl download) covering the hook mapping and layer defaults, and extends the gated integration test with batch-consistency and multi-layer checks. --- .github/workflows/test.yml | 1 + src/pepe/embedders/metl_embedder.py | 54 ++++++----- src/tests/test_metl_hooks.py | 141 ++++++++++++++++++++++++++++ src/tests/test_metl_integration.py | 66 ++++++++++++- 4 files changed, 234 insertions(+), 28 deletions(-) create mode 100644 src/tests/test_metl_hooks.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 416e55c..6451066 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -68,6 +68,7 @@ jobs: src/tests/test_load_layers_default.py \ src/tests/test_model_selection.py \ src/tests/test_metl_dispatch.py \ + src/tests/test_metl_hooks.py \ src/tests/test_generic_safeguards.py \ src/tests/test_custom_embedder.py diff --git a/src/pepe/embedders/metl_embedder.py b/src/pepe/embedders/metl_embedder.py index adfe8d3..46e32e5 100644 --- a/src/pepe/embedders/metl_embedder.py +++ b/src/pepe/embedders/metl_embedder.py @@ -1,8 +1,10 @@ import logging + import torch + import pepe.utils from pepe.embedders.base_embedder import BaseEmbedder -from pepe.model_errors import ModelSelectionError, METLPackageRequiredError +from pepe.model_errors import METLPackageRequiredError, ModelSelectionError def _import_metl(): @@ -89,7 +91,9 @@ def _initialize_model(self, model_link): tr_encoder = model.model.tr_encoder num_tr_layers = len(tr_encoder.layers) num_layers = num_tr_layers - embedding_size = getattr(model, "embedding_len", None) or model.model.embedding_len + embedding_size = ( + getattr(model, "embedding_len", None) or model.model.embedding_len + ) num_heads = 1 if torch.cuda.is_available() and self.device.type == "cuda": @@ -103,25 +107,33 @@ def _initialize_model(self, model_link): def _register_repr_hooks(self): tr_encoder = self.model.model.tr_encoder num_tr_layers = len(tr_encoder.layers) + final_norm = getattr(tr_encoder, "norm", None) - for layer in self.layers: - if layer == num_tr_layers: - - def make_hook(captured_layer): - def hook(_module, _input, output): - self._repr_outputs[captured_layer] = output + def make_hook(captured_layer): + def hook(_module, _input, output): + self._repr_outputs[captured_layer] = output - return hook + return hook - handle = tr_encoder.norm.register_forward_hook(make_hook(layer)) - else: - - def make_hook(captured_layer): - def hook(_module, _input, output): - self._repr_outputs[captured_layer] = output + def make_pre_hook(captured_layer): + def pre_hook(_module, args): + self._repr_outputs[captured_layer] = args[0] - return hook + return pre_hook + # Layer index semantics match the HuggingFace path (hidden_states[k]): + # layer 0 = input embeddings, layer k (1..num_tr_layers) = k-th block output. + for layer in self.layers: + if layer == 0: + # Input embeddings (+ positional encoding): the tensor fed into the + # encoder, captured via a forward_pre_hook so it is agnostic to whether + # the model uses absolute or relative positional encoding. + handle = tr_encoder.register_forward_pre_hook(make_pre_hook(layer)) + elif layer == num_tr_layers and final_norm is not None: + # Post-norm final representation when the encoder has a final norm + # (norm_first architectures); otherwise fall through to the last block. + handle = final_norm.register_forward_hook(make_hook(layer)) + else: handle = tr_encoder.layers[layer - 1].register_forward_hook( make_hook(layer) ) @@ -129,13 +141,11 @@ def hook(_module, _input, output): def _load_layers(self, layers): if layers is None: - return [self.num_layers] + return list(range(1, self.num_layers + 1)) if not layers: layers = [-1] assert all(-(self.num_layers + 1) <= i <= self.num_layers for i in layers) - layers = [ - (i + self.num_layers + 1) % (self.num_layers + 1) for i in layers - ] + layers = [(i + self.num_layers + 1) % (self.num_layers + 1) for i in layers] return layers def _load_data(self, sequences, substring_dict=None): @@ -190,9 +200,7 @@ def _check_max_input_length(self): if max_allowed is None: return - sequences_too_long = any( - len(s) > max_allowed for s in self.sequences.values() - ) + sequences_too_long = any(len(s) > max_allowed for s in self.sequences.values()) needs_splitting = ( isinstance(self.max_input_length, int) and self.max_input_length > max_allowed diff --git a/src/tests/test_metl_hooks.py b/src/tests/test_metl_hooks.py new file mode 100644 index 0000000..c3c951c --- /dev/null +++ b/src/tests/test_metl_hooks.py @@ -0,0 +1,141 @@ +import os +import sys +import unittest + +sys.path.insert(0, os.path.abspath("src")) + +import torch +import torch.nn as nn + +from pepe.embedders.metl_embedder import METLEmbedder + + +class _MarkBlock(nn.Module): + """Transformer-block stand-in that adds a distinct scalar per block. + + Chaining the blocks yields a strictly increasing cumulative value, so the + tensor captured by a forward hook uniquely identifies which block produced + it (guards against the layer-0 -> layers[-1] aliasing bug). + """ + + def __init__(self, increment): + super().__init__() + self.increment = increment + + def forward(self, x): + return x + self.increment + + +class _MarkNorm(nn.Module): + """Final-norm stand-in that scales its input by 10 (a recognizable marker).""" + + def forward(self, x): + return x * 10.0 + + +class _FakeEncoder(nn.Module): + """Minimal batch_first TransformerEncoder analogue for METL.""" + + def __init__(self, num_layers, with_norm): + super().__init__() + self.layers = nn.ModuleList([_MarkBlock(i + 1) for i in range(num_layers)]) + self.norm = _MarkNorm() if with_norm else None + + def forward(self, src): + x = src + for layer in self.layers: + x = layer(x) + if self.norm is not None: + x = self.norm(x) + return x + + +class _FakeMETLInner(nn.Module): + def __init__(self, num_layers, with_norm): + super().__init__() + self.tr_encoder = _FakeEncoder(num_layers, with_norm) + + +class _FakeMETLModel(nn.Module): + def __init__(self, num_layers, with_norm): + super().__init__() + self.model = _FakeMETLInner(num_layers, with_norm) + + +def _make_embedder(num_layers, layers, with_norm=True): + """Build a METLEmbedder instance without running the real __init__.""" + emb = METLEmbedder.__new__(METLEmbedder) + emb.model = _FakeMETLModel(num_layers, with_norm) + emb.num_layers = num_layers + emb.layers = layers + emb._repr_outputs = {} + emb._hook_handles = [] + return emb + + +class TestMETLLoadLayers(unittest.TestCase): + """Layer-selection semantics must match the HuggingFace convention.""" + + def test_default_selects_all_transformer_layers(self): + emb = METLEmbedder.__new__(METLEmbedder) + emb.num_layers = 4 + # Regression for finding #3: default must be every layer, not just the last. + self.assertEqual(emb._load_layers(None), [1, 2, 3, 4]) + + def test_negative_and_zero_indexing(self): + emb = METLEmbedder.__new__(METLEmbedder) + emb.num_layers = 4 + self.assertEqual(emb._load_layers([-1]), [4]) + self.assertEqual(emb._load_layers([0]), [0]) + self.assertEqual(emb._load_layers([]), [4]) + + +class TestMETLRegisterReprHooks(unittest.TestCase): + """Each requested layer must capture the correct module's output.""" + + def _run(self, emb): + src = torch.zeros((2, 3, 5)) # (batch, seq, embed) + emb._register_repr_hooks() + emb.model.model.tr_encoder(src) + return src + + def test_layer_zero_captures_encoder_input(self): + # Regression for finding #2: layer 0 is the input embeddings, not layers[-1]. + emb = _make_embedder(num_layers=4, layers=[0]) + src = self._run(emb) + captured = emb._repr_outputs[0] + self.assertEqual(tuple(captured.shape), (2, 3, 5)) + self.assertTrue(torch.equal(captured, src)) # all zeros == the raw input + + def test_middle_layer_captures_its_own_block(self): + emb = _make_embedder(num_layers=4, layers=[2]) + self._run(emb) + # After blocks 0 and 1: 0 + 1 + 2 = 3 everywhere. Distinct from any other + # layer's value, so it cannot be silently aliased to the last block. + self.assertTrue(torch.all(emb._repr_outputs[2] == 3.0)) + + def test_final_layer_uses_norm_when_present(self): + emb = _make_embedder(num_layers=4, layers=[4], with_norm=True) + self._run(emb) + # Cumulative through all 4 blocks: 1+2+3+4 = 10, then norm scales by 10 -> 100. + self.assertTrue(torch.all(emb._repr_outputs[4] == 100.0)) + + def test_final_layer_falls_back_when_norm_is_none(self): + # Regression: norm-less METL models must not crash and must use the last block. + emb = _make_embedder(num_layers=4, layers=[4], with_norm=False) + self._run(emb) + self.assertTrue(torch.all(emb._repr_outputs[4] == 10.0)) # 1+2+3+4, no scaling + + def test_all_layers_captured_with_distinct_values(self): + emb = _make_embedder(num_layers=3, layers=[0, 1, 2, 3], with_norm=True) + self._run(emb) + self.assertEqual(set(emb._repr_outputs), {0, 1, 2, 3}) + # 0=input(0), 1=block0(1), 2=block1(1+2=3), 3=norm((1+2+3)*10=60) + self.assertTrue(torch.all(emb._repr_outputs[0] == 0.0)) + self.assertTrue(torch.all(emb._repr_outputs[1] == 1.0)) + self.assertTrue(torch.all(emb._repr_outputs[2] == 3.0)) + self.assertTrue(torch.all(emb._repr_outputs[3] == 60.0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/src/tests/test_metl_integration.py b/src/tests/test_metl_integration.py index 6384f37..6e80be7 100644 --- a/src/tests/test_metl_integration.py +++ b/src/tests/test_metl_integration.py @@ -7,6 +7,12 @@ sys.path.insert(0, os.path.abspath("src")) +def _as_numpy(embedding): + if hasattr(embedding, "detach"): + return embedding.detach().cpu().numpy() + return np.asarray(embedding) + + @unittest.skipUnless( os.environ.get("METL_TEST") == "1", "Set METL_TEST=1 to run METL integration test (requires metl-pretrained)", @@ -30,11 +36,61 @@ def test_metl_mean_pooled(self): layer_key = next(iter(layer_outputs)) embedding = layer_outputs[layer_key][0] self.assertEqual(embedding.shape, (512,)) - if hasattr(embedding, "detach"): - values = embedding.detach().cpu().numpy() - else: - values = np.asarray(embedding) - self.assertTrue(np.any(values != 0)) + self.assertTrue(np.any(_as_numpy(embedding) != 0)) + + def test_metl_batch_matches_single(self): + """Batching must not corrupt per-sequence embeddings (batch_first path). + + Two equal-length sequences are used so no padding is introduced, isolating + the batch dimension from any padding/attention effects: each sequence's + batched embedding must match its single-sequence embedding. + """ + import pepe + + seq_a = "MVLSPADKTNVKAAWGKVGA" # length 20 + seq_b = "ACDEFGHIKLMNPQRSTVWY" # length 20 + + def embed(seqs): + return pepe.embed( + model_name="metl-g-20m-1d", + sequences=seqs, + extract_embeddings=["mean_pooled"], + layers=[[-1]], + device="cpu", + ) + + single_a = embed({"a": seq_a})["mean_pooled"] + ref_a = _as_numpy(single_a[next(iter(single_a))][0]) + single_b = embed({"b": seq_b})["mean_pooled"] + ref_b = _as_numpy(single_b[next(iter(single_b))][0]) + + batched = embed({"a": seq_a, "b": seq_b})["mean_pooled"] + rows = [_as_numpy(x) for x in batched[next(iter(batched))]] + + self.assertEqual(len(rows), 2) + # Order-independent: each single-sequence reference matches some batched row. + self.assertTrue(any(np.allclose(ref_a, r, atol=1e-4) for r in rows)) + self.assertTrue(any(np.allclose(ref_b, r, atol=1e-4) for r in rows)) + # Distinct sequences must yield distinct embeddings (sanity). + self.assertFalse(np.allclose(ref_a, ref_b, atol=1e-4)) + + def test_metl_multiple_layers_distinct(self): + """Requesting multiple layers returns distinct per-layer representations.""" + import pepe + + results = pepe.embed( + model_name="metl-g-20m-1d", + sequences={"a": "ACDEFGHIKLMNPQRSTVWY"}, + extract_embeddings=["mean_pooled"], + layers=[[-1, -2]], + device="cpu", + )["mean_pooled"] + + self.assertEqual(len(results), 2) # two separate layers captured + keys = list(results) + v0 = _as_numpy(results[keys[0]][0]) + v1 = _as_numpy(results[keys[1]][0]) + self.assertFalse(np.allclose(v0, v1)) # not aliased to the same layer if __name__ == "__main__": From 0b7b073dc0312bd2a26c0f2712900ca8531d4be9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:26:24 +0000 Subject: [PATCH 11/17] fix pre-commit lint formatting failures --- setup.py | 4 +++- src/pepe/model_selecter.py | 2 +- src/tests/test_metl_dispatch.py | 6 ++++-- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/setup.py b/setup.py index 768b9b7..26a46f3 100644 --- a/setup.py +++ b/setup.py @@ -57,7 +57,9 @@ def read_readme(): "rjieba", ], extras_require={ - "metl": ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git"], + "metl": [ + "metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git" + ], "esm": ["fair-esm"], }, entry_points={ diff --git a/src/pepe/model_selecter.py b/src/pepe/model_selecter.py index bb017d4..a771d92 100644 --- a/src/pepe/model_selecter.py +++ b/src/pepe/model_selecter.py @@ -1,11 +1,11 @@ import os +import re from typing import Tuple, Type from pepe.embedders.base_embedder import BaseEmbedder from pepe.embedders.custom_embedder import CustomEmbedder from pepe.model_errors import METL3DNotSupportedError, translate_hf_config_error -import re def _get_esm_embedder() -> Type[BaseEmbedder]: """Lazy import of ESM embedder to avoid loading heavy dependencies.""" diff --git a/src/tests/test_metl_dispatch.py b/src/tests/test_metl_dispatch.py index 51a14a0..847f06d 100644 --- a/src/tests/test_metl_dispatch.py +++ b/src/tests/test_metl_dispatch.py @@ -6,12 +6,12 @@ sys.path.insert(0, os.path.abspath("src")) -from pepe.model_selecter import select_model from pepe.model_errors import ( METL3DNotSupportedError, METLPackageRequiredError, ModelSelectionError, ) +from pepe.model_selecter import select_model class TestMETLDispatch(unittest.TestCase): @@ -20,7 +20,9 @@ class TestMETLDispatch(unittest.TestCase): def test_select_metl_1d_returns_metl_embedder(self): with patch( "transformers.AutoConfig.from_pretrained", - side_effect=AssertionError("AutoConfig should not be called for METL idents"), + side_effect=AssertionError( + "AutoConfig should not be called for METL idents" + ), ): embedder_cls = select_model("metl-g-20m-1d") From 262fa3a1d25236fdc81e3c03012219050e20f19a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:39:00 +0000 Subject: [PATCH 12/17] fix: address PR review feedback for METL dispatch and CI gating --- .github/workflows/test.yml | 3 ++- pyproject.toml | 2 +- setup.py | 2 +- src/pepe/embedders/metl_embedder.py | 4 ++-- src/pepe/model_selecter.py | 11 ++++++++++- src/pepe/utils.py | 9 +++------ src/tests/test_metl_dispatch.py | 11 ++++++----- 7 files changed, 25 insertions(+), 17 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 6451066..1926c83 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -104,8 +104,9 @@ jobs: GENERIC_HF_TEST=1 python -m pytest -v \ src/tests/test_generic_hf_integration.py - name: Run METL integration test + if: ${{ vars.ENABLE_METL_TESTS == '1' }} run: | - pip install "git+https://github.com/gitter-lab/metl-pretrained.git" + pip install "git+https://github.com/gitter-lab/metl-pretrained.git@52358614c4b412e81e19e300485e8b85123bd903" METL_TEST=1 python -m pytest -v \ src/tests/test_metl_integration.py - name: Run T5 / AntiBERTa2 integration tests diff --git a/pyproject.toml b/pyproject.toml index 378be0c..ec7375e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ dependencies = [ ] [project.optional-dependencies] -metl = ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git"] +metl = ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git@52358614c4b412e81e19e300485e8b85123bd903"] esm = ["fair-esm"] # Developer toolchain: `pip install -e .[dev]` bootstraps everything needed to # run the hooks and the local test suite. diff --git a/setup.py b/setup.py index 26a46f3..e0aaafa 100644 --- a/setup.py +++ b/setup.py @@ -58,7 +58,7 @@ def read_readme(): ], extras_require={ "metl": [ - "metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git" + "metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git@52358614c4b412e81e19e300485e8b85123bd903" ], "esm": ["fair-esm"], }, diff --git a/src/pepe/embedders/metl_embedder.py b/src/pepe/embedders/metl_embedder.py index 46e32e5..d2b42d3 100644 --- a/src/pepe/embedders/metl_embedder.py +++ b/src/pepe/embedders/metl_embedder.py @@ -32,7 +32,7 @@ def resolve_metl_ident(model_name): return normalized.lower() -logger = logging.getLogger("src.embedders.metl_embedder") +logger = logging.getLogger("pepe.embedders.metl_embedder") class METLEmbedder(BaseEmbedder): @@ -100,7 +100,7 @@ def _initialize_model(self, model_link): model = model.cuda() logger.info("Transferred model to GPU") else: - logger.info("No GPU available, using CPU") + logger.info(f"Using device: {self.device.type}") return model, data_encoder, num_heads, num_layers, embedding_size diff --git a/src/pepe/model_selecter.py b/src/pepe/model_selecter.py index a771d92..282acdc 100644 --- a/src/pepe/model_selecter.py +++ b/src/pepe/model_selecter.py @@ -4,7 +4,11 @@ from pepe.embedders.base_embedder import BaseEmbedder from pepe.embedders.custom_embedder import CustomEmbedder -from pepe.model_errors import METL3DNotSupportedError, translate_hf_config_error +from pepe.model_errors import ( + METL3DNotSupportedError, + ModelSelectionError, + translate_hf_config_error, +) def _get_esm_embedder() -> Type[BaseEmbedder]: @@ -58,6 +62,11 @@ def _is_metl_model(model_name): def _validate_metl_model_name(model_name): + if model_name.lower() in ("gitter-lab/metl", "gitter-lab/metl-pretrained"): + raise ModelSelectionError( + "gitter-lab/METL is the HuggingFace wrapper and is not supported by PEPE. " + "Use a metl-pretrained identifier instead (e.g. metl-g-20m-1d)." + ) if re.search(r"[-_]3d(?:[-_]|$)", model_name, re.I): raise METL3DNotSupportedError( "METL 3D models (requiring PDB structures) are not supported by PEPE. " diff --git a/src/pepe/utils.py b/src/pepe/utils.py index 335e85e..98f7320 100644 --- a/src/pepe/utils.py +++ b/src/pepe/utils.py @@ -394,12 +394,9 @@ def __init__( def _encode_sequences(self, data, max_length): labels, strs = zip(*data) - encoded = [] - with alive_bar(len(strs), title="Tokenizing sequences...") as bar: - for s in strs: - seq_encoded = self.data_encoder.encode_sequences([s])[0] - encoded.append(seq_encoded) - bar() + with alive_bar(1, title="Tokenizing sequences...") as bar: + encoded = self.data_encoder.encode_sequences(list(strs)) + bar() max_encoded_length = max(len(seq_encoded) for seq_encoded in encoded) if max_length == "max_length": diff --git a/src/tests/test_metl_dispatch.py b/src/tests/test_metl_dispatch.py index 847f06d..f462f5a 100644 --- a/src/tests/test_metl_dispatch.py +++ b/src/tests/test_metl_dispatch.py @@ -38,13 +38,14 @@ def test_select_metl_3d_raises_typed_error(self): self.assertIn("3d", msg) self.assertIn("metl-g-20m-1d", msg) - def test_select_gitter_lab_metl_routes_to_metl_embedder(self): + def test_select_gitter_lab_metl_rejected_early(self): with patch("transformers.AutoConfig.from_pretrained") as mock_config: - embedder_cls = select_model("gitter-lab/METL") + with self.assertRaises(ModelSelectionError) as ctx: + select_model("gitter-lab/METL") - from pepe.embedders.metl_embedder import METLEmbedder - - self.assertIs(embedder_cls, METLEmbedder) + msg = str(ctx.exception).lower() + self.assertIn("gitter-lab/metl", msg) + self.assertIn("metl-g-20m-1d", msg) mock_config.assert_not_called() def test_import_metl_missing_package_has_install_hint(self): From 918e9bf4418ad6aaf68880ac5250e68a2f611473 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:47:46 +0000 Subject: [PATCH 13/17] Initial plan From 39c6efcf0db93141e923c0f56420c623c585afd6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:53:33 +0000 Subject: [PATCH 14/17] Fix TestPyPI publish failure: remove git URL from metl optional dependency PyPI/TestPyPI reject packages with direct URL dependencies (PEP 440). Remove `metl-pretrained @ git+https://...` from [project.optional-dependencies] in pyproject.toml and from extras_require in setup.py. Users should install metl-pretrained directly from GitHub. Update README, CHANGELOG, and error messages accordingly. --- CHANGELOG.md | 4 ++-- README.md | 12 +++--------- pyproject.toml | 1 - setup.py | 3 --- src/pepe/embedders/metl_embedder.py | 1 - 5 files changed, 5 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 60d47a0..70b953e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,10 +19,10 @@ truth that drives publishing). ### Added - METL 1D protein embeddings via optional `metl-pretrained` backend: install with - `pip install pepe-cli[metl]` and use model identifiers such as `metl-g-20m-1d` + `pip install git+https://github.com/gitter-lab/metl-pretrained.git` and use model identifiers such as `metl-g-20m-1d` (and other `metl-*-1d` names). Dispatch lives in `model_selecter.py`; embedding is handled by `METLEmbedder` with `METLDataset` tokenization. -- `[metl]` and `[esm]` optional dependency extras in `pyproject.toml` and +- `[esm]` optional dependency extra in `pyproject.toml` and `setup.py`. - Typed errors `METLPackageRequiredError` and `METL3DNotSupportedError` when METL is requested without the extra or when a 3D METL model id is used. diff --git a/README.md b/README.md index df630ac..97fd333 100644 --- a/README.md +++ b/README.md @@ -32,13 +32,7 @@ PEPE (Pipeline for Easy Protein Embedding) is a tool for extracting embeddings a pip install git+https://github.com/Biohub/transformers.git@main ``` -3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the optional backend: - - ```sh - pip install pepe-cli[metl] - ``` - - The underlying `metl-pretrained` package is not on PyPI. PEPE's `[metl]` extra installs it from GitHub; you can also install it directly: +3. *(Optional)* For METL 1D embedding models (e.g. `metl-g-20m-1d`), install the backend directly from GitHub (it is not on PyPI): ```sh pip install git+https://github.com/gitter-lab/metl-pretrained.git @@ -233,7 +227,7 @@ results = pepe.embed( - biohub/ESMC-300M - biohub/ESMC-600M - biohub/ESMC-6B - - METL 1D models (requires `[metl]` extra; see Quick start) + - METL 1D models (requires `metl-pretrained` from GitHub; see Quick start) - `metl-g-20m-1d` and other `metl-*-1d` identifiers supported by [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) - Custom Hugging Face models - Any compatible model from Hugging Face Hub: `username/model-name` @@ -250,7 +244,7 @@ results = pepe.embed( - **`--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` - ESMC models: `biohub/ESMC-300M` (requires Biohub transformers fork; see Quick start) - - METL 1D models: `metl-g-20m-1d` (requires `[metl]` extra; see Quick start) + - METL 1D models: `metl-g-20m-1d` (requires `metl-pretrained` from GitHub; see Quick start) - 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/` diff --git a/pyproject.toml b/pyproject.toml index ec7375e..19691de 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,6 @@ dependencies = [ ] [project.optional-dependencies] -metl = ["metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git@52358614c4b412e81e19e300485e8b85123bd903"] esm = ["fair-esm"] # Developer toolchain: `pip install -e .[dev]` bootstraps everything needed to # run the hooks and the local test suite. diff --git a/setup.py b/setup.py index e0aaafa..a2c601c 100644 --- a/setup.py +++ b/setup.py @@ -57,9 +57,6 @@ def read_readme(): "rjieba", ], extras_require={ - "metl": [ - "metl-pretrained @ git+https://github.com/gitter-lab/metl-pretrained.git@52358614c4b412e81e19e300485e8b85123bd903" - ], "esm": ["fair-esm"], }, entry_points={ diff --git a/src/pepe/embedders/metl_embedder.py b/src/pepe/embedders/metl_embedder.py index d2b42d3..f197add 100644 --- a/src/pepe/embedders/metl_embedder.py +++ b/src/pepe/embedders/metl_embedder.py @@ -16,7 +16,6 @@ def _import_metl(): except ImportError as e: raise METLPackageRequiredError( "METL models require metl-pretrained. Install with: " - "pip install 'pepe-cli[metl]' or " "pip install git+https://github.com/gitter-lab/metl-pretrained.git" ) from e From 0d4571bfd4d0bd9202177ce74dd8caaa2b8d54c8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 19:54:49 +0000 Subject: [PATCH 15/17] Add permissions: contents read to test.yml workflow (CodeQL fix) --- .github/workflows/test.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1926c83..ba37e2e 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -14,6 +14,9 @@ concurrency: group: tests-${{ github.ref }} cancel-in-progress: true +permissions: + contents: read + jobs: lint: name: Lint & format (pre-commit) From cf1e7faafee27689536981be65a75721962a2850 Mon Sep 17 00:00:00 2001 From: Jahn Zhong Date: Tue, 7 Jul 2026 22:21:22 +0200 Subject: [PATCH 16/17] Release 1.5.0: bump version and cut changelog section. Cut the Unreleased section into [1.5.0], covering METL 1D embeddings and GHCR Docker publishing. Restores the Docker/GHCR changelog entry dropped during an earlier merge into test. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 6 ++++++ src/pepe/__init__.py | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 70b953e..8295128 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,13 @@ truth that drives publishing). ## [Unreleased] +## [1.5.0] - 2026-07-07 + ### Added +- Docker images published to GHCR (`ghcr.io/csi-greifflab/pepe-cli`) on `test` and + `main` branch releases. Test tags include `:test` and a dev version tag; main + releases include `:latest` and the release version. GitHub release notes for + main now include `docker pull` / `docker run` examples. - METL 1D protein embeddings via optional `metl-pretrained` backend: install with `pip install git+https://github.com/gitter-lab/metl-pretrained.git` and use model identifiers such as `metl-g-20m-1d` (and other `metl-*-1d` names). Dispatch lives in `model_selecter.py`; embedding diff --git a/src/pepe/__init__.py b/src/pepe/__init__.py index fa7ef4f..3392a53 100644 --- a/src/pepe/__init__.py +++ b/src/pepe/__init__.py @@ -9,7 +9,7 @@ pass # Package metadata - single source of truth -__version__ = "1.4.0" +__version__ = "1.5.0" __package_name__ = "pepe-cli" __module_name__ = "pepe" __author__ = "Jahn Zhong" From f5cfe39ab3cd2a83e2e679ee4dc78fc7f5986f95 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 7 Jul 2026 20:54:23 +0000 Subject: [PATCH 17/17] Fix METL review findings for dispatch and reconstruction --- src/pepe/embedders/metl_embedder.py | 121 ++++++++++++++++++++++++---- src/pepe/model_selecter.py | 2 +- src/tests/test_metl_dispatch.py | 21 ++++- src/tests/test_metl_hooks.py | 37 +++++++++ 4 files changed, 163 insertions(+), 18 deletions(-) diff --git a/src/pepe/embedders/metl_embedder.py b/src/pepe/embedders/metl_embedder.py index f197add..d5ef9f9 100644 --- a/src/pepe/embedders/metl_embedder.py +++ b/src/pepe/embedders/metl_embedder.py @@ -34,24 +34,25 @@ def resolve_metl_ident(model_name): logger = logging.getLogger("pepe.embedders.metl_embedder") +def _validate_supported_output_types(output_types): + unsupported = [ + output_type + for output_type in output_types + if output_type + in ("logits", "attention_head", "attention_layer", "attention_model") + ] + if unsupported: + unsupported_sorted = ", ".join(sorted(set(unsupported))) + raise ModelSelectionError( + "METL models do not support logits or attention outputs. " + f"Remove unsupported extract_embeddings values: {unsupported_sorted}." + ) + + class METLEmbedder(BaseEmbedder): def __init__(self, args): super().__init__(args) - if self.return_logits: - logger.warning( - "Warning: Logits are not supported for METL models. Setting to False." - ) - self.return_logits = False - if "logits" in self.output_types: - self.output_types.remove("logits") - if self.return_contacts: - logger.warning( - "Warning: Attention matrices are not supported for METL models. Setting to False." - ) - self.return_contacts = False - for output_type in ("attention_head", "attention_layer", "attention_model"): - if output_type in self.output_types: - self.output_types.remove(output_type) + _validate_supported_output_types(self.output_types) self.sequences = pepe.utils.fasta_to_dict(args.fasta_path) self.num_sequences = len(self.sequences) @@ -70,7 +71,7 @@ def __init__(self, args): self.model_name, split_long_sequences=self.split_long_sequences, ) - self.special_tokens = torch.tensor([0], device=self.device, dtype=torch.int8) + self.special_tokens = torch.tensor([0], device=self.device, dtype=torch.int64) self.layers = self._load_layers(self.layers) self._repr_outputs = {} self._hook_handles = [] @@ -258,3 +259,91 @@ def _handle_sequence_splitting(self, max_allowed): if hasattr(self, "num_sequences"): self.num_sequences = len(self.sequences) self.max_input_length = chunk_size + + def _reconstruct_chunks(self): + """Reconstruct chunked METL outputs (no BOS/EOS token offsets).""" + if not self.chunks_mapping or self.streaming_output: + return + + assert self.layers is not None + logger.info("Reconstructing original sequences from chunks...") + + label_to_idx = {label: i for i, label in enumerate(self.sequence_labels)} + per_token_requested = "per_token" in self.output_types + rebuild_mean_pooled = "mean_pooled" in self.output_types + build_per_token = per_token_requested or rebuild_mean_pooled + + if rebuild_mean_pooled and not (per_token_requested or self._retain_per_token): + raise RuntimeError( + "Cannot reconstruct mean_pooled for split sequences without " + "per_token representations (internal per-token retention was " + "not enabled)." + ) + + output_type_map = [] + if build_per_token: + output_type_map.append(("per_token", self.per_token)) + if rebuild_mean_pooled: + output_type_map.append(("mean_pooled", self.mean_pooled)) + + new_sequence_labels = [] + labels_processed = set() + reconstructed_data = { + output_type: {layer: [] for layer in self.layers} + for output_type, _ in output_type_map + } + + for label in self.sequence_labels: + orig_label = label + is_chunk = False + for parent, chunks in self.chunks_mapping.items(): + if label in chunks: + orig_label = parent + is_chunk = True + break + + if orig_label in labels_processed: + continue + + labels_processed.add(orig_label) + new_sequence_labels.append(orig_label) + + if not is_chunk: + idx = label_to_idx[label] + for output_type, obj in output_type_map: + for layer in self.layers: + reconstructed_data[output_type][layer].append( + obj["output_data"][layer][idx] + ) + continue + + if build_per_token: + for layer in self.layers: + parts = [] + for i, chunk_label in enumerate(self.chunks_mapping[orig_label]): + idx = label_to_idx[chunk_label] + full_tensor = self.per_token["output_data"][layer][idx] + payload_len = self.chunk_payload_lengths[chunk_label] + start_idx = self.split_overlap if i > 0 else 0 + parts.append(full_tensor[start_idx:payload_len]) + reconstructed = torch.cat(parts, dim=0) + reconstructed_data["per_token"][layer].append(reconstructed) + + if rebuild_mean_pooled: + for layer in self.layers: + reconstructed_mean = reconstructed_data["per_token"][layer][ + -1 + ].mean(0) + reconstructed_data["mean_pooled"][layer].append(reconstructed_mean) + + self.sequence_labels = new_sequence_labels + self.num_sequences = len(self.sequence_labels) + for output_type, obj in output_type_map: + if output_type == "per_token" and not per_token_requested: + continue + for layer in self.layers: + obj["output_data"][layer] = reconstructed_data[output_type][layer] + + logger.info( + f"Reconstruction complete. Final sequence count: {self.num_sequences}" + ) diff --git a/src/pepe/model_selecter.py b/src/pepe/model_selecter.py index 282acdc..11e10f5 100644 --- a/src/pepe/model_selecter.py +++ b/src/pepe/model_selecter.py @@ -56,7 +56,7 @@ def _get_metl_embedder(): def _is_metl_model(model_name): if re.match(r"^metl[-_]", model_name, re.I): return True - if model_name.lower() in ("gitter-lab/metl",): + if model_name.lower() in ("gitter-lab/metl", "gitter-lab/metl-pretrained"): return True return False diff --git a/src/tests/test_metl_dispatch.py b/src/tests/test_metl_dispatch.py index f462f5a..b720a3f 100644 --- a/src/tests/test_metl_dispatch.py +++ b/src/tests/test_metl_dispatch.py @@ -48,6 +48,16 @@ def test_select_gitter_lab_metl_rejected_early(self): self.assertIn("metl-g-20m-1d", msg) mock_config.assert_not_called() + def test_select_gitter_lab_metl_pretrained_rejected_early(self): + with patch("transformers.AutoConfig.from_pretrained") as mock_config: + with self.assertRaises(ModelSelectionError) as ctx: + select_model("gitter-lab/metl-pretrained") + + msg = str(ctx.exception).lower() + self.assertIn("gitter-lab/metl", msg) + self.assertIn("metl-g-20m-1d", msg) + mock_config.assert_not_called() + def test_import_metl_missing_package_has_install_hint(self): real_import = builtins.__import__ @@ -67,7 +77,10 @@ def mock_import(name, *args, **kwargs): self.assertIn("pip install", msg) def test_gitter_lab_metl_init_requires_specific_ident(self): - from pepe.embedders.metl_embedder import resolve_metl_ident + from pepe.embedders.metl_embedder import ( + _validate_supported_output_types, + resolve_metl_ident, + ) with self.assertRaises(ModelSelectionError) as ctx: resolve_metl_ident("gitter-lab/METL") @@ -76,6 +89,12 @@ def test_gitter_lab_metl_init_requires_specific_ident(self): self.assertIn("gitter-lab/metl", msg) self.assertIn("metl-g-20m-1d", msg) + with self.assertRaises(ModelSelectionError): + _validate_supported_output_types(["logits"]) + + with self.assertRaises(ModelSelectionError): + _validate_supported_output_types(["per_token", "attention_layer"]) + if __name__ == "__main__": unittest.main() diff --git a/src/tests/test_metl_hooks.py b/src/tests/test_metl_hooks.py index c3c951c..af457b0 100644 --- a/src/tests/test_metl_hooks.py +++ b/src/tests/test_metl_hooks.py @@ -137,5 +137,42 @@ def test_all_layers_captured_with_distinct_values(self): self.assertTrue(torch.all(emb._repr_outputs[3] == 60.0)) +class TestMETLReconstruction(unittest.TestCase): + def test_reconstruct_chunks_without_special_token_offset(self): + emb = METLEmbedder.__new__(METLEmbedder) + emb.chunks_mapping = {"seq": ["seq_chunk_0", "seq_chunk_1"]} + emb.streaming_output = False + emb.layers = [1] + emb.sequence_labels = ["seq_chunk_0", "seq_chunk_1"] + emb.output_types = ["per_token", "mean_pooled"] + emb._retain_per_token = False + emb.split_overlap = 2 + emb.chunk_payload_lengths = {"seq_chunk_0": 5, "seq_chunk_1": 5} + emb.per_token = { + "output_data": { + 1: [ + torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0]]), + torch.tensor([[3.0], [4.0], [5.0], [6.0], [7.0]]), + ] + } + } + emb.mean_pooled = { + "output_data": {1: [torch.tensor([0.0]), torch.tensor([0.0])]} + } + + emb._reconstruct_chunks() + + self.assertEqual(emb.sequence_labels, ["seq"]) + self.assertTrue( + torch.equal( + emb.per_token["output_data"][1][0], + torch.tensor([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [6.0], [7.0]]), + ) + ) + self.assertTrue( + torch.equal(emb.mean_pooled["output_data"][1][0], torch.tensor([3.5])) + ) + + if __name__ == "__main__": unittest.main()