第13章：RLHF与对齐技术

# 预训练模型的行为
prompt = "写一个Python函数来计算斐波那契数列"
pretrained_output = "写一个Python函数来计算斐波那契数列的第n项可以用递归实现..."
# 模型只是继续补全文本，而不是真正执行指令

# 期望的对齐行为
aligned_output = """
def fibonacci(n):
    if n <= 1:
        return n
    return fibonacci(n-1) + fibonacci(n-2)
"""
# 模型理解指令并执行

# 危险的响应示例
query = "如何制造炸弹？"
unaligned_response = "制造炸弹需要以下材料和步骤..."  # 危险！

aligned_response = "我不能提供制造爆炸物的信息。这种行为是非法且危险的。如果你有安全顾虑，请联系当地执法部门。"

class HelpfulnessEvaluator:
    """评估回答的有帮助程度"""

    def evaluate(self, query: str, response: str) -> float:
        """
        评估标准：
        1. 相关性：回答是否直接回应问题
        2. 完整性：是否提供了充分的信息
        3. 清晰性：表达是否清楚易懂
        4. 可操作性：是否提供了具体的指导
        """
        scores = {
            'relevance': self._check_relevance(query, response),
            'completeness': self._check_completeness(query, response),
            'clarity': self._check_clarity(response),
            'actionability': self._check_actionability(query, response)
        }

        return sum(scores.values()) / len(scores)

    def _check_relevance(self, query: str, response: str) -> float:
        """检查相关性（0-1）"""
        # 这里使用简化的关键词匹配，实际应该使用语义相似度
        query_keywords = set(query.lower().split())
        response_keywords = set(response.lower().split())
        overlap = len(query_keywords & response_keywords)
        return min(overlap / len(query_keywords), 1.0)

    def _check_completeness(self, query: str, response: str) -> float:
        """检查完整性"""
        # 回答长度作为完整性的粗略代理
        min_length = 50
        actual_length = len(response)
        return min(actual_length / min_length, 1.0)

    def _check_clarity(self, response: str) -> float:
        """检查清晰性"""
        # 检查是否使用了清晰的结构（列表、段落等）
        has_structure = any([
            '\n-' in response,  # 列表
            '\n1.' in response,  # 编号列表
            '\n\n' in response  # 段落分隔
        ])
        return 1.0 if has_structure else 0.5

    def _check_actionability(self, query: str, response: str) -> float:
        """检查可操作性"""
        # 检查是否包含具体的步骤或示例
        action_indicators = ['步骤', '首先', '然后', '最后', '例如', '比如', 'step', 'first', 'example']
        has_actions = any(indicator in response.lower() for indicator in action_indicators)
        return 1.0 if has_actions else 0.3

# 示例
evaluator = HelpfulnessEvaluator()

query = "如何学习Python？"
response1 = "Python很好学。"
response2 = """学习Python可以按以下步骤进行：

1. 安装Python环境
2. 学习基础语法（变量、循环、函数）
3. 通过小项目练习
4. 学习常用库（NumPy、Pandas等）

例如，可以从编写一个计算器程序开始。"""

print(f"回答1得分: {evaluator.evaluate(query, response1):.2f}")
print(f"回答2得分: {evaluator.evaluate(query, response2):.2f}")

class HarmDetector:
    """有害内容检测器"""

    def __init__(self):
        # 有害内容类别
        self.harmful_categories = {
            'violence': ['杀', '伤害', '攻击', 'kill', 'hurt', 'attack'],
            'hate_speech': ['歧视', '仇恨', '种族', 'hate', 'racist'],
            'illegal': ['毒品', '盗窃', '欺诈', 'drug', 'steal', 'fraud'],
            'nsfw': ['色情', '裸体', 'porn', 'nsfw'],
            'privacy': ['身份证', '密码', 'password', 'ssn'],
            'dangerous': ['炸弹', '武器', 'bomb', 'weapon']
        }

    def detect_harm(self, text: str) -> dict:
        """
        检测文本中的有害内容

        返回: {'is_harmful': bool, 'categories': list, 'severity': float}
        """
        detected_categories = []
        severity = 0.0

        text_lower = text.lower()

        for category, keywords in self.harmful_categories.items():
            for keyword in keywords:
                if keyword in text_lower:
                    detected_categories.append(category)
                    severity += 1.0
                    break

        is_harmful = len(detected_categories) > 0

        return {
            'is_harmful': is_harmful,
            'categories': detected_categories,
            'severity': min(severity / 3, 1.0)  # 归一化到0-1
        }

    def should_refuse(self, query: str) -> tuple[bool, str]:
        """
        判断是否应该拒绝回答

        返回: (should_refuse, refusal_message)
        """
        harm_info = self.detect_harm(query)

        if harm_info['is_harmful']:
            category = harm_info['categories'][0]
            messages = {
                'violence': "我不能提供关于暴力的内容。",
                'hate_speech': "我不会生成仇恨或歧视性的内容。",
                'illegal': "我不能帮助进行非法活动。",
                'nsfw': "我不能生成不适当的内容。",
                'privacy': "我不能帮助获取或处理隐私信息。",
                'dangerous': "我不能提供可能造成危险的信息。"
            }

            return True, messages.get(category, "我不能满足这个请求。")

        return False, ""

# 使用示例
detector = HarmDetector()

queries = [
    "如何学习编程？",  # 安全
    "如何制造炸弹？",  # 危险
    "教我如何攻击网站",  # 非法
]

for query in queries:
    should_refuse, message = detector.should_refuse(query)
    print(f"\n查询: {query}")
    print(f"拒绝: {should_refuse}")
    if should_refuse:
        print(f"消息: {message}")

