main.py

"""
main.py: Controller - 编排推理-训练解耦的数据流。

数据流：
1. main 读取图片路径
2. inference.extract_features.remote(paths) -> vis_ref (存入 Object Store)
3. inference.generate_rollouts.remote(prompts, paths) -> rollouts
4. training.train_step.remote(vis_ref, rollouts, rewards) -> 更新 LoRA

架构：
- GPU 1 (InferenceActor): vLLM 高速生成 + CLIP 特征提取
- GPU 0 (TrainingActor): 原生 PyTorch，手动拼接 Embedding 并计算 LoRA 梯度
"""

import os
import re
import json
import argparse
from typing import List, Tuple, Optional, Dict, Any

import ray
from ray.util.placement_group import placement_group, PlacementGroupSchedulingStrategy


# ---------------------------------------------------------------------------
# MultimodalMathVerifier: 数学题答案验证器
# ---------------------------------------------------------------------------

class MultimodalMathVerifier:
    """多模态数学题验证器：对比模型输出与 Ground Truth，返回 0/1 奖励。"""

    def __init__(self, normalize_number: bool = True):
        self.normalize_number = normalize_number

    @staticmethod
    def _extract_last_number(s: str) -> Optional[str]:
        """提取字符串中最后一个数字。"""
        numbers = re.findall(r"-?\d+\.?\d*", s)
        return numbers[-1] if numbers else None

    @staticmethod
    def _normalize(s: str) -> str:
        return s.strip().lower()

    def verify(self, model_answer: str, ground_truth: str) -> float:
        """单条验证，返回 0.0 或 1.0。"""
        gt_norm = self._normalize(ground_truth)
        ans_norm = self._normalize(model_answer)

        if self.normalize_number:
            gt_num = self._extract_last_number(gt_norm)
            ans_num = self._extract_last_number(ans_norm)
            if gt_num is not None and ans_num is not None:
                return 1.0 if gt_num == ans_num else 0.0

        return 1.0 if gt_norm in ans_norm or ans_norm in gt_norm else 0.0

    def verify_batch(
        self,
        responses_per_prompt: List[List[str]],
        ground_truths: List[str],
    ) -> List[List[float]]:
        """批量验证。"""
        return [
            [self.verify(resp, gt) for resp in resp_list]
            for resp_list, gt in zip(responses_per_prompt, ground_truths)
        ]


# ---------------------------------------------------------------------------
# 数据加载
# ---------------------------------------------------------------------------

def load_dataset(
    data_path: Optional[str] = None,
    max_samples: Optional[int] = None,
) -> List[Dict[str, Any]]:
    """
    加载数据集。
    格式：[{"image_path": "...", "question": "...", "answer": "..."}, ...]
    """
    if data_path and os.path.isfile(data_path):
        with open(data_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        samples = data if isinstance(data, list) else data.get("data", [data])
        if max_samples:
            samples = samples[:max_samples]
        return samples

    # 示例数据（联调用）
    return [
        {"image_path": "", "question": "What is 2 + 3?", "answer": "5"},
        {"image_path": "", "question": "Calculate 7 * 8.", "answer": "56"},
    ]


def build_prompts(
    samples: List[Dict[str, Any]],
    image_token: str = "<image>",
) -> Tuple[List[str], List[str], List[str]]:
    """构建 prompt、图片路径、ground truth 列表。"""
    prompts = []
    image_paths = []
    ground_truths = []

    for s in samples:
        question = s.get("question", "")
        # LLaVA 风格 prompt：<image> 在前
        prompt = f"{image_token}\nUSER: {question}\nASSISTANT:"
        prompts.append(prompt)
        image_paths.append(s.get("image_path", ""))
        ground_truths.append(s.get("answer", ""))

    return prompts, image_paths, ground_truths


# ---------------------------------------------------------------------------
# Placement Group 创建
# ---------------------------------------------------------------------------

def create_placement_group():
    """创建 2 GPU 的 placement group。"""
    bundles = [{"GPU": 1, "CPU": 2}, {"GPU": 1, "CPU": 2}]
    pg = placement_group(bundles, strategy="STRICT_PACK")
    return pg


# ---------------------------------------------------------------------------
# 主训练循环
# ---------------------------------------------------------------------------

def train(
    llm_model_name: str = "llava-hf/llava-1.5-7b-hf",
    clip_model_name: str = "openai/clip-vit-large-patch14-336",
    base_model_name: str = "meta-llama/Llama-2-7b-hf",
    data_path: Optional[str] = None,
    max_samples: int = 4,
    num_epochs: int = 2,
    grpo_n: int = 8,
    save_dir: str = "./checkpoints",
    projector_path: Optional[str] = None,
):
    """
    主训练函数。

    数据流：
    1. InferenceActor.extract_features(image_paths) -> visual_features (Object Store)
    2. InferenceActor.generate_rollouts(prompts, image_paths, n=8) -> responses
    3. Verifier 计算 rewards
    4. TrainingActor.train_step(visual_features, prompts, responses, rewards)
    """
    ray.init(ignore_reinit_error=True)

    # 创建 Placement Group
    pg = create_placement_group()
    ray.get(pg.ready())

    # 导入 Actor 类
    from inference_actor import InferenceActor
    from training_actor import TrainingActor

    # GPU 0 -> TrainingActor, GPU 1 -> InferenceActor
    training_actor = ray.remote(TrainingActor).options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_bundle_index=0,
        ),
    ).remote(
        base_model_name=base_model_name,
        projector_path=projector_path,
        clip_hidden_size=1024,
        llm_hidden_size=4096,
        freeze_projector=True,  # 阶段二：冻结 Projector
    )

    inference_actor = ray.remote(InferenceActor).options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_bundle_index=1,
        ),
    ).remote(
        llm_model_name=llm_model_name,
        clip_model_name=clip_model_name,
        gpu_memory_utilization=0.6,
    )

    # 等待 Actor 初始化
    ray.get(inference_actor.get_clip_hidden_size.remote())
    print("Actors initialized.")

    # 加载数据
    samples = load_dataset(data_path, max_samples)
    prompts, image_paths, ground_truths = build_prompts(samples)
    print(f"Loaded {len(samples)} samples.")

    # 验证器
    verifier = MultimodalMathVerifier(normalize_number=True)
    os.makedirs(save_dir, exist_ok=True)

    # 训练循环
    for epoch in range(num_epochs):
        print(f"\n{'='*50}")
        print(f"Epoch {epoch + 1}/{num_epochs}")
        print(f"{'='*50}")

        # Step 1: 提取 CLIP 特征 (GPU 1 -> Object Store)
        print("[1] Extracting visual features...")
        vis_ref = inference_actor.extract_features.remote(image_paths)

        # Step 2: vLLM 生成 rollouts (GPU 1)
        print(f"[2] Generating rollouts (n={grpo_n})...")
        responses_per_prompt = ray.get(
            inference_actor.generate_rollouts.remote(
                prompts=prompts,
                image_paths=image_paths,
                n=grpo_n,
            )
        )

        # Step 3: 计算 rewards
        print("[3] Computing rewards...")
        rewards = verifier.verify_batch(responses_per_prompt, ground_truths)
        avg_reward = sum(sum(r) for r in rewards) / sum(len(r) for r in rewards)
        print(f"    Average reward: {avg_reward:.4f}")

        # Step 4: 训练 (GPU 0)
        # 先 ray.get 获取 visual features
        visual_features = ray.get(vis_ref)
        print(f"[4] Training... (visual_features shape: {visual_features.shape})")

        loss_info = ray.get(
            training_actor.train_step.remote(
                visual_embeds_cpu=visual_features,
                prompt_texts=prompts,
                response_texts=responses_per_prompt,
                rewards=rewards,
                accumulation_steps=4,
            )
        )
        print(f"    Loss: {loss_info['loss']:.4f}, n_items: {loss_info['n_items']}")

        # Step 5: 保存 checkpoint
        state_dict = ray.get(training_actor.get_state_dict.remote())
        ckpt_path = os.path.join(save_dir, f"epoch_{epoch + 1}.pt")
        import torch
        torch.save(state_dict, ckpt_path)
        print(f"[5] Saved checkpoint: {ckpt_path}")

    ray.shutdown()
    print("\nTraining completed!")


