preprocess/extract_example_feat.py

"""Pre-extract ProtT5-XL features for the built-in example sequences.

Mirrors the offline pre-computation step in the training pipeline
(LLPSense/preprocess/extract_feat.py): run this once so the Gradio app never
has to hit the T5 model / GPU for a sequence it already knows about (see
`examples.find_example_by_seq` and `cb_extract` in app.py). This matters
most on the ZeroGPU-backed Space, where every GPU call consumes quota.

Usage:
    python preprocess/extract_example_feat.py
"""

import sys
from pathlib import Path

from huggingface_hub import snapshot_download
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from examples import EXAMPLES, feature_path, ASSETS_DIR  # noqa: E402
from t5_utils import T5_REPO_ID, extract_t5_feature, write_feature_h5  # noqa: E402


def main():
    snapshot_download(T5_REPO_ID)
    ASSETS_DIR.mkdir(parents=True, exist_ok=True)

    for example in tqdm(EXAMPLES, desc="Extracting example T5 features"):
        out_path = feature_path(example["id"])
        if out_path.exists():
            continue
        feat = extract_t5_feature(example["seq"])
        write_feature_h5(out_path, feat)


if __name__ == "__main__":
    main()