scripts/runner/inference_unified.py

import logging
import os
import sys
import traceback
from contextlib import nullcontext
from os.path import exists as opexists
from os.path import join as opjoin
from pathlib import Path
from typing import Any, Mapping

import torch
import torch.distributed as dist

PACKAGE_ROOT = Path(__file__).resolve().parents[2]
if str(PACKAGE_ROOT) not in sys.path:
    sys.path.insert(0, str(PACKAGE_ROOT))

from models.protenix.config import parse_configs, parse_sys_args
from models.protenix.protenix import Protenix
from onescience.utils.protenix.distributed import DIST_WRAPPER
from onescience.utils.protenix.seed import seed_everything
from onescience.utils.protenix.torch_utils import to_device
from scripts.runner.dumper import DataDumper

# Biology 数据加载器导入
from onescience.datapipes.biology.dataloader import get_protein_dataloader
from onescience.utils.YParams import YParams

logger = logging.getLogger(__name__)


class BiologyInferenceRunner:
    """使用 Biology Dataloader 的推理运行器"""

    def __init__(self, configs: Any) -> None:
        self.configs = configs
        self.init_env()
        self.init_basics()
        self.init_dataloader()
        self.init_model()
        self.load_checkpoint()
        self.init_dumper(
            need_atom_confidence=configs.need_atom_confidence,
            sorted_by_ranking_score=configs.sorted_by_ranking_score,
        )

    def init_env(self) -> None:
        """初始化环境"""
        self.print(
            f"Distributed: world_size={DIST_WRAPPER.world_size}, "
            f"rank={DIST_WRAPPER.rank}, local_rank={DIST_WRAPPER.local_rank}"
        )
        self.use_cuda = torch.cuda.device_count() > 0
        if self.use_cuda:
            self.device = torch.device(f"cuda:{DIST_WRAPPER.local_rank}")
            torch.cuda.set_device(self.device)
        else:
            self.device = torch.device("cpu")
        if DIST_WRAPPER.world_size > 1:
            dist.init_process_group(backend="nccl")
        logging.info("Finished init ENV.")

    def init_basics(self) -> None:
        """初始化基础设置"""
        self.dump_dir = self.configs.dump_dir
        self.error_dir = opjoin(self.dump_dir, "ERR")
        os.makedirs(self.dump_dir, exist_ok=True)
        os.makedirs(self.error_dir, exist_ok=True)

    def init_dataloader(self) -> None:
        """初始化 Biology Dataloader"""
        logger.info("Initializing Biology Dataloader")
        self.dataloader = get_protein_dataloader(configs=self.configs)
        logger.info(f"Dataloader initialized with {len(self.dataloader.dataset)} samples")

    def init_model(self) -> None:
        """初始化模型"""
        self.model = Protenix(self.configs).to(self.device)

    def load_checkpoint(self) -> None:
        """加载检查点"""
        checkpoint_path = self.configs.load_checkpoint_path
        if not os.path.exists(checkpoint_path):
            raise Exception(f"Checkpoint not found: {checkpoint_path}")
        self.print(f"Loading from {checkpoint_path}")
        checkpoint = torch.load(checkpoint_path, self.device)
        state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint
        state_dict = self._strip_module_prefix(state_dict)
        state_dict = self._select_checkpoint_key_format(state_dict)
        self.model.load_state_dict(state_dict, strict=self.configs.load_strict)
        self.model.eval()
        self.print("Checkpoint loaded")

    @staticmethod
    def _strip_module_prefix(state_dict: Mapping[str, Any]) -> dict:
        """Remove DistributedDataParallel's module. prefix when present."""
        if state_dict and all(k.startswith("module.") for k in state_dict.keys()):
            return {k[len("module.") :]: v for k, v in state_dict.items()}
        return dict(state_dict)

    def _select_checkpoint_key_format(self, state_dict: Mapping[str, Any]) -> dict:
        """Pick the checkpoint key format that best matches this model."""
        model_keys = set(self.model.state_dict().keys())
        candidates = {
            "original": dict(state_dict),
            "distogram_linear_wrapped": self._distogram_linear_wrapped_key_format(state_dict),
            "legacy_wrapped": self._legacy_wrapped_key_format(state_dict),
        }

        def score(candidate: Mapping[str, Any]) -> tuple[int, int]:
            candidate_keys = set(candidate.keys())
            matched = len(model_keys & candidate_keys)
            missing_or_unexpected = len(model_keys - candidate_keys) + len(candidate_keys - model_keys)
            return matched, -missing_or_unexpected

        best_name, best_state_dict = max(candidates.items(), key=lambda item: score(item[1]))
        best_keys = set(best_state_dict.keys())
        logger.info(
            "Selected checkpoint key format: %s (matched=%d, missing=%d, unexpected=%d)",
            best_name,
            len(model_keys & best_keys),
            len(model_keys - best_keys),
            len(best_keys - model_keys),
        )
        return best_state_dict

    @staticmethod
    def _distogram_linear_wrapped_key_format(state_dict: Mapping[str, Any]) -> dict:
        """Compatibility mapping for checkpoints with an unwrapped distogram linear."""
        remapped = {}
        for k, v in state_dict.items():
            new_key = k
            if new_key.startswith("distogram_head.linear.") and not new_key.startswith("distogram_head.linear.Linear."):
                new_key = new_key.replace("distogram_head.linear.", "distogram_head.linear.Linear.", 1)
            remapped[new_key] = v
        return remapped

    @staticmethod
    def _legacy_wrapped_key_format(state_dict: Mapping[str, Any]) -> dict:
        """Compatibility mapping for older wrappers used in some extracted packages."""
        remapped = {}
        for k, v in state_dict.items():
            new_key = k
            if new_key.startswith("input_embedder.") and not new_key.startswith("input_embedder.embedder."):
                new_key = new_key.replace("input_embedder.", "input_embedder.embedder.", 1)
            elif new_key.startswith("template_embedder.") and not new_key.startswith("template_embedder.embedder."):
                new_key = new_key.replace("template_embedder.", "template_embedder.embedder.", 1)
            elif new_key.startswith("relative_position_encoding.") and not new_key.startswith("relative_position_encoding.encoder."):
                new_key = new_key.replace("relative_position_encoding.", "relative_position_encoding.encoder.", 1)
            elif new_key.startswith("msa_module.") and not new_key.startswith("msa_module.msa."):
                new_key = new_key.replace("msa_module.", "msa_module.msa.", 1)
            elif new_key.startswith("pairformer_stack.") and not new_key.startswith("pairformer_stack.Pairformer."):
                new_key = new_key.replace("pairformer_stack.", "pairformer_stack.Pairformer.", 1)
            elif new_key.startswith("diffusion_module.") and not new_key.startswith("diffusion_module.Diffusion."):
                new_key = new_key.replace("diffusion_module.", "diffusion_module.Diffusion.", 1)
            elif new_key.startswith("distogram_head.linear.") and not new_key.startswith("distogram_head.linear.Linear."):
                new_key = new_key.replace("distogram_head.linear.", "distogram_head.linear.Linear.", 1)

            for prefix in (
                "linear_no_bias_sinit.",
                "linear_no_bias_zinit1.",
                "linear_no_bias_zinit2.",
                "linear_no_bias_token_bond.",
                "linear_no_bias_z_cycle.",
                "linear_no_bias_s.",
            ):
                if new_key.startswith(prefix) and not new_key.startswith(f"{prefix}Linear."):
                    new_key = new_key.replace(prefix, f"{prefix}Linear.", 1)
                    break

            if new_key.startswith("msa_module.msa.blocks.") and ".pair_stack." in new_key and ".pair_stack.Pairformer." not in new_key:
                new_key = new_key.replace(".pair_stack.", ".pair_stack.Pairformer.", 1)

            remapped[new_key] = v
        return remapped

    def init_dumper(self, need_atom_confidence=False, sorted_by_ranking_score=True):
        self.dumper = DataDumper(
            base_dir=self.dump_dir,
            need_atom_confidence=need_atom_confidence,
            sorted_by_ranking_score=sorted_by_ranking_score,
        )

    @torch.no_grad()
    def predict(self, data: Mapping[str, Any]) -> dict:
        eval_precision = {
            "fp32": torch.float32,
            "bf16": torch.bfloat16,
            "fp16": torch.float16,
        }[self.configs.dtype]

        enable_amp = (
            torch.autocast(device_type="cuda", dtype=eval_precision)
            if torch.cuda.is_available() else nullcontext()
        )

        data = to_device(data, self.device)
        with enable_amp:
            prediction, _, _ = self.model(
                input_feature_dict=data["input_feature_dict"],
                label_full_dict=None,
                label_dict=None,
                mode="inference",
            )
        return prediction

    def print(self, message: str) -> None:
        if DIST_WRAPPER.rank == 0:
            logger.info(message)

    def update_model_configs(self, new_configs: Any) -> None:
        self.model.configs = new_configs


