main_controller.py

"""
main_controller.py: Ray 调度中心，练推一体。
数据流始终为：GPU 1 提特征 → 存入 Object Store → GPU 0 拉取并算梯度/算生成 Token。
第一步：训练投影层（GPU 1 辅助，GPU 0 只训 Projector）。
第二步：指令微调与 GRPO（GPU 0 多开 LoRA，同数据流下生成 8 组 + train_step）。
"""

import os
import re
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, fuzzy_match: bool = True):
        self.normalize_number = normalize_number
        self.fuzzy_match = fuzzy_match

    @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:
        s = s.strip().lower()
        s = re.sub(r"\s+", " ", s)
        return s

    def verify(self, model_answers: List[str], ground_truths: List[str]) -> List[float]:
        gt = ground_truths[0] if isinstance(ground_truths[0], str) else str(ground_truths[0])
        gt_norm = self._normalize(gt)
        gt_num = self._extract_last_number(gt_norm) if self.normalize_number else None
        rewards = []
        for ans in model_answers:
            ans_norm = self._normalize(ans)
            ans_num = self._extract_last_number(ans_norm) if self.normalize_number else None
            if self.normalize_number and gt_num is not None and ans_num is not None:
                match = gt_num == ans_num
            else:
                match = gt_norm in ans_norm or (self.fuzzy_match and ans_norm in gt_norm)
            rewards.append(1.0 if match else 0.0)
        return rewards

    def verify_batch(
        self,
        responses_per_prompt: List[List[str]],
        ground_truth_list: List[str],
    ) -> List[List[float]]:
        return [
            self.verify(resp_list, [gt])
            for resp_list, gt in zip(responses_per_prompt, ground_truth_list)
        ]


# ---------------------------------------------------------------------------
# 数据集与 Prompt 构建
# ---------------------------------------------------------------------------

def load_multimodal_math_dataset(
    data_path: Optional[str] = None,
    image_key: str = "image_path",
    question_key: str = "question",
    answer_key: str = "answer",
    max_samples: Optional[int] = None,
) -> List[Dict[str, Any]]:
    if data_path and os.path.isfile(data_path):
        import json
        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.get("samples", [data]))
        if max_samples:
            samples = samples[:max_samples]
        return samples
    return [
        {"image_path": "", "question": "Solve: 2 + 3 = ?", "answer": "5"},
        {"image_path": "", "question": "What is 7 * 8?", "answer": "56"},
    ]


def build_prompts(
    samples: List[Dict[str, Any]],
    question_key: str = "question",
    img_placeholder: str = "[IMG]",
) -> Tuple[List[str], List[str], List[str]]:
    prompts, image_paths, ground_truths = [], [], []
    for s in samples:
        q = s.get(question_key, "")
        prompts.append(f"{img_placeholder}\n{q}\nAnswer:")
        image_paths.append(s.get("image_path", ""))
        ground_truths.append(s.get("answer", ""))
    return prompts, image_paths, ground_truths


def create_placement_group_for_two_gpus():
    bundles = [{"GPU": 1, "CPU": 2}, {"GPU": 1, "CPU": 2}]
    return placement_group(bundles, strategy="STRICT_PACK")


# ---------------------------------------------------------------------------
# 第一步：训练投影层（GPU 1 提特征 → Object Store → GPU 0 拉取并只训 Projector）
# ---------------------------------------------------------------------------

def run_step1_train_projection(
    inference_actor,
    training_actor,
    samples: List[Dict[str, Any]],
    base_model_name: str,
    img_placeholder: str,
    num_epochs: int = 2,
):
    """第一步：GPU 1 只跑 CLIP，产出特征存 Object Store；GPU 0 冻结 LLM，只训 Projector。"""
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
    tokenizer.add_special_tokens({"additional_special_tokens": [img_placeholder]})

    prompts, image_paths, ground_truths = build_prompts(samples)
    valid = [i for i, p in enumerate(image_paths) if p and os.path.isfile(p)]
    if not valid:
        valid = list(range(len(prompts)))
    prompts = [prompts[i] for i in valid]
    image_paths = [image_paths[i] for i in valid]
    ground_truths = [ground_truths[i] for i in valid]

    prompt_token_ids = [tokenizer.encode(p, add_special_tokens=False) for p in prompts]
    response_token_ids = [tokenizer.encode(gt, add_special_tokens=False) for gt in ground_truths]

    for epoch in range(num_epochs):
        # 数据流：GPU 1 提特征 → Object Store
        visual_emb = ray.get(inference_actor.extract_visual_embeddings.remote(image_paths))
        # GPU 0 从 Object Store 拉取，只训 Projector
        loss_info = ray.get(
            training_actor.train_projection_step.remote(
                visual_embeddings=visual_emb,
                prompt_token_ids=prompt_token_ids,
                response_token_ids=response_token_ids,
            )
        )
        print(f"[Step1 投影层] Epoch {epoch + 1}/{num_epochs} loss: {loss_info['loss']:.4f}")
    print("第一步完成：投影层已训练，可进入第二步 GRPO。")


# ---------------------------------------------------------------------------
# 第二步：GRPO（同数据流，GPU 0 生成 8 组 + 算奖励 + train_step）
# ---------------------------------------------------------------------------

