hf_loader.py

# -*- coding: utf-8 -*-
import os
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
from functools import lru_cache
from typing import Any, Dict

import torch
from huggingface_hub import hf_hub_download
from molscribe import MolScribe

MODEL_REPO = os.getenv("MODEL_REPO", "yujieq/MolScribe")
MODEL_FILE = os.getenv("MODEL_FILE", "swin_base_char_aux_1m.pth")
DEVICE_NAME = os.getenv("DEVICE") or ("cuda" if torch.cuda.is_available() else "cpu")
REQUEST_TIMEOUT_SECONDS = float(os.getenv("REQUEST_TIMEOUT_SECONDS", "180") or 0)

_PREDICT_EXECUTOR = ThreadPoolExecutor(
    max_workers=1,
    thread_name_prefix="molscribe_predict",
)


def model_descriptor() -> Dict[str, str]:
    return {
        "repo": MODEL_REPO,
        "file": MODEL_FILE,
        "device": DEVICE_NAME,
    }


@lru_cache(maxsize=1)
def _load_model() -> MolScribe:
    print(f"[hf_loader] Loading MolScribe checkpoint: {MODEL_REPO}/{MODEL_FILE} on {DEVICE_NAME}")
    ckpt_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    model = MolScribe(ckpt_path, device=torch.device(DEVICE_NAME))
    print("[hf_loader] MolScribe loaded.")
    return model


def warmup() -> Dict[str, str]:
    _load_model()
    return model_descriptor()


def _predict_once(
    image_path: str,
    return_atoms_bonds: bool,
    return_confidence: bool,
) -> Dict[str, Any]:
    model = _load_model()
    result = model.predict_image_file(
        image_path,
        return_atoms_bonds=return_atoms_bonds,
        return_confidence=return_confidence,
    )
    if not isinstance(result, dict):
        raise TypeError(f"MolScribe 返回了非字典结果: {type(result)!r}")
    return result


def _normalized_timeout_seconds(timeout_seconds: float | None) -> float | None:
    if timeout_seconds is None:
        timeout_seconds = REQUEST_TIMEOUT_SECONDS
    try:
        value = float(timeout_seconds)
    except Exception:
        value = 0.0
    if value <= 0:
        return None
    return value


def predict_image_file(
    image_path: str,
    return_atoms_bonds: bool = True,
    return_confidence: bool = True,
    timeout_seconds: float | None = None,
) -> Dict[str, Any]:
    timeout_value = _normalized_timeout_seconds(timeout_seconds)
    future = _PREDICT_EXECUTOR.submit(
        _predict_once,
        image_path,
        return_atoms_bonds,
        return_confidence,
    )
    try:
        result = future.result(timeout=timeout_value)
    except FutureTimeoutError as exc:
        future.cancel()
        if timeout_value is None:
            raise TimeoutError("MolScribe 推理超时") from exc
        raise TimeoutError(
            f"MolScribe 推理超过 {timeout_value:.0f} 秒，已触发超时保护。"
        ) from exc
    if not isinstance(result, dict):
        raise TypeError(f"MolScribe 返回了非字典结果: {type(result)!r}")
    return result