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

class RewardModel(nn.Module):
    def __init__(self, model_name="bert-base-chinese"):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.reward_head = nn.Linear(self.model.config.hidden_size, 1)
    
    def forward(self, prompts, responses):
        """计算奖励分数"""
        # 拼接prompt和response
        texts = [f"{p} {r}" for p, r in zip(prompts, responses)]
        
        # 编码
        inputs = self.tokenizer(
            texts, 
            return_tensors="pt", 
            padding=True, 
            truncation=True,
            max_length=512
        )
        
        # 获取表示
        outputs = self.model(**inputs)
        pooled_output = outputs.pooler_output
        
        # 计算奖励分数
        rewards = self.reward_head(pooled_output)
        return rewards.squeeze(-1)
    
    def train_step(self, batch):
        """训练步骤"""
        prompts, responses_a, responses_b, preferences = batch
        
        # 计算奖励分数
        rewards_a = self.forward(prompts, responses_a)
        rewards_b = self.forward(prompts, responses_b)
        
        # 计算损失
        # 使用Bradley-Terry模型
        logits = rewards_a - rewards_b
        loss = -torch.mean(preferences * torch.logsigmoid(logits) + 
                          (1 - preferences) * torch.logsigmoid(-logits))
        
        return loss