def run_step2_grpo(
    inference_actor,
    training_actor,
    samples: List[Dict[str, Any]],
    base_model_name: str,
    img_placeholder: str,
    verifier: MultimodalMathVerifier,
    save_dir: str,
    num_epochs: int = 2,
):
    """第二步：GPU 1 仍只提特征；GPU 0 拉取特征后 generate_rollouts(8 组) → 算奖励 → train_step。"""
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
    tokenizer.add_special_tokens({"additional_special_tokens": [img_placeholder]})

    prompts, image_paths, ground_truths = build_prompts(samples)
    valid = [i for i, p in enumerate(image_paths) if p and os.path.isfile(p)]
    if not valid:
        valid = list(range(len(prompts)))
    prompts = [prompts[i] for i in valid]
    image_paths = [image_paths[i] for i in valid]
    ground_truths = [ground_truths[i] for i in valid]

    for epoch in range(num_epochs):
        # 数据流：GPU 1 提特征 → Object Store
        visual_emb = ray.get(inference_actor.extract_visual_embeddings.remote(image_paths))
        # GPU 0 拉取特征，算生成 Token（与推理同数据流）
        responses_per_prompt = ray.get(
            training_actor.generate_rollouts.remote(
                visual_embeddings=visual_emb,
                prompts=prompts,
                sampling_n=8,
            )
        )
        rewards_batch = verifier.verify_batch(responses_per_prompt, ground_truths)

        prompt_token_ids = [tokenizer.encode(p, add_special_tokens=False) for p in prompts]
        response_token_ids = [
            [tokenizer.encode(r, add_special_tokens=False) for r in resp_list]
            for resp_list in responses_per_prompt
        ]

        loss_info = ray.get(
            training_actor.train_step.remote(
                prompt_token_ids=prompt_token_ids,
                response_token_ids=response_token_ids,
                rewards=rewards_batch,
                visual_embeddings=visual_emb,
                accumulation_steps=4,
            )
        )
        print(f"[Step2 GRPO] Epoch {epoch + 1}/{num_epochs} loss: {loss_info['loss']:.4f}")

        lora_sd = ray.get(training_actor.get_lora_state_dict.remote())
        import torch
        path = os.path.join(save_dir, f"lora_epoch_{epoch + 1}.pt")
        torch.save(lora_sd, path)
        print(f"Saved {path}")


# ---------------------------------------------------------------------------
# 推理测试（练推一体：同一数据流，无需改底层代码）
# ---------------------------------------------------------------------------

def run_inference_test(
    inference_actor,
    training_actor,
    prompts: List[str],
    image_paths: List[str],
    sampling_n: int = 1,
) -> List[List[str]]:
    """
    推理时：GPU 1 提特征 → Object Store → GPU 0 拉取并算生成的 Token。
    与训练时数据流一致，无需改动底层数据流代码。
    """
    visual_emb = ray.get(inference_actor.extract_visual_embeddings.remote(image_paths))
    responses_per_prompt = ray.get(
        training_actor.generate_rollouts.remote(
            visual_embeddings=visual_emb,
            prompts=prompts,
            sampling_n=sampling_n,
        )
    )
    return responses_per_prompt


def main(
    clip_model_name: str = "openai/clip-vit-large-patch14",
    base_model_name: str = "meta-llama/Llama-2-7b-hf",
    data_path: Optional[str] = None,
    max_samples: Optional[int] = 4,
    save_dir: str = "./grpo_checkpoints",
    step1_epochs: int = 2,
    step2_epochs: int = 2,
    run_step1: bool = True,
    run_step2: bool = True,
):
    ray.init(ignore_reinit_error=True)
    pg = create_placement_group_for_two_gpus()
    ray.get(pg.ready())

    from inference_actor import InferenceActor
    from training_actor import TrainingActor

    training_actor = ray.remote(TrainingActor).options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=pg,
            placement_group_bundle_index=0,
        ),
    ).remote(
        base_model_name=base_model_name,
        clip_hidden_size=1024,
        use_lora=True,
    )

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

    ray.get(inference_actor.get_clip_hidden_size.remote())
    os.makedirs(save_dir, exist_ok=True)
    samples = load_multimodal_math_dataset(data_path=data_path, max_samples=max_samples)
    verifier = MultimodalMathVerifier(normalize_number=True, fuzzy_match=True)

    if run_step1:
        run_step1_train_projection(
            inference_actor,
            training_actor,
            samples,
            base_model_name=base_model_name,
            img_placeholder="[IMG]",
            num_epochs=step1_epochs,
        )
    if run_step2:
        run_step2_grpo(
            inference_actor,
            training_actor,
            samples,
            base_model_name=base_model_name,
            img_placeholder="[IMG]",
            verifier=verifier,
            save_dir=save_dir,
            num_epochs=step2_epochs,
        )

    ray.shutdown()
    print("主流程结束。")


if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--data_path", type=str, default=None)
    p.add_argument("--max_samples", type=int, default=4)
    p.add_argument("--save_dir", type=str, default="./grpo_checkpoints")
    p.add_argument("--step1_epochs", type=int, default=2)
    p.add_argument("--step2_epochs", type=int, default=2)
    p.add_argument("--no_step1", action="store_true", help="跳过第一步投影层训练")
    p.add_argument("--no_step2", action="store_true", help="跳过第二步 GRPO")
    args = p.parse_args()
    main(
        data_path=args.data_path,
        max_samples=args.max_samples,
        save_dir=args.save_dir,
        step1_epochs=args.step1_epochs,
        step2_epochs=args.step2_epochs,
        run_step1=not args.no_step1,
        run_step2=not args.no_step2,
    )