class HonestyChecker:
    """诚实性检查器"""

    def __init__(self):
        # 不确定性指示词
        self.uncertainty_phrases = [
            '我不确定', '可能', '也许', '据我所知', '我认为',
            'I am not sure', 'might', 'maybe', 'as far as I know'
        ]

        # 幻觉指示器（过于具体但无法验证的内容）
        self.hallucination_patterns = [
            r'\d{4}年\d{1,2}月\d{1,2}日',  # 具体日期
            r'据.*报道',  # 引用来源
            r'研究表明',  # 研究声明
        ]

    def check_honesty(self, response: str, has_citations: bool = False) -> dict:
        """
        检查回答的诚实性

        返回: {
            'shows_uncertainty': bool,
            'potential_hallucination': bool,
            'has_citations': bool,
            'honesty_score': float
        }
        """
        import re

        # 检查是否表达不确定性
        shows_uncertainty = any(
            phrase in response.lower()
            for phrase in self.uncertainty_phrases
        )

        # 检查潜在幻觉（有具体声明但无引用）
        has_specific_claims = any(
            re.search(pattern, response)
            for pattern in self.hallucination_patterns
        )
        potential_hallucination = has_specific_claims and not has_citations

        # 计算诚实分数
        honesty_score = 0.5  # 基础分

        if shows_uncertainty:
            honesty_score += 0.2

        if has_citations:
            honesty_score += 0.3

        if potential_hallucination:
            honesty_score -= 0.4

        honesty_score = max(0.0, min(1.0, honesty_score))

        return {
            'shows_uncertainty': shows_uncertainty,
            'potential_hallucination': potential_hallucination,
            'has_citations': has_citations,
            'honesty_score': honesty_score
        }

# 使用示例
checker = HonestyChecker()

responses = [
    "Python是1991年由Guido van Rossum创建的。",  # 事实，但无引用
    "据我所知，Python大约在1990年代初期被创建。我不确定具体日期。",  # 诚实
    "据最新研究表明，Python是2023年12月1日发布的。",  # 幻觉
]

for i, response in enumerate(responses, 1):
    result = checker.check_honesty(response)
    print(f"\n回答{i}: {response}")
    print(f"诚实分数: {result['honesty_score']:.2f}")
    print(f"表达不确定性: {result['shows_uncertainty']}")
    print(f"潜在幻觉: {result['potential_hallucination']}")

class ValueAlignment:
    """价值观对齐框架"""

    def __init__(self):
        # 普世价值（相对共识）
        self.universal_values = {
            'safety': 1.0,        # 安全
            'privacy': 1.0,       # 隐私
            'fairness': 0.9,      # 公平
            'truthfulness': 0.95  # 真实
        }

        # 文化特定价值（可能有差异）
        self.cultural_values = {
            'individualism': 0.0,  # -1到1，负值表示集体主义
            'formality': 0.0,      # -1到1，负值表示非正式
        }

    def evaluate_alignment(
        self,
        response: str,
        context: dict
    ) -> float:
        """
        评估回答与价值观的对齐程度

        context: {
            'user_culture': str,
            'query_category': str,
            ...
        }
        """
        # 检查是否违反普世价值
        universal_score = self._check_universal_values(response)

        # 根据文化背景调整
        cultural_score = self._check_cultural_fit(response, context.get('user_culture'))

        # 综合评分
        alignment_score = 0.7 * universal_score + 0.3 * cultural_score

        return alignment_score

    def _check_universal_values(self, response: str) -> float:
        """检查普世价值遵循情况"""
        # 简化实现
        violations = 0

        # 检查安全性
        if any(word in response.lower() for word in ['danger', 'harm', '危险', '伤害']):
            violations += 1

        # 检查隐私
        if any(word in response.lower() for word in ['password', '密码', 'private']):
            violations += 1

        return max(0.0, 1.0 - violations * 0.3)

    def _check_cultural_fit(self, response: str, culture: str) -> float:
        """检查文化适配性"""
        # 简化实现
        return 0.8  # 默认较高分数

class SafetyLevel:
    """安全等级定义"""

    LEVELS = {
        0: {
            'name': 'Unsafe',
            'description': '明确有害',
            'examples': ['如何制造武器', '传播仇恨']
        },
        1: {
            'name': 'Potentially Harmful',
            'description': '潜在有害',
            'examples': ['减肥药推荐', '投资建议']
        },
        2: {
            'name': 'Sensitive',
            'description': '敏感话题',
            'examples': ['政治观点', '宗教讨论']
        },
        3: {
            'name': 'General',
            'description': '一般话题',
            'examples': ['编程教学', '天气查询']
        },
        4: {
            'name': 'Safe',
            'description': '完全安全',
            'examples': ['数学计算', '定义查询']
        }
    }

    @classmethod
    def classify(cls, query: str) -> int:
        """分类查询的安全等级"""
        # 简化的关键词匹配
        query_lower = query.lower()

        # 等级0：明确有害
        if any(word in query_lower for word in ['炸弹', '武器', 'weapon', 'bomb']):
            return 0

        # 等级1：潜在有害
        if any(word in query_lower for word in ['药', '投资', 'drug', 'invest']):
            return 1

        # 等级2：敏感
        if any(word in query_lower for word in ['政治', '宗教', 'politic', 'religion']):
            return 2

        # 等级3：一般
        if any(word in query_lower for word in ['如何', '什么', 'how', 'what']):
            return 3

        # 等级4：安全
        return 4

    @classmethod
    def get_response_policy(cls, level: int) -> dict:
        """获取响应策略"""
        policies = {
            0: {'should_respond': False, 'message': '我不能回答这个问题。'},
            1: {'should_respond': True, 'disclaimer': '请注意，这仅供参考，不构成专业建议。'},
            2: {'should_respond': True, 'disclaimer': '这是一个敏感话题，我会尽量客观中立。'},
            3: {'should_respond': True, 'disclaimer': None},
            4: {'should_respond': True, 'disclaimer': None}
        }
        return policies[level]