def verify_required_local_files(configs: Any) -> None:
    for cache_name in ("ccd_components_file", "ccd_components_rdkit_mol_file"):
        cur_cache_fpath = configs["data"][cache_name]
        if not opexists(cur_cache_fpath):
            raise FileNotFoundError(
                f"Missing required local data cache: {cur_cache_fpath}. "
                "Set DATA_ROOT_DIR to a prepared Protenix dataset directory."
            )

    checkpoint_path = configs.load_checkpoint_path
    if not opexists(checkpoint_path):
        raise FileNotFoundError(
            f"Missing required local checkpoint: {checkpoint_path}. "
            "This standalone package expects weight/model_v0.5.0.pt."
        )


def update_inference_configs(configs: Any, N_token: int):
    if N_token > 3840:
        configs.skip_amp.confidence_head = False
        configs.skip_amp.sample_diffusion = False
    elif N_token > 2560:
        configs.skip_amp.confidence_head = False
        configs.skip_amp.sample_diffusion = True
    else:
        configs.skip_amp.confidence_head = True
        configs.skip_amp.sample_diffusion = True
    return configs


def run_inference(runner: BiologyInferenceRunner, configs: Any) -> None:
    """运行推理"""
    for seed in configs.seeds:
        seed_everything(seed=seed, deterministic=configs.deterministic)

        # 使用 dataloader 遍历数据
        for batch_idx, batch in enumerate(runner.dataloader):
            # batch 是列表，取第一个元素（因为 batch_size=1）
            data_item = batch[0] if isinstance(batch, list) else batch

            # 处理返回的元组格式 (result, atom_array, error_message)
            if isinstance(data_item, tuple):
                result_dict, atom_array, error_message = data_item
                sample_name = result_dict.get("sample_name", f"sample_{batch_idx}")
            else:
                # 兼容旧格式
                result_dict = data_item
                sample_name = data_item.get("sample_name", f"sample_{batch_idx}")
                atom_array = data_item.get("atom_array")
                error_message = ""

            try:
                if error_message:
                    raise RuntimeError(f"Dataset processing error: {error_message}")

                logger.info(f"[{batch_idx+1}/{len(runner.dataloader)}] Processing {sample_name}")

                # 从 dataloader 获取的数据中提取特征
                features_dict = result_dict["features"]
                if atom_array is None:
                    atom_array = result_dict.get("atom_array")

                # 构建输入数据
                N_token = features_dict["token_index"].shape[0]
                N_atom = features_dict["atom_to_token_idx"].shape[0]
                N_msa = features_dict["msa"].shape[0]
                N_asym = len(torch.unique(features_dict["asym_id"]))

                logger.info(
                    f"{sample_name}: N_asym={N_asym}, N_token={N_token}, "
                    f"N_atom={N_atom}, N_msa={N_msa}"
                )

                # 准备输入数据
                data = {
                    "input_feature_dict": features_dict,
                    "sample_name": sample_name,
                    "sample_index": batch_idx,
                }

                # 更新配置
                new_configs = update_inference_configs(configs, N_token)
                runner.update_model_configs(new_configs)

                # 预测
                prediction = runner.predict(data)

                # 保存结果
                runner.dumper.dump(
                    dataset_name="",
                    pdb_id=sample_name,
                    seed=seed,
                    pred_dict=prediction,
                    atom_array=atom_array,
                    entity_poly_type=features_dict.get("entity_poly_type", {}),
                )

                logger.info(f"{sample_name} succeeded. Results saved to {configs.dump_dir}")
                torch.cuda.empty_cache()

            except Exception as e:
                error_message = f"{sample_name}: {e}\n{traceback.format_exc()}"
                logger.error(error_message)
                with open(opjoin(runner.error_dir, f"{sample_name}.txt"), "w") as f:
                    f.write(error_message)
                if hasattr(torch.cuda, "empty_cache"):
                    torch.cuda.empty_cache()


def main():
    """主函数"""
    LOG_FORMAT = (
        "%(asctime)s,%(msecs)-3d %(levelname)-8s "
        "[%(filename)s:%(lineno)s %(funcName)s] %(message)s"
    )
    logging.basicConfig(
        format=LOG_FORMAT,
        level=logging.INFO,
        datefmt="%Y-%m-%d %H:%M:%S",
    )

    # 从YAML配置文件加载所有配置
    config_file_path = PACKAGE_ROOT / "configs" / "inference_config.yaml"
    cfg = YParams(config_file_path, "inference", print_params=True)

    # 解析命令行参数（命令行参数优先级最高）
    configs = parse_configs(
        configs=cfg.params,
        arg_str=parse_sys_args(),
        fill_required_with_null=True,
    )

    verify_required_local_files(configs)

    # 运行推理
    logger.info("=" * 60)
    logger.info("Starting Biology Dataloader Inference")
    logger.info("=" * 60)

    runner = BiologyInferenceRunner(configs)
    run_inference(runner, configs)

    logger.info("=" * 60)
    logger.info("Inference Completed")
    logger.info("=" * 60)


if __name__ == "__main__":
    main()