models/medgemma.py

# MedGemma 主模型类
# 继承 OneScience Module 基类，集成到 OneScience 框架

import logging
from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn

from onescience.modules.module import Module
from models.model_runner import VLLMModelRunner, TransformersModelRunner
from models.predictor_wrapper import MedGemmaPredictor

logger = logging.getLogger(__name__)


class MedGemma(Module):
    """
    MedGemma: 医学大语言模型

    支持:
    - 4B 多模态模型（文本 + 医学图像）
    - 27B 文本模型
    - DICOM/CT/CXR/WSI 图像输入
    - OpenAI Chat Completion API 格式

    继承 OneScience Module 基类，提供统一接口
    """

    def __init__(self, configs: Any) -> None:
        """
        初始化 MedGemma 模型

        Args:
            configs: 配置对象（ConfigDict）
        """
        super(MedGemma, self).__init__()
        self.configs = configs

        # 模型变体（4B 或 27B）
        self.model_variant = configs.model.variant
        self.is_multimodal = configs.model.is_multimodal

        logger.info(f"Initializing MedGemma {self.model_variant} model")
        logger.info(f"Multimodal: {self.is_multimodal}")

        # 初始化模型运行器
        self._init_model_runner()

        # 初始化推理包装器
        self.predictor = MedGemmaPredictor(
            model_runner=self.model_runner,
            configs=configs,
        )

        logger.info("MedGemma model initialized successfully")

    def _init_model_runner(self):
        """初始化模型运行器（vLLM 或 Transformers）"""
        if self.configs.inference.use_vllm:
            try:
                self.model_runner = VLLMModelRunner(
                    model_path=self.configs.model.model_path,
                    tokenizer_path=self.configs.model.tokenizer_path,
                    gpu_memory_utilization=self.configs.inference.gpu_memory_utilization,
                    max_model_len=self.configs.inference.max_model_len,
                    tensor_parallel_size=self.configs.inference.tensor_parallel_size,
                )
                logger.info("Using vLLM model runner")
            except Exception as e:
                logger.warning(f"vLLM initialization failed: {e}")
                logger.info("Falling back to Transformers model runner")
                self._init_transformers_runner()
        else:
            self._init_transformers_runner()

    def _init_transformers_runner(self):
        """初始化 Transformers 运行器"""
        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model_runner = TransformersModelRunner(
            model_path=self.configs.model.model_path,
            tokenizer_path=self.configs.model.tokenizer_path,
            device=device,
        )
        logger.info("Using Transformers model runner")

    def forward(
        self,
        messages: List[Dict[str, Any]],
        max_tokens: Optional[int] = None,
        temperature: float = 0.7,
        top_p: float = 0.9,
        n: int = 1,
    ) -> Dict[str, Any]:
        """
        前向传播（使用 OpenAI Chat Completion 格式）

        Args:
            messages: 消息列表，格式为 [{"role": "user", "content": "..."}]
            max_tokens: 最大生成 token 数
            temperature: 采样温度（0-2）
            top_p: Nucleus 采样参数（0-1）
            n: 生成数量

        Returns:
            OpenAI 格式的响应字典
        """
        if max_tokens is None:
            max_tokens = self.configs.inference.default_max_tokens

        return self.predictor.predict(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            n=n,
        )

    @torch.no_grad()
    def inference(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """
        推理方法（兼容 BiologyInferenceRunner）

        Args:
            data: 输入数据字典，可包含:
                - messages: 消息列表
                - instances: 实例列表（用于批处理）
                - parameters: 推理参数

        Returns:
            预测结果字典
        """
        # 转换数据格式
        if "instances" in data:
            messages = self._convert_instances_to_messages(data["instances"])
        elif "messages" in data:
            messages = data["messages"]
        else:
            raise ValueError("Input data must contain 'messages' or 'instances'")

        # 提取推理参数
        parameters = data.get("parameters", {})
        max_tokens = parameters.get("max_tokens", self.configs.inference.default_max_tokens)
        temperature = parameters.get("temperature", self.configs.inference.temperature)
        top_p = parameters.get("top_p", self.configs.inference.top_p)
        n = parameters.get("n", 1)

        # 运行推理
        return self.forward(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            n=n,
        )

    def _convert_instances_to_messages(
        self,
        instances: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        将实例列表转换为消息格式

        Args:
            instances: 实例列表

        Returns:
            消息列表
        """
        messages = []
        for instance in instances:
            if "role" in instance and "content" in instance:
                messages.append(instance)
            elif "text" in instance:
                messages.append({"role": "user", "content": instance["text"]})
            elif "question" in instance:
                messages.append({"role": "user", "content": instance["question"]})
            else:
                logger.warning(f"Unknown instance format: {instance}")

        return messages

    def predict_text(
        self,
        text: str,
        max_tokens: int = 500,
        temperature: float = 0.7,
    ) -> str:
        """
        简化的文本预测接口

        Args:
            text: 输入文本
            max_tokens: 最大生成 token 数
            temperature: 采样温度

        Returns:
            生成的文本
        """
        messages = [{"role": "user", "content": text}]
        response = self.forward(messages, max_tokens, temperature)

        if response["choices"]:
            return response["choices"][0]["message"]["content"]
        return ""

    def predict_multimodal(
        self,
        text: str,
        images: List[Any],
        max_tokens: int = 500,
        temperature: float = 0.7,
    ) -> str:
        """
        多模态预测接口（文本 + 图像）

        Args:
            text: 输入文本
            images: 图像列表
            max_tokens: 最大生成 token 数
            temperature: 采样温度

        Returns:
            生成的文本
        """
        if not self.is_multimodal:
            logger.warning("Model is not multimodal, ignoring images")
            return self.predict_text(text, max_tokens, temperature)

        # TODO: 实现多模态推理
        # 需要集成图像编码器和多模态提示格式
        messages = [{"role": "user", "content": text}]
        response = self.predictor.predict_with_images(
            messages=messages,
            images=images,
            max_tokens=max_tokens,
            temperature=temperature,
        )

        if response["choices"]:
            return response["choices"][0]["message"]["content"]
        return ""

    def __repr__(self) -> str:
        return f"MedGemma(variant={self.model_variant}, multimodal={self.is_multimodal})"