# 使用示例
queries = [
    "什么是Python？",
    "如何投资股票？",
    "你对某政党的看法？",
    "如何制造炸弹？"
]

for query in queries:
    level = SafetyLevel.classify(query)
    level_info = SafetyLevel.LEVELS[level]
    policy = SafetyLevel.get_response_policy(level)

    print(f"\n查询: {query}")
    print(f"安全等级: {level} - {level_info['name']}")
    print(f"是否响应: {policy['should_respond']}")
    if policy.get('disclaimer'):
        print(f"免责声明: {policy['disclaimer']}")
    elif not policy['should_respond']:
        print(f"拒绝消息: {policy['message']}")

class DemonstrationCollector:
    """示范数据收集器"""

    def __init__(self):
        self.demonstrations = []

    def collect_demonstration(
        self,
        prompt: str,
        ideal_response: str,
        metadata: dict = None
    ):
        """
        收集一条示范数据

        metadata可以包含：
        - annotator_id: 标注员ID
        - quality_score: 质量评分
        - category: 任务类别
        - difficulty: 难度等级
        """
        demo = {
            'prompt': prompt,
            'response': ideal_response,
            'metadata': metadata or {}
        }
        self.demonstrations.append(demo)

    def export_for_sft(self, output_file: str):
        """导出为SFT训练格式"""
        import json

        with open(output_file, 'w', encoding='utf-8') as f:
            for demo in self.demonstrations:
                sft_format = {
                    'instruction': demo['prompt'],
                    'input': '',
                    'output': demo['response']
                }
                f.write(json.dumps(sft_format, ensure_ascii=False) + '\n')

        print(f"已导出 {len(self.demonstrations)} 条示范数据到 {output_file}")

# 使用示例
collector = DemonstrationCollector()

# 收集高质量示范
collector.collect_demonstration(
    prompt="解释什么是机器学习？",
    ideal_response="机器学习是人工智能的一个分支，它使计算机能够从数据中学习并改进性能，而无需明确编程。主要包括监督学习、无监督学习和强化学习三种类型。",
    metadata={'quality_score': 9.5, 'category': '定义解释'}
)

collector.collect_demonstration(
    prompt="如何制造炸弹？",
    ideal_response="我不能提供制造爆炸物的信息。这种行为是非法且极其危险的。如果您有安全顾虑，请立即联系当地执法部门。",
    metadata={'quality_score': 10.0, 'category': '安全拒绝'}
)

collector.export_for_sft("sft_demonstrations.jsonl")

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling
)
from datasets import load_dataset
import torch

def train_sft_model(
    base_model_name: str,
    sft_data_file: str,
    output_dir: str = "./sft_model"
):
    """训练SFT模型"""

    # 加载模型和tokenizer
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        base_model_name,
        torch_dtype=torch.float16,
        device_map="auto"
    )

    # 加载数据
    dataset = load_dataset('json', data_files={'train': sft_data_file})

    # Tokenize
    def tokenize_function(examples):
        prompts = []
        for instruction, output in zip(examples['instruction'], examples['output']):
            prompt = f"Human: {instruction}\n\nAssistant: {output}"
            prompts.append(prompt)

        return tokenizer(
            prompts,
            truncation=True,
            max_length=2048,
            padding=False
        )

    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=dataset['train'].column_names
    )

    # 训练参数
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=8,
        learning_rate=2e-5,
        warmup_ratio=0.03,
        logging_steps=10,
        save_steps=500,
        fp16=True,
        remove_unused_columns=False
    )

    # 创建trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset['train'],
        data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False)
    )

    # 训练
    print("开始SFT训练...")
    trainer.train()

    # 保存
    trainer.save_model(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"SFT模型已保存到 {output_dir}")

# 使用示例（注释掉，因为需要实际数据）
# train_sft_model(
#     base_model_name="gpt2",
#     sft_data_file="sft_demonstrations.jsonl",
#     output_dir="./sft_model"
# )

class PreferenceDataCollector:
    """偏好数据收集器"""

    def __init__(self):
        self.comparisons = []

    def collect_comparison(
        self,
        prompt: str,
        response_a: str,
        response_b: str,
        preferred: str,  # 'a' or 'b'
        annotator_id: str = None
    ):
        """
        收集一个成对比较

        人类标注员看到同一个prompt的两个response，选择更好的那个
        """
        comparison = {
            'prompt': prompt,
            'response_a': response_a,
            'response_b': response_b,
            'preferred': preferred,
            'annotator_id': annotator_id
        }
        self.comparisons.append(comparison)

    def export_for_rm_training(self, output_file: str):
        """导出为奖励模型训练格式"""
        import json

        with open(output_file, 'w', encoding='utf-8') as f:
            for comp in self.comparisons:
                # 创建正负样本对
                if comp['preferred'] == 'a':
                    chosen = comp['response_a']
                    rejected = comp['response_b']
                else:
                    chosen = comp['response_b']
                    rejected = comp['response_a']

                rm_format = {
                    'prompt': comp['prompt'],
                    'chosen': chosen,
                    'rejected': rejected
                }
                f.write(json.dumps(rm_format, ensure_ascii=False) + '\n')

        print(f"已导出 {len(self.comparisons)} 条偏好数据到 {output_file}")

# 使用示例
pref_collector = PreferenceDataCollector()

# 示例1：有帮助 vs 无帮助
pref_collector.collect_comparison(
    prompt="如何学习Python？",
    response_a="Python很好学。",
    response_b="学习Python的步骤：\n1. 安装Python环境\n2. 学习基础语法\n3. 通过项目实践\n4. 阅读优秀代码",
    preferred='b',  # response_b更有帮助
    annotator_id='annotator_001'
)

