main_bert_fine_tune.py

#!/usr/bin/env python3
# main.py - 支持原有流程 + Bert HF fine-tune 模式

import os
import re
import shutil
import csv
import sys
import json
from pathlib import Path

import torch
import torch.distributed as dist
import math
from dotenv import load_dotenv

import hydra
from omegaconf import OmegaConf

# project imports (adjust paths if needed)
from lmr.config import initialize_config
from lmr.tokenizer import Tokenizer
from lmr.models import get_model 
from lmr.checkpointing import Checkpointing
from lmr.utils.seed import set_seed
from lmr.training import Bert_Trainer, Trainer
from lmr.generation import Generator
from lmr.benchmark import Benchmark
from lmr.ddp import unwrap_model

# import the BertFineTuneTrainer (ensure this module exists at this path)
# 如果你把 BertFineTuneTrainer 放在 lmr/training/bert_finetune_trainer.py，则如下导入：
try:
    from lmr.training.bert_finetune_trainer import BertFineTuneTrainer
except Exception as e:
    # 如果没有该文件，提醒并继续（后续会报错）
    BertFineTuneTrainer = None
    print("⚠️ Warning: Could not import BertFineTuneTrainer: ", e)

DATASET_DIR = Path("datasets")
CHECKPOINT_DIR = Path("/work/jf381/checkpoints")
BENCHMARK_DIR = Path("output")

# -------------------------
# Helper functions (unchanged)
# -------------------------
def load_weight_data(path_obj, device="cpu"):
    from safetensors.torch import load_file
    model_state = {}
    
    path_obj = Path(path_obj)
    if path_obj.is_dir():
        index_file = path_obj / "model.safetensors.index.json"
        if index_file.exists():
            print(f"🔹 Detected sharded safetensors folder: {path_obj.name}")
            with open(index_file, 'r') as f:
                index_data = json.load(f)
            weight_map = index_data.get("weight_map", {})
            shards = set(weight_map.values())
            for shard_name in shards:
                shard_path = path_obj / shard_name
                model_state.update(load_file(str(shard_path), device=str(device)))
            return model_state
        else:
            possible = list(path_obj.glob("*.safetensors")) + list(path_obj.glob("*.pt"))
            if not possible:
                return None
            path_obj = possible[0]

    if path_obj.suffix == ".safetensors":
        print(f"🔹 Loading single safetensors: {path_obj.name}")
        return load_file(str(path_obj), device=str(device))
    else:
        print(f"🔹 Loading pickle (.pt): {path_obj.name}")
        ckpt = torch.load(path_obj, map_location=device)
        return ckpt.get("model", ckpt.get("state_dict", ckpt))

def load_state_dict_robust(model, checkpoint_path, strict=False):
    path_obj = Path(checkpoint_path)
    if not path_obj.exists():
        print(f"❌ Path not found: {checkpoint_path}")
        return False

    try:
        model_state = load_weight_data(path_obj)
        if model_state is None: return False

        ckpt_keys_map = {}
        for k in model_state.keys():
            clean_k = k.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            ckpt_keys_map[clean_k] = k

        target_model = unwrap_model(model)
        target_state = target_model.state_dict()
        
        filtered_state = {}
        matched_count = 0

        for k_target, v_target in target_state.items():
            k_target_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            if k_target_clean in ckpt_keys_map:
                real_ckpt_key = ckpt_keys_map[k_target_clean]
                v_ckpt = model_state[real_ckpt_key]
                if v_ckpt.shape == v_target.shape:
                    filtered_state[k_target] = v_ckpt
                    matched_count += 1

        msg = target_model.load_state_dict(filtered_state, strict=strict)
        print(f"✅ Success! Loaded {matched_count} parameters. Status: {msg}")
        return True
    except Exception as e:
        print(f"❌ Failed to load checkpoint: {e}")
        import traceback
        traceback.print_exc()
        return False

def setup_model_and_tokenizer(config):
    print(f"🔧 Initializing Tokenizer: {config.tokenizer_base}")
    tokenizer = Tokenizer(config.tokenizer_base)
    print(f"🔧 Initializing Model: {config.model}")
    model = get_model(config.model, tokenizer.vocab_size, tokenizer=tokenizer)
    return tokenizer, model

def run_generation_task(config, model, tokenizer, output_dir, ckpt_name_tag=""):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    
    generator = Generator(config, model, tokenizer, device=device, output_dir=output_dir)
    results = generator.generate()
    
    metrics = {}
    if results and "verification" in results:
        metrics = {
            'acc': results['verification']['accuracy'],
            'correct': results['verification']['correct'],
            'total': results['verification']['total']
        }
        if results.get('token_accuracies'):
            import numpy as np
            metrics['token_acc'] = float(np.mean(results['token_accuracies']))

        if ckpt_name_tag:
            for fname in ["gsm8k_metrics.txt", "gsm8k_generations.txt"]:
                src = output_dir / fname
                if src.exists():
                    dst = output_dir / f"{src.stem}_{ckpt_name_tag}{src.suffix}"
                    shutil.move(src, dst)
    return metrics

