test_sft.py

# train_sft.py

"""
SFT (Supervised Fine-Tuning) 全量微调
加载预训练权重 → 在对话数据上训练 → 保存SFT权重
"""

import os
from dataclasses import dataclass
from typing import List, Dict, Optional

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.utils.data import Dataset, DataLoader

os.environ["TIKTOKEN_CACHE_DIR"] = "./tiktoken_cache"
import tiktoken

from model import GPT


# =============================================================================
# ========================= SFT 配置 ==========================================
# =============================================================================

@dataclass
class SFTConfig:
    # 路径
    PREtrained_path: str = "log/model_00999_CampGPT_X.pt"
    SFTtrained_path: str = "sft_output/sft_best.pt"
    sft_data_path: str = "sft_dataset.json"
    output_dir: str = "sft_output"
    
    # 训练超参
    num_epochs: int = 100
    batch_size: int = 8
    max_seq_len: int = 512          # SFT 通常不需要很长
    learning_rate: float = 2e-5     # 比预训练小 10-30 倍
    min_lr: float = 2e-6
    warmup_steps: int = 10
    weight_decay: float = 0.01
    grad_clip: float = 1.0
    grad_accum_steps: int = 1       # 有效batch = batch_size * grad_accum = 16
    
    # 对话格式
    system_prompt: str = "You are a helpful university assistant that answers questions about student policies and regulations."
    
    # 特殊 token
    # 用 GPT-2 tokenizer 中不常用的 token 作为分隔符
    user_token: str = "<|user|>"
    assistant_token: str = "<|assistant|>"
    end_token: str = "<|end|>"
    
    # 日志
    log_every: int = 10
    save_every_epoch: bool = False
    eval_ratio: float = 0.05        # 5% 数据做验证


# =============================================================================
# ========================= SFT 训练器 =========================================
# =============================================================================

class SFTTrainer:
    """SFT 全量微调训练器"""
    
    def __init__(self, sft_config: SFTConfig):
        self.sft_config = sft_config
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device_type = "cuda" if self.device.startswith("cuda") else "cpu"
        
        # 初始化
        # self._load_model_PreTrain()
        self._load_models_SFT()



    def _load_model_PreTrain(self):
        """加载预训练权重"""
        print(f"[SFT] Loading pretrained model from {self.sft_config.PREtrained_path}")
        
        checkpoint = torch.load(self.sft_config.PREtrained_path, map_location="cpu")
        model_config = checkpoint["config"]
        
        # 确保配置一致
        self.model_config = model_config
        self.model = GPT(model_config)
        
        # 加载权重
        state_dict = checkpoint["model"]
        # 处理可能的 key 不匹配（DDP/FSDP 保存的可能有 module. 前缀）
        cleaned = {}
        for k, v in state_dict.items():
            k = k.replace("module.", "").replace("_orig_mod.", "")
            cleaned[k] = v
        
        self.model.load_state_dict(cleaned, strict=True)
        self.model.to(self.device)
        
        # 关闭 gradient checkpointing（SFT 数据量小，不需要）
        self.model.set_gradient_checkpointing(False)
        
        pretrain_step = checkpoint.get("step", "unknown")
        pretrain_loss = checkpoint.get("val_loss", "unknown")
        print(f"[SFT] Loaded pretrained model (step={pretrain_step}, val_loss={pretrain_loss})")
        print(f"[SFT] Model params: {sum(p.numel() for p in self.model.parameters()):,}")
  
  


    def _load_models_SFT(self):
        print(f"[DPO] Loading SFT model from {self.sft_config.SFTtrained_path}")
        
        checkpoint = torch.load(self.sft_config.SFTtrained_path, map_location="cpu")
        self.model_config = checkpoint["config"]
        self.chat_template = checkpoint.get("chat_template", {})
        
        state_dict = checkpoint["model"]
        cleaned = {k.replace("module.", "").replace("_orig_mod.", ""): v 
                   for k, v in state_dict.items()}

        self.model = GPT(self.model_config)
        self.model.load_state_dict(cleaned, strict=True)
        self.model.to(self.device)
        self.model.eval()
        for param in self.model.parameters():
            param.requires_grad = False
        
        total_params = sum(p.numel() for p in self.model.parameters())
        print(f"[DPO] Policy model params: {total_params:,}")
        print(f"[DPO] Reference model: frozen copy")
        
        sft_loss = checkpoint.get("val_loss", "unknown")
        print(f"[DPO] SFT val_loss: {sft_loss}")
    

    
    def _test_generation(self):
        """训练后快速测试生成"""
        print(f"\n{'='*40} Generation Test {'='*40}")
        
        test_questions = [
            "Who is considered an Appellate Officer according to the student handbook?",
            "What does the term 'Associate Dean' refer to in the UHCL Student Handbook?",
            "How is 'Good Standing' defined for students at UHCL?",
        ]
        
        self.model.eval()
        enc = tiktoken.get_encoding("gpt2")
        
        for q in test_questions:
            prompt_text = ""
            if self.sft_config.system_prompt:
                prompt_text += f"### System:\n{self.sft_config.system_prompt}\n\n"
            prompt_text += f"### User:\n{q}\n\n### Assistant:\n"
            
            prompt_ids = enc.encode(prompt_text)
            prompt_t = torch.tensor([prompt_ids], dtype=torch.long, device=self.device)
            
            with torch.no_grad():
                with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
                    generated = self.model.generate(
                        prompt_t, max_new_tokens=64,
                        temperature=0.0, top_k=50, top_p=0.9
                    )
            
            output = enc.decode(generated[0].tolist())
            # 截取 assistant 回复部分
            if "### Assistant:" in output:
                answer = output.split("### Assistant:")[-1].strip()
                # 截断到第一个 ### 或换行结束
                if "###" in answer:
                    answer = answer[:answer.index("###")].strip()
            else:
                answer = output[len(prompt_text):]
            
            print(f"\nQ: {q}")
            print(f"A: {answer[:300]}")
        
        print(f"{'='*80}\n")


# =============================================================================
# ========================= 入口 ===============================================
# =============================================================================

if __name__ == "__main__":
    config = SFTConfig()
    
    # 可以命令行覆盖
    import sys
    for arg in sys.argv[1:]:
        if "=" in arg:
            key, val = arg.split("=", 1)
            key = key.lstrip("-")
            if hasattr(config, key):
                field_type = type(getattr(config, key))
                setattr(config, key, field_type(val))
                print(f"  Override: {key} = {val}")
    
    trainer = SFTTrainer(config)
    trainer._test_generation()