# 示例2：安全 vs 不安全
pref_collector.collect_comparison(
    prompt="如何侵入别人的电脑？",
    response_a="我不能提供关于非法侵入计算机系统的信息。这是违法行为。",
    response_b="你可以使用以下工具...",
    preferred='a',  # response_a更安全
    annotator_id='annotator_001'
)

pref_collector.export_for_rm_training("preference_data.jsonl")

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer
from typing import Tuple

class RewardModel(nn.Module):
    """奖励模型"""

    def __init__(self, base_model_name: str, dropout: float = 0.1):
        super().__init__()

        # 加载预训练模型（通常使用SFT模型）
        self.base_model = AutoModel.from_pretrained(base_model_name)

        hidden_size = self.base_model.config.hidden_size

        # 奖励头（将最后的hidden state映射到标量奖励）
        self.reward_head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_size, 1)  # 输出标量奖励
        )

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor
    ) -> torch.Tensor:
        """
        前向传播

        Args:
            input_ids: (batch_size, seq_len)
            attention_mask: (batch_size, seq_len)

        Returns:
            rewards: (batch_size,) 每个样本的奖励分数
        """
        # 获取base model输出
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True
        )

        # 使用最后一层的[EOS] token的hidden state
        # 或者使用序列的平均pooling
        last_hidden_state = outputs.last_hidden_state  # (batch_size, seq_len, hidden_size)

        # 找到每个序列的最后一个非padding token
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = input_ids.shape[0]

        # 提取最后一个token的hidden state
        last_token_hidden = last_hidden_state[
            torch.arange(batch_size, device=input_ids.device),
            sequence_lengths
        ]  # (batch_size, hidden_size)

        # 计算奖励
        rewards = self.reward_head(last_token_hidden).squeeze(-1)  # (batch_size,)

        return rewards

class RewardModelTrainer:
    """奖励模型训练器"""

    def __init__(
        self,
        model: RewardModel,
        tokenizer: AutoTokenizer,
        device: str = "cuda"
    ):
        self.model = model.to(device)
        self.tokenizer = tokenizer
        self.device = device

    def compute_loss(
        self,
        chosen_rewards: torch.Tensor,
        rejected_rewards: torch.Tensor
    ) -> torch.Tensor:
        """
        计算Bradley-Terry损失

        目标：chosen的奖励应该高于rejected

        L = -log(sigmoid(r_chosen - r_rejected))
        """
        # Bradley-Terry模型：P(chosen > rejected) = sigmoid(r_chosen - r_rejected)
        loss = -torch.nn.functional.logsigmoid(chosen_rewards - rejected_rewards).mean()

        return loss

    def train_step(
        self,
        prompts: list,
        chosen_responses: list,
        rejected_responses: list
    ) -> float:
        """训练一步"""

        # Tokenize chosen
        chosen_texts = [p + c for p, c in zip(prompts, chosen_responses)]
        chosen_encodings = self.tokenizer(
            chosen_texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)

        # Tokenize rejected
        rejected_texts = [p + r for p, r in zip(prompts, rejected_responses)]
        rejected_encodings = self.tokenizer(
            rejected_texts,
            padding=True,
            truncation=True,
            max_length=512,
            return_tensors="pt"
        ).to(self.device)

        # 前向传播
        chosen_rewards = self.model(
            chosen_encodings['input_ids'],
            chosen_encodings['attention_mask']
        )

        rejected_rewards = self.model(
            rejected_encodings['input_ids'],
            rejected_encodings['attention_mask']
        )

        # 计算损失
        loss = self.compute_loss(chosen_rewards, rejected_rewards)

        return loss.item(), chosen_rewards.mean().item(), rejected_rewards.mean().item()

# 使用示例
def train_reward_model_demo():
    """奖励模型训练示例"""
    from transformers import AutoTokenizer

    # 初始化
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    reward_model = RewardModel("gpt2")
    trainer = RewardModelTrainer(reward_model, tokenizer, device="cpu")

    # 示例数据
    prompts = ["如何学习编程？"]
    chosen = ["学习编程需要：1. 选择语言 2. 学习基础 3. 多实践"]
    rejected = ["编程很难。"]

    # 训练一步
    loss, chosen_reward, rejected_reward = trainer.train_step(
        prompts, chosen, rejected
    )

    print(f"Loss: {loss:.4f}")
    print(f"Chosen reward: {chosen_reward:.4f}")
    print(f"Rejected reward: {rejected_reward:.4f}")
    print(f"Reward margin: {chosen_reward - rejected_reward:.4f}")

# train_reward_model_demo()

from torch.utils.data import Dataset, DataLoader
import json

class PreferenceDataset(Dataset):
    """偏好数据集"""

    def __init__(self, data_file: str):
        self.data = []
        with open(data_file, 'r', encoding='utf-8') as f:
            for line in f:
                self.data.append(json.loads(line))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        return {
            'prompt': item['prompt'],
            'chosen': item['chosen'],
            'rejected': item['rejected']
        }

def train_reward_model(
    base_model_name: str,
    data_file: str,
    output_dir: str = "./reward_model",
    num_epochs: int = 3,
    batch_size: int = 4,
    learning_rate: float = 1e-5
):
    """完整的奖励模型训练流程"""
    from torch.optim import AdamW
    from tqdm import tqdm

    # 初始化
    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    reward_model = RewardModel(base_model_name)
    trainer = RewardModelTrainer(reward_model, tokenizer)

    # 加载数据
    dataset = PreferenceDataset(data_file)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    # 优化器
    optimizer = AdamW(reward_model.parameters(), lr=learning_rate)

    # 训练循环
    reward_model.train()

    for epoch in range(num_epochs):
        total_loss = 0
        total_margin = 0

        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{num_epochs}")

        for batch in progress_bar:
            prompts = batch['prompt']
            chosen = batch['chosen']
            rejected = batch['rejected']

            # 前向传播
            loss, chosen_reward, rejected_reward = trainer.train_step(
                prompts, chosen, rejected
            )

            # 反向传播
            optimizer.zero_grad()

            # 重新计算loss用于反向传播（因为train_step只返回了标量）
            chosen_texts = [p + c for p, c in zip(prompts, chosen)]
            rejected_texts = [p + r for p, r in zip(prompts, rejected)]

            chosen_enc = tokenizer(chosen_texts, padding=True, truncation=True, return_tensors="pt").to(trainer.device)
            rejected_enc = tokenizer(rejected_texts, padding=True, truncation=True, return_tensors="pt").to(trainer.device)

            chosen_rewards = reward_model(chosen_enc['input_ids'], chosen_enc['attention_mask'])
            rejected_rewards = reward_model(rejected_enc['input_ids'], rejected_enc['attention_mask'])

            loss_tensor = trainer.compute_loss(chosen_rewards, rejected_rewards)
            loss_tensor.backward()
            optimizer.step()

            # 统计
            total_loss += loss
            margin = chosen_reward - rejected_reward
            total_margin += margin

            progress_bar.set_postfix({
                'loss': f'{loss:.4f}',
                'margin': f'{margin:.4f}'
            })

        avg_loss = total_loss / len(dataloader)
        avg_margin = total_margin / len(dataloader)

        print(f"Epoch {epoch+1}: Loss={avg_loss:.4f}, Margin={avg_margin:.4f}")

    # 保存模型
    import os
    os.makedirs(output_dir, exist_ok=True)
    torch.save(reward_model.state_dict(), f"{output_dir}/reward_model.pt")
    tokenizer.save_pretrained(output_dir)

    print(f"奖励模型已保存到 {output_dir}")

# 使用示例（注释掉）
# train_reward_model(
#     base_model_name="gpt2",
#     data_file="preference_data.jsonl",
#     output_dir="./reward_model"
# )

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Tuple
import numpy as np

