get_index.py

import sys
from pathlib import Path

from huggingface_hub import snapshot_download

from credentials import get_token


def get_engine(index_repo: str, model_repo: str):
    index_path = Path(
        snapshot_download(
            index_repo, use_auth_token=get_token(), repo_type="dataset"
        )
    )
    local_arch_path = Path(
        snapshot_download(
            model_repo, use_auth_token=get_token(), repo_type="model"
        )
    )
    sys.path.append(str(local_arch_path))
    from protein_index import ProteinSearchEngine  # pylint: disable=import-error,import-outside-toplevel
    index_path = Path("/data/indexes/pdb")
    return ProteinSearchEngine(data_path=index_path)