# -------------------------
# Train / Generate orchestration
# -------------------------

def generate(config):
    tokenizer, model = setup_model_and_tokenizer(config)
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    load_mode = getattr(config.benchmark, "checkpoint_mode", "recent")
    print(f"🔍 Looking for [{load_mode}] weights in {ckpt_dir}...")

    all_items = list(ckpt_dir.iterdir()) if ckpt_dir.exists() else []
    checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]

    def sort_key(f):
        nums = re.findall(r'\d+', f.name)
        return int(nums[-1]) if nums else 0
    checkpoints.sort(key=sort_key)

    target_ckpt = None
    if load_mode == "best":
        best_candidates = [f for f in checkpoints if "best" in f.name.lower()]
        target_ckpt = best_candidates[0] if best_candidates else (checkpoints[-1] if checkpoints else None)
    else:
        target_ckpt = checkpoints[-1] if checkpoints else None

    if target_ckpt:
        print(f"🚀 Found target: {target_ckpt.name}")
        success = load_state_dict_robust(model, target_ckpt)
        if not success:
            print("⚠️ Load failed, check file integrity.")
    else:
        print(f"❌ No checkpoints found in {ckpt_dir}")

    run_generation_task(config, model, tokenizer, output_dir=ckpt_dir)

def generate_all(config):
    tokenizer, model = setup_model_and_tokenizer(config)
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    all_items = list(ckpt_dir.iterdir()) if ckpt_dir.exists() else []
    checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]
    
    def sort_key(f):
        nums = re.findall(r'\d+', f.name)
        return int(nums[-1]) if nums else 0
    checkpoints.sort(key=sort_key)
    
    print(f"\n🔎 Found {len(checkpoints)} checkpoints.")
    summary_path = ckpt_dir / "metrics_summary.csv"
    
    with open(summary_path, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["checkpoint", "accuracy", "token_accuracy", "correct", "total"])

    for ckpt_path in checkpoints:
        print(f"\n{'-'*40}\nProcessing: {ckpt_path.name}\n{'-'*40}")
        if load_state_dict_robust(model, ckpt_path):
            metrics = run_generation_task(config, model, tokenizer, ckpt_dir, ckpt_name_tag=ckpt_path.name)
            if metrics:
                with open(summary_path, mode='a', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow([
                        ckpt_path.name,
                        f"{metrics.get('acc', 0):.4f}",
                        f"{metrics.get('token_acc', 0):.4f}",
                        metrics.get('correct', 0),
                        metrics.get('total', 0)
                    ])

def train_model(config):
    # if user wants to run HF Bert fine-tune mode, handle specially
    # 配置约定： config.mode == "finetune_bert" 或 config.finetune.bert.enabled == true
    use_bert_finetune_mode = False
    if getattr(config, "mode", "") == "finetune_bert":
        use_bert_finetune_mode = True
    elif hasattr(config, "finetune") and getattr(config.finetune, "bert", None) and getattr(config.finetune.bert, "enabled", False):
        use_bert_finetune_mode = True

    if use_bert_finetune_mode:
        if BertFineTuneTrainer is None:
            raise RuntimeError("BertFineTuneTrainer not available. 请确认 lmr.training.bert_finetune_trainer.py 存在并且可导入。")

        # 构造 BertFineTuneTrainer 需要的 cfg 字典
        # 优先使用 config.finetune.bert 下的参数，其次从 config.training/全局取默认
        ft_cfg = {}
        # required fields sample:
        # model_name_or_path, dataset, dataset_config_name (optional), task, num_labels, output_dir, batch_size, eval_batch_size, num_epochs, lr, weight_decay, gradient_accumulation_steps, max_length, max_train_samples, max_eval_samples, fp16, use_ddp
        fin = getattr(config, "finetune", None)
        if fin and getattr(fin, "bert", None):
            fin_bert = fin.bert
        else:
            # backward compat: maybe config.training contains some entries
            fin_bert = getattr(config, "training", {})

        # map possible fields — 这里尽量宽容
        ft_cfg["model_name_or_path"] = getattr(fin_bert, "model_name_or_path", None) or getattr(config, "model_name", None) or getattr(config, "model", None)
        ft_cfg["dataset"] = getattr(fin_bert, "dataset", None) or getattr(config, "dataset", None)
        ft_cfg["dataset_config_name"] = getattr(fin_bert, "dataset_config_name", None)
        ft_cfg["task"] = getattr(fin_bert, "task", "sentence_pair")
        ft_cfg["num_labels"] = getattr(fin_bert, "num_labels", None)
        ft_cfg["output_dir"] = str(Path(getattr(fin_bert, "output_dir", config.get("output_dir", "./outputs/finetune")))) if isinstance(config, dict) else str(Path(getattr(fin_bert, "output_dir", "./outputs/finetune")))
        ft_cfg["batch_size"] = getattr(fin_bert, "batch_size", getattr(config, "batch_size", 16))
        ft_cfg["eval_batch_size"] = getattr(fin_bert, "eval_batch_size", max(32, int(ft_cfg["batch_size"])))
        ft_cfg["num_epochs"] = getattr(fin_bert, "num_epochs", getattr(config, "num_epochs", 3))
        ft_cfg["lr"] = getattr(fin_bert, "lr", getattr(config, "learning_rate", 2e-5))
        ft_cfg["weight_decay"] = getattr(fin_bert, "weight_decay", 0.01)
        ft_cfg["gradient_accumulation_steps"] = getattr(fin_bert, "gradient_accumulation_steps", 1)
        ft_cfg["max_length"] = getattr(fin_bert, "max_length", 128)
        ft_cfg["fp16"] = getattr(fin_bert, "fp16", True)
        ft_cfg["use_ddp"] = bool(getattr(config, "distributed", False) or getattr(config, "use_ddp", False))
        ft_cfg["seed"] = getattr(config, "seed", 42)
        ft_cfg["logging_steps"] = getattr(fin_bert, "logging_steps", 100)
        ft_cfg["eval_steps"] = getattr(fin_bert, "eval_steps", 500)
        ft_cfg["save_steps"] = getattr(fin_bert, "save_steps", 1000)
        ft_cfg["warmup_steps"] = getattr(fin_bert, "warmup_steps", 0)
        ft_cfg["max_train_samples"] = getattr(fin_bert, "max_train_samples", None)
        ft_cfg["max_eval_samples"] = getattr(fin_bert, "max_eval_samples", None)
        ft_cfg["nsp_negatives_ratio"] = getattr(fin_bert, "nsp_negatives_ratio", 1)

        # Validation of required args
        if not ft_cfg["model_name_or_path"] or not ft_cfg["dataset"]:
            raise ValueError("finetune_bert 模式需要在配置中指定 model_name_or_path 和 dataset（Hugging Face dataset id）. 示例: finetune.bert.model_name_or_path='bert-base-uncased', finetune.bert.dataset='glue/mrpc'")

        # Convert OmegaConf objects to plain types if necessary
        if not isinstance(ft_cfg["dataset"], str) and hasattr(ft_cfg["dataset"], "__str__"):
            ft_cfg["dataset"] = str(ft_cfg["dataset"])

        # Create and run the BertFineTuneTrainer
        print(f"🔧 Starting HuggingFace BERT finetune with cfg: model={ft_cfg['model_name_or_path']}, dataset={ft_cfg['dataset']}, task={ft_cfg.get('task')}")
        trainer = BertFineTuneTrainer(ft_cfg)
        trainer.train()
        return

    # -------------------------
    # 原有的训练流程（不做改动）：
    # -------------------------
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        if not dist.is_initialized():
            dist.init_process_group(backend="nccl")
            torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())

    tokenizer, model = setup_model_and_tokenizer(config)
    tokenized_dataset_dir = DATASET_DIR / config.tokenizer_base
    splits = None
    try:
        # 尝试调用老的 get_dataset_splits（如果存在）
        from lmr.data import get_dataset_splits
        splits = get_dataset_splits(config.dataset, 1024, tokenized_dataset_dir)
    except Exception:
        print("⚠️ get_dataset_splits not available or failed — continuing without it.")

    checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
    
    if "bert" in str(config.model).lower():
        from lmr.training import Bert_Trainer
        trainer = Bert_Trainer(config.training, model, tokenizer, splits, checkpointing, None)
    else:
        trainer = Trainer(config.training, model, tokenizer, splits, checkpointing)
    trainer.train()

# -------------------------
# Entrypoint
# -------------------------
@hydra.main(config_path="config", config_name="config", version_base="1.3")
def main(config):
    # config is an OmegaConf object
    load_dotenv()
    set_seed(config)
    initialize_config(config)

    mode = getattr(config, "mode", "train")
    if mode == "train":
        train_model(config)
    elif mode == "generate":
        generate(config)
    elif mode == "generate_all":
        generate_all(config)
    elif mode == "benchmark":
        tokenizer, model = setup_model_and_tokenizer(config)
        checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
        benchmarking = Benchmark(config.benchmark, model, tokenizer, checkpointing, BENCHMARK_DIR / config.checkpoint_name)
        benchmarking.run_benchmarks()
    elif mode == "finetune_bert":
        # 进入我们上面实现的 fine-tune 分支
        train_model(config)
    else:
        print(f"❌ Unknown mode: {mode}")

if __name__ == "__main__":
    main()