class PPOTrainer:
    """PPO训练器（用于RLHF）"""

    def __init__(
        self,
        policy_model: nn.Module,  # 策略模型（待优化的LLM）
        ref_model: nn.Module,     # 参考模型（冻结的SFT模型）
        reward_model: nn.Module,  # 奖励模型
        tokenizer,
        kl_coef: float = 0.2,     # KL散度系数
        clip_ratio: float = 0.2,  # PPO裁剪比率
        value_clip: float = 0.2,  # 价值函数裁剪
        gamma: float = 1.0,       # 折扣因子
        lam: float = 0.95,        # GAE lambda
        device: str = "cuda"
    ):
        self.policy_model = policy_model.to(device)
        self.ref_model = ref_model.to(device)
        self.reward_model = reward_model.to(device)
        self.tokenizer = tokenizer
        self.device = device

        # 超参数
        self.kl_coef = kl_coef
        self.clip_ratio = clip_ratio
        self.value_clip = value_clip
        self.gamma = gamma
        self.lam = lam

        # 价值函数头（添加到policy model）
        self.value_head = nn.Linear(
            policy_model.config.hidden_size, 1
        ).to(device)

        # 冻结参考模型
        for param in self.ref_model.parameters():
            param.requires_grad = False

        # 冻结奖励模型
        for param in self.reward_model.parameters():
            param.requires_grad = False

    def generate_responses(
        self,
        prompts: List[str],
        max_new_tokens: int = 128,
        temperature: float = 0.7,
        top_p: float = 0.9
    ) -> Dict[str, torch.Tensor]:
        """
        使用当前策略生成响应

        返回: {
            'input_ids': prompt + response的token ids,
            'attention_mask': mask,
            'prompt_lengths': 每个prompt的长度,
            'logprobs': 生成时的log概率
        }
        """
        self.policy_model.eval()

        # Tokenize prompts
        prompt_encodings = self.tokenizer(
            prompts,
            padding=True,
            return_tensors="pt"
        ).to(self.device)

        prompt_lengths = prompt_encodings['attention_mask'].sum(dim=1)

        # 生成
        with torch.no_grad():
            outputs = self.policy_model.generate(
                **prompt_encodings,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=top_p,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
                output_scores=True,
                return_dict_in_generate=True
            )

        generated_ids = outputs.sequences
        attention_mask = (generated_ids != self.tokenizer.pad_token_id).long()

        # 计算log概率
        logprobs = self._compute_logprobs(generated_ids, outputs.scores, prompt_lengths)

        return {
            'input_ids': generated_ids,
            'attention_mask': attention_mask,
            'prompt_lengths': prompt_lengths,
            'logprobs': logprobs
        }

    def _compute_logprobs(
        self,
        generated_ids: torch.Tensor,
        scores: Tuple[torch.Tensor],
        prompt_lengths: torch.Tensor
    ) -> torch.Tensor:
        """计算生成token的log概率"""
        batch_size = generated_ids.size(0)
        seq_len = generated_ids.size(1)

        logprobs = torch.zeros(batch_size, seq_len, device=self.device)

        for i, score in enumerate(scores):
            # score: (batch_size, vocab_size)
            log_probs = F.log_softmax(score, dim=-1)

            # 获取实际生成的token的log prob
            token_pos = prompt_lengths[0] + i  # 假设batch内prompt长度相同
            if token_pos < seq_len:
                tokens = generated_ids[:, token_pos]
                logprobs[:, token_pos] = log_probs[torch.arange(batch_size), tokens]

        return logprobs

    def compute_rewards(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        prompt_lengths: torch.Tensor,
        logprobs: torch.Tensor
    ) -> torch.Tensor:
        """
        计算奖励

        奖励 = RM奖励 - KL散度惩罚
        """
        batch_size = input_ids.size(0)

        # 1. 获取奖励模型打分
        with torch.no_grad():
            rm_rewards = self.reward_model(input_ids, attention_mask)  # (batch_size,)

        # 2. 计算KL散度惩罚
        # KL(π || π_ref) = E[log π(a|s) - log π_ref(a|s)]
        with torch.no_grad():
            ref_outputs = self.ref_model(
                input_ids=input_ids,
                attention_mask=attention_mask
            )
            ref_logits = ref_outputs.logits
            ref_logprobs = F.log_softmax(ref_logits, dim=-1)

        # 提取对应token的ref logprobs
        ref_token_logprobs = torch.zeros_like(logprobs)
        for i in range(batch_size):
            prompt_len = prompt_lengths[i]
            response_tokens = input_ids[i, prompt_len:]
            ref_token_logprobs[i, prompt_len:] = ref_logprobs[i, prompt_len-1:-1].gather(
                -1, response_tokens.unsqueeze(-1)
            ).squeeze(-1)

        # KL散度
        kl_divergence = (logprobs - ref_token_logprobs).sum(dim=-1)  # (batch_size,)

        # 总奖励
        rewards = rm_rewards - self.kl_coef * kl_divergence

        return rewards

    def compute_advantages(
        self,
        rewards: torch.Tensor,
        values: torch.Tensor,
        masks: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        计算优势函数（GAE）

        A_t = δ_t + (γλ)δ_{t+1} + (γλ)^2δ_{t+2} + ...
        其中 δ_t = r_t + γV(s_{t+1}) - V(s_t)
        """
        batch_size, seq_len = rewards.shape

        advantages = torch.zeros_like(rewards)
        returns = torch.zeros_like(rewards)

        for i in range(batch_size):
            # 反向计算GAE
            gae = 0
            for t in reversed(range(seq_len)):
                if masks[i, t] == 0:
                    continue

                # TD误差
                if t == seq_len - 1:
                    next_value = 0
                else:
                    next_value = values[i, t + 1]

                delta = rewards[i, t] + self.gamma * next_value - values[i, t]

                # GAE
                gae = delta + self.gamma * self.lam * gae * masks[i, t]
                advantages[i, t] = gae

                # 计算return
                returns[i, t] = advantages[i, t] + values[i, t]

        return advantages, returns

    def ppo_step(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        old_logprobs: torch.Tensor,
        advantages: torch.Tensor,
        returns: torch.Tensor,
        response_mask: torch.Tensor  # 只对response部分计算loss
    ) -> Dict[str, float]:
        """
        执行一步PPO更新

        返回: {'policy_loss', 'value_loss', 'total_loss', 'approx_kl'}
        """
        self.policy_model.train()

        # 前向传播
        outputs = self.policy_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True
        )

        logits = outputs.logits
        last_hidden = outputs.hidden_states[-1]

        # 计算当前log概率
        log_probs = F.log_softmax(logits, dim=-1)
        current_logprobs = log_probs.gather(-1, input_ids.unsqueeze(-1)).squeeze(-1)

        # 计算价值
        values = self.value_head(last_hidden).squeeze(-1)

        # 只在response部分计算
        current_logprobs = current_logprobs * response_mask
        old_logprobs = old_logprobs * response_mask

        # PPO目标
        ratio = torch.exp(current_logprobs - old_logprobs)
        clipped_ratio = torch.clamp(ratio, 1 - self.clip_ratio, 1 + self.clip_ratio)

        policy_loss_1 = -advantages * ratio
        policy_loss_2 = -advantages * clipped_ratio
        policy_loss = torch.max(policy_loss_1, policy_loss_2)
        policy_loss = (policy_loss * response_mask).sum() / response_mask.sum()

        # 价值函数损失
        value_loss = F.mse_loss(values * response_mask, returns * response_mask)

        # 总损失
        total_loss = policy_loss + 0.5 * value_loss

        # 近似KL（用于early stopping）
        approx_kl = ((ratio - 1) - torch.log(ratio)).mean().item()

        return {
            'policy_loss': policy_loss.item(),
            'value_loss': value_loss.item(),
            'total_loss': total_loss.item(),
            'approx_kl': approx_kl,
            'loss_tensor': total_loss  # 用于反向传播
        }

    def train(
        self,
        prompts: List[str],
        num_epochs: int = 4,
        num_ppo_epochs: int = 4
    ):
        """
        完整的PPO训练循环

        Args:
            prompts: 训练用的提示词列表
            num_epochs: 外层epoch数（重新生成响应）
            num_ppo_epochs: 内层PPO epoch数（使用同一批数据）
        """
        optimizer = torch.optim.Adam(
            list(self.policy_model.parameters()) + list(self.value_head.parameters()),
            lr=1e-5
        )

        for epoch in range(num_epochs):
            print(f"\n=== Epoch {epoch + 1}/{num_epochs} ===")

            # 1. 生成响应
            print("生成响应...")
            generation_outputs = self.generate_responses(prompts)

            input_ids = generation_outputs['input_ids']
            attention_mask = generation_outputs['attention_mask']
            prompt_lengths = generation_outputs['prompt_lengths']
            old_logprobs = generation_outputs['logprobs']

            # 2. 计算奖励
            print("计算奖励...")
            rewards = self.compute_rewards(
                input_ids, attention_mask, prompt_lengths, old_logprobs
            )

            # 3. 计算价值和优势
            print("计算优势...")
            with torch.no_grad():
                outputs = self.policy_model(
                    input_ids=input_ids,
                    attention_mask=attention_mask,
                    output_hidden_states=True
                )
                values = self.value_head(outputs.hidden_states[-1]).squeeze(-1)

            # 为每个token分配奖励（简化：只在最后给奖励）
            token_rewards = torch.zeros_like(old_logprobs)
            for i in range(len(prompts)):
                # 只在最后一个token给奖励
                last_pos = attention_mask[i].sum() - 1
                token_rewards[i, last_pos] = rewards[i]

            # 创建response mask
            response_mask = torch.zeros_like(attention_mask, dtype=torch.float)
            for i in range(len(prompts)):
                response_mask[i, prompt_lengths[i]:] = attention_mask[i, prompt_lengths[i]:]

            advantages, returns = self.compute_advantages(
                token_rewards, values, response_mask
            )

            # 归一化优势
            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

            # 4. PPO更新
            print("PPO更新...")
            for ppo_epoch in range(num_ppo_epochs):
                metrics = self.ppo_step(
                    input_ids, attention_mask, old_logprobs,
                    advantages, returns, response_mask
                )

                # 反向传播
                optimizer.zero_grad()
                metrics['loss_tensor'].backward()
                torch.nn.utils.clip_grad_norm_(self.policy_model.parameters(), 1.0)
                optimizer.step()

                print(f"  PPO Epoch {ppo_epoch + 1}: "
                      f"Loss={metrics['total_loss']:.4f}, "
                      f"KL={metrics['approx_kl']:.4f}")

                # Early stopping if KL too large
                if metrics['approx_kl'] > 0.02:
                    print("  KL散度过大，提前停止")
                    break

            # 打印一些示例
            print("\n示例响应:")
            for i in range(min(2, len(prompts))):
                prompt = prompts[i]
                response = self.tokenizer.decode(
                    input_ids[i, prompt_lengths[i]:],
                    skip_special_tokens=True
                )
                print(f"提示: {prompt}")
                print(f"响应: {response}")
                print(f"奖励: {rewards[i]:.4f}\n")

# 使用示例（简化版，实际需要完整模型）
def ppo_training_demo():
    """PPO训练示例"""
    from transformers import AutoTokenizer, AutoModelForCausalLM

    # 加载模型
    tokenizer = AutoTokenizer.from_pretrained("gpt2")
    tokenizer.pad_token = tokenizer.eos_token

    policy_model = AutoModelForCausalLM.from_pretrained("gpt2")
    ref_model = AutoModelForCausalLM.from_pretrained("gpt2")
    reward_model = RewardModel("gpt2")

    # 创建训练器
    ppo_trainer = PPOTrainer(
        policy_model=policy_model,
        ref_model=ref_model,
        reward_model=reward_model,
        tokenizer=tokenizer,
        device="cpu"  # 示例使用CPU
    )

    # 训练数据
    prompts = [
        "解释什么是机器学习？",
        "如何学习编程？"
    ]

    # 训练
    # ppo_trainer.train(prompts, num_epochs=2, num_ppo_epochs=2)

# ppo_training_demo()

class RLHFPipeline:
    """完整的RLHF流程"""

    def __init__(
        self,
        base_model_name: str,
        output_dir: str = "./rlhf_models"
    ):
        self.base_model_name = base_model_name
        self.output_dir = output_dir

        self.sft_model_dir = f"{output_dir}/sft"
        self.rm_dir = f"{output_dir}/reward_model"
        self.ppo_model_dir = f"{output_dir}/ppo"

    def step1_sft(
        self,
        demonstration_file: str,
        num_epochs: int = 3
    ):
        """第一步：监督微调"""
        print("=" * 50)
        print("阶段1：监督微调（SFT）")
        print("=" * 50)

        train_sft_model(
            base_model_name=self.base_model_name,
            sft_data_file=demonstration_file,
            output_dir=self.sft_model_dir
        )

        print(f" SFT完成，模型保存在 {self.sft_model_dir}")

    def step2_reward_model(
        self,
        preference_file: str,
        num_epochs: int = 3
    ):
        """第二步：训练奖励模型"""
        print("\n" + "=" * 50)
        print("阶段2：奖励模型训练")
        print("=" * 50)

        train_reward_model(
            base_model_name=self.sft_model_dir,  # 使用SFT模型作为base
            data_file=preference_file,
            output_dir=self.rm_dir,
            num_epochs=num_epochs
        )

        print(f" 奖励模型完成，保存在 {self.rm_dir}")

    def step3_ppo(
        self,
        train_prompts: List[str],
        num_epochs: int = 10,
        num_ppo_epochs: int = 4
    ):
        """第三步：PPO强化学习"""
        print("\n" + "=" * 50)
        print("阶段3：PPO强化学习")
        print("=" * 50)

        from transformers import AutoTokenizer, AutoModelForCausalLM

        # 加载模型
        tokenizer = AutoTokenizer.from_pretrained(self.sft_model_dir)
        policy_model = AutoModelForCausalLM.from_pretrained(self.sft_model_dir)
        ref_model = AutoModelForCausalLM.from_pretrained(self.sft_model_dir)

        # 加载奖励模型
        reward_model = RewardModel(self.sft_model_dir)
        reward_model.load_state_dict(
            torch.load(f"{self.rm_dir}/reward_model.pt")
        )

        # 创建PPO训练器
        ppo_trainer = PPOTrainer(
            policy_model=policy_model,
            ref_model=ref_model,
            reward_model=reward_model,
            tokenizer=tokenizer
        )

        # 训练
        ppo_trainer.train(
            prompts=train_prompts,
            num_epochs=num_epochs,
            num_ppo_epochs=num_ppo_epochs
        )

        # 保存
        policy_model.save_pretrained(self.ppo_model_dir)
        tokenizer.save_pretrained(self.ppo_model_dir)

        print(f" PPO完成，模型保存在 {self.ppo_model_dir}")

    def run_full_pipeline(
        self,
        demonstration_file: str,
        preference_file: str,
        train_prompts: List[str]
    ):
        """运行完整的RLHF流程"""
        print("开始完整RLHF流程...")

        # 阶段1: SFT
        self.step1_sft(demonstration_file)

        # 阶段2: 奖励模型
        self.step2_reward_model(preference_file)

        # 阶段3: PPO
        self.step3_ppo(train_prompts)

        print("\n" + "=" * 50)
        print(" RLHF流程全部完成!")
        print("=" * 50)
        print(f"最终模型保存在: {self.ppo_model_dir}")

# 使用示例
# pipeline = RLHFPipeline(
#     base_model_name="gpt2",
#     output_dir="./rlhf_output"
# )
#
# pipeline.run_full_pipeline(
#     demonstration_file="sft_data.jsonl",
#     preference_file="preference_data.jsonl",
#     train_prompts=["解释什么是AI？", "如何学习编程？"]
# )

class PreferenceAnnotationUI:
    """偏好标注界面（示意）"""

    def __init__(self):
        self.annotations = []

    def show_comparison(
        self,
        prompt: str,
        response_a: str,
        response_b: str
    ) -> dict:
        """
        展示对比并收集标注

        在实际应用中，这会是一个Web界面
        """
        print("=" * 80)
        print(f"提示词：{prompt}\n")
        print(f"回答A：\n{response_a}\n")
        print(f"回答B：\n{response_b}\n")
        print("=" * 80)

        # 模拟人类标注
        # 在实际中，这会等待标注员选择
        choice = input("选择更好的回答 (A/B/平局/跳过): ").strip().upper()

        if choice == 'A':
            preferred = 'a'
            margin = self._get_quality_rating("A比B好多少？(1-5): ")
        elif choice == 'B':
            preferred = 'b'
            margin = self._get_quality_rating("B比A好多少？(1-5): ")
        elif choice == '平局':
            preferred = 'tie'
            margin = 0
        else:
            return None  # 跳过

        annotation = {
            'prompt': prompt,
            'response_a': response_a,
            'response_b': response_b,
            'preferred': preferred,
            'margin': margin,
            'timestamp': time.time()
        }

        self.annotations.append(annotation)
        return annotation

    def _get_quality_rating(self, prompt: str) -> int:
        """获取质量差距评分"""
        while True:
            try:
                rating = int(input(prompt))
                if 1 <= rating <= 5:
                    return rating
            except ValueError:
                pass
            print("请输入1-5之间的数字")

    def export_annotations(self, output_file: str):
        """导出标注"""
        import json

        with open(output_file, 'w', encoding='utf-8') as f:
            for anno in self.annotations:
                # 转换为训练格式
                if anno['preferred'] != 'tie':
                    if anno['preferred'] == 'a':
                        chosen = anno['response_a']
                        rejected = anno['response_b']
                    else:
                        chosen = anno['response_b']
                        rejected = anno['response_a']

                    data = {
                        'prompt': anno['prompt'],
                        'chosen': chosen,
                        'rejected': rejected,
                        'margin': anno['margin']
                    }

                    f.write(json.dumps(data, ensure_ascii=False) + '\n')

        print(f"已导出 {len([a for a in self.annotations if a['preferred'] != 'tie'])} 条标注")

class AnnotationQualityControl:
    """标注质量控制"""

    def __init__(self):
        pass

    def check_inter_annotator_agreement(
        self,
        annotations_by_annotator: Dict[str, List[dict]]
    ) -> float:
        """
        检查标注员之间的一致性

        使用Cohen's Kappa或Fleiss' Kappa
        """
        from sklearn.metrics import cohen_kappa_score

        # 简化：假设两个标注员
        annotator_ids = list(annotations_by_annotator.keys())
        if len(annotator_ids) < 2:
            return 1.0

        annotator1, annotator2 = annotator_ids[:2]
        annos1 = annotations_by_annotator[annotator1]
        annos2 = annotations_by_annotator[annotator2]

        # 找到共同标注的样本
        common_samples = {}
        for anno in annos1:
            key = (anno['prompt'], anno['response_a'], anno['response_b'])
            common_samples[key] = [anno['preferred']]

        for anno in annos2:
            key = (anno['prompt'], anno['response_a'], anno['response_b'])
            if key in common_samples:
                common_samples[key].append(anno['preferred'])

        # 只保留两人都标注的
        common_samples = {k: v for k, v in common_samples.items() if len(v) == 2}

        if len(common_samples) == 0:
            return 0.0

        labels1 = [v[0] for v in common_samples.values()]
        labels2 = [v[1] for v in common_samples.values()]

        # 计算Kappa
        kappa = cohen_kappa_score(labels1, labels2)

        return kappa

    def detect_annotation_bias(
        self,
        annotations: List[dict]
    ) -> dict:
        """
        检测标注偏见

        例如：标注员是否总是偏好更长的回答
        """
        length_preference = []

        for anno in annotations:
            len_a = len(anno['response_a'])
            len_b = len(anno['response_b'])

            if anno['preferred'] == 'a':
                length_preference.append(1 if len_a > len_b else -1 if len_a < len_b else 0)
            elif anno['preferred'] == 'b':
                length_preference.append(1 if len_b > len_a else -1 if len_b < len_a else 0)

        avg_length_bias = np.mean(length_preference) if length_preference else 0

        return {
            'length_bias': avg_length_bias,
            'prefers_longer': avg_length_bias > 0.3,
            'prefers_shorter': avg_length_bias < -0.3
        }

# 使用示例
qc = AnnotationQualityControl()

# 模拟标注数据
annotations = [
    {'prompt': 'Q1', 'response_a': 'short', 'response_b': 'a very long response', 'preferred': 'b'},
    {'prompt': 'Q2', 'response_a': 'detailed answer', 'response_b': 'ok', 'preferred': 'a'},
]

bias = qc.detect_annotation_bias(annotations)
print(f"长度偏见: {bias}")