# ---------------------------------------------------------------------------
# 推理测试
# ---------------------------------------------------------------------------

def inference(
    llm_model_name: str = "llava-hf/llava-1.5-7b-hf",
    clip_model_name: str = "openai/clip-vit-large-patch14-336",
    prompts: List[str] = None,
    image_paths: List[str] = None,
    n: int = 1,
):
    """
    推理测试：验证数据流。

    数据流与训练一致：
    1. InferenceActor.extract_features(image_paths) -> visual_features
    2. InferenceActor.generate_rollouts(prompts, image_paths, n) -> responses
    """
    ray.init(ignore_reinit_error=True)

    pg = create_placement_group()
    ray.get(pg.ready())

    from inference_actor import InferenceActor

    inference_actor = ray.remote(InferenceActor).options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_bundle_index=1,
        ),
    ).remote(
        llm_model_name=llm_model_name,
        clip_model_name=clip_model_name,
        gpu_memory_utilization=0.6,
    )

    prompts = prompts or ["<image>\nUSER: What is 2+3?\nASSISTANT:"]
    image_paths = image_paths or [""] * len(prompts)

    print("Extracting features...")
    vis_ref = inference_actor.extract_features.remote(image_paths)
    visual_features = ray.get(vis_ref)
    print(f"Visual features shape: {visual_features.shape}")

    print(f"Generating responses (n={n})...")
    responses = ray.get(
        inference_actor.generate_rollouts.remote(prompts, image_paths, n=n)
    )

    for i, (p, resps) in enumerate(zip(prompts, responses)):
        print(f"\nPrompt {i}: {p[:60]}...")
        for j, r in enumerate(resps):
            print(f"  [{j}] {r[:100]}...")

    ray.shutdown()
    return responses


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Tiny-R1-Vision: 推理-训练解耦的分布式多模态系统"
    )
    parser.add_argument(
        "--mode", type=str, default="train", choices=["train", "inference"],
        help="运行模式"
    )
    parser.add_argument("--data_path", type=str, default=None, help="数据集路径")
    parser.add_argument("--max_samples", type=int, default=4, help="最大样本数")
    parser.add_argument("--num_epochs", type=int, default=2, help="训练轮数")
    parser.add_argument("--grpo_n", type=int, default=8, help="GRPO 采样数")
    parser.add_argument("--save_dir", type=str, default="./checkpoints", help="保存目录")
    parser.add_argument("--projector_path", type=str, default=None, help="预训练 Projector 路径")

    # 推理参数
    parser.add_argument("--prompts", type=str, nargs="+", default=None)
    parser.add_argument("--image_paths", type=str, nargs="+", default=None)
    parser.add_argument("--n", type=int, default=1, help="推理采样数")

    args = parser.parse_args()

    if args.mode == "train":
        train(
            data_path=args.data_path,
            max_samples=args.max_samples,
            num_epochs=args.num_epochs,
            grpo_n=args.grpo_n,
            save_dir=args.save_dir,
            projector_path=args.projector_path,
        )
    else:
        inference(
            prompts=args.prompts,
            image_paths=args.image_paths,
            n=args.n,
        )