runner/medical_inference_runner.py

# MedGemma 医学推理运行器
# 类似 Protenix 的 BiologyInferenceRunner
import sys
from pathlib import Path

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

import logging
import os
import json
from typing import Any, Dict, List, Optional
import torch

from models.medgemma import MedGemma
from models.config import parse_configs, load_config

logger = logging.getLogger(__name__)

class AttrDict(dict):
    """同时支持 config.key 和 config['key']。"""

    def __getattr__(self, key):
        try:
            return self[key]
        except KeyError as exc:
            raise AttributeError(key) from exc

    def __setattr__(self, key, value):
        self[key] = value

def to_attr_dict(value):
    if isinstance(value, dict):
        return AttrDict({
            key: to_attr_dict(item)
            for key, item in value.items()
        })

    if isinstance(value, list):
        return [to_attr_dict(item) for item in value]

    if isinstance(value, tuple):
        return tuple(to_attr_dict(item) for item in value)

    return value


class MedicalInferenceRunner:
    """
    MedGemma 医学推理运行器
    提供统一的推理接口
    """

    def __init__(self, configs: Any) -> None:
        """
        初始化推理运行器

        Args:
            configs: 配置对象
        """
        #self.configs = configs
        self.configs = to_attr_dict(configs)
        self.init_env()
        self.init_basics()
        self.init_model()
        self.init_dumper()

        logger.info("MedicalInferenceRunner initialized")

    def init_env(self) -> None:
        """初始化环境"""
        self.use_cuda = torch.cuda.is_available()
        if self.use_cuda:
            self.device = torch.device("cuda:0")
            torch.cuda.set_device(self.device)
            logger.info(f"Using GPU: {torch.cuda.get_device_name(0)}")
        else:
            self.device = torch.device("cpu")
            logger.info("Using CPU")

    def init_basics(self) -> None:
        """初始化基础设置"""
        self.dump_dir = self.configs.output.dump_dir
        self.error_dir = os.path.join(self.dump_dir, "errors")
        os.makedirs(self.dump_dir, exist_ok=True)
        os.makedirs(self.error_dir, exist_ok=True)
        logger.info(f"Output directory: {self.dump_dir}")

    def init_model(self) -> None:
        """初始化 MedGemma 模型"""
        logger.info("Loading MedGemma model...")
        try:
            self.model = MedGemma(self.configs)
            logger.info("Model loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise

    def init_dumper(self) -> None:
        """初始化结果保存器"""
        self.output_format = self.configs.output.output_format
        self.save_predictions = self.configs.output.save_predictions
        logger.info(f"Output format: {self.output_format}")

    @torch.no_grad()
    def predict(
        self,
        messages: List[Dict[str, Any]],
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
    ) -> Dict[str, Any]:
        """
        运行推理

        Args:
            messages: 消息列表
            max_tokens: 最大生成 token 数
            temperature: 采样温度

        Returns:
            预测结果
        """
        if max_tokens is None:
            max_tokens = self.configs.inference.default_max_tokens
        if temperature is None:
            temperature = self.configs.inference.temperature

        try:
            result = self.model.forward(
                messages=messages,
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return result
        except Exception as e:
            logger.error(f"Prediction failed: {e}")
            return {"error": str(e)}

    def run_from_file(self, input_path: str) -> None:
        """
        从文件运行推理

        Args:
            input_path: 输入文件路径（JSON 或 JSONL）
        """
        logger.info(f"Loading input from: {input_path}")

        # 读取输入数据
        if input_path.endswith('.jsonl'):
            samples = self._load_jsonl(input_path)
        elif input_path.endswith('.json'):
            samples = self._load_json(input_path)
        else:
            raise ValueError(f"Unsupported file format: {input_path}")

        logger.info(f"Loaded {len(samples)} samples")

        # 处理每个样本
        results = []
        for idx, sample in enumerate(samples):
            logger.info(f"Processing sample {idx + 1}/{len(samples)}")

            try:
                # 提取消息
                if "messages" in sample:
                    messages = sample["messages"]
                elif "text" in sample:
                    messages = [{"role": "user", "content": sample["text"]}]
                elif "question" in sample:
                    messages = [{"role": "user", "content": sample["question"]}]
                else:
                    logger.warning(f"Sample {idx} has no valid input")
                    continue

                # 运行推理
                result = self.predict(messages)

                # 添加样本 ID
                result["sample_id"] = sample.get("id", idx)

                # 保存结果
                if self.save_predictions:
                    self._save_result(result, idx)

                results.append(result)

            except Exception as e:
                logger.error(f"Error processing sample {idx}: {e}")
                self._save_error(sample, idx, str(e))

        logger.info(f"Completed processing {len(results)} samples")

        # 保存汇总结果
        self._save_summary(results)

    def _load_json(self, filepath: str) -> List[Dict[str, Any]]:
        """加载 JSON 文件"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)

        if isinstance(data, list):
            return data
        else:
            return [data]

    def _load_jsonl(self, filepath: str) -> List[Dict[str, Any]]:
        """加载 JSONL 文件"""
        samples = []
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    samples.append(json.loads(line))
        return samples

    def _save_result(self, result: Dict[str, Any], idx: int) -> None:
        """保存单个结果"""
        output_path = os.path.join(
            self.dump_dir,
            f"prediction_{idx}.{self.output_format}"
        )

        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(result, f, indent=2, ensure_ascii=False)

    def _save_error(self, sample: Dict[str, Any], idx: int, error: str) -> None:
        """保存错误信息"""
        error_path = os.path.join(
            self.error_dir,
            f"error_{idx}.json"
        )

        error_data = {
            "sample": sample,
            "error": error,
        }

        with open(error_path, 'w', encoding='utf-8') as f:
            json.dump(error_data, f, indent=2, ensure_ascii=False)

    def _save_summary(self, results: List[Dict[str, Any]]) -> None:
        """保存汇总结果"""
        summary_path = os.path.join(self.dump_dir, "summary.json")

        summary = {
            "total_samples": len(results),
            "successful": sum(1 for r in results if "error" not in r),
            "failed": sum(1 for r in results if "error" in r),
            "results": results,
        }

        with open(summary_path, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)

        logger.info(f"Summary saved to: {summary_path}")

    def run_interactive(self) -> None:
        """交互式推理"""
        logger.info("Starting interactive mode. Type 'quit' to exit.")

        while True:
            try:
                user_input = input("\nUser: ")
                if user_input.lower() in ['quit', 'exit', 'q']:
                    break

                messages = [{"role": "user", "content": user_input}]
                result = self.predict(messages)

                # 提取响应
                if "choices" in result and result["choices"]:
                    response = result["choices"][0]["message"]["content"]
                    print(f"\nAssistant: {response}")
                else:
                    print(f"\nError: {result}")

            except KeyboardInterrupt:
                break
            except Exception as e:
                logger.error(f"Error: {e}")

        logger.info("Exiting interactive mode")


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description="MedGemma Medical Inference Runner")
    parser.add_argument("--config", type=str, required=True, help="Config file path")
    parser.add_argument("--input", type=str, help="Input file path (JSON/JSONL)")
    parser.add_argument("--interactive", action="store_true", help="Interactive mode")
    parser.add_argument("--model_path", type=str, help="Override model path")
    parser.add_argument("--dump_dir", type=str, help="Override output directory")

    args = parser.parse_args()

    # 加载配置
    configs = load_config(args.config)

    # 覆盖配置
    if args.model_path:
        configs.model.model_path = args.model_path
    if args.dump_dir:
        configs.output.dump_dir = args.dump_dir

    # 创建运行器
    runner = MedicalInferenceRunner(configs)

    # 运行推理
    if args.interactive:
        runner.run_interactive()
    elif args.input:
        runner.run_from_file(args.input)
    else:
        logger.error("Either --input or --interactive must be specified")


if __name__ == "__main__":
    main()