许可协议: apache-2.0
支持语言:

• 英语
• 中文
• 德语
• 俄语
  基础模型:
• meta-llama/Llama-3.1-8B
  标签:
• 翻译
• 推理
• 测试时优化

Plan2Align奖励模型，用于中文->英语、中文->德语、中文->俄语语言对的测试时翻译任务。

@article{wang2025plan2align,
  title={Plan2Align: 基于预测规划的段落级机器翻译测试时偏好对齐},
  author={王匡达、陈腾瑞、洪宇恒、丁硕阳、吴岳华、王裕强、杨朝翰、彭文志、谢品骏},
  journal={arXiv预印本 arXiv:2502.20795},
  year={2025}
}

使用奖励模型

RM = AutoModelForCausalLMWithValueHead.from_pretrained('ray24724919/plan2align_rm', torch_dtype=torch_dtype)
RM.eval()
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B-Instruct")
RM.gradient_checkpointing_enable() # 如需启用
        
value_head_weights = load_file("valuehead权重文件路径.safetensors")
new_state_dict = {key.replace("v_head.", "") if key.startswith("v_head.") else key: value for key, value in value_head_weights.items()}
RM.v_head.load_state_dict(new_state_dict)

奖励函数

def reward(目标语言, 原文, 译文, device='cuda:0'):
    对话记录=[{"role": "system", "content":'你是一个专业的翻译助手，只需输出翻译结果。'},
          {"role": "user", "content": f"### 将以下中文翻译成{目标语言}，中文:\n{原文}\n### {目标语言}:"},
          {"role": "assistant", "content": 译文}]
    输入标记 = tokenizer.apply_chat_template(
        对话记录,
        add_generation_prompt=False,
        return_tensors="pt"
    ).to(device)

    输入数据 = {
        "input_ids": 输入标记,
        "attention_mask": torch.ones_like(输入标记, device=device)
    }

    with torch.no_grad():
        输出结果 = model(**输入数据, return_value=True)
        奖励分数 = 输出结果[2]

    最终奖励 = 奖励分数[:, -1].item()

    return 最终奖励

翻译奖励建模的系统提示

对话模板 = [{"role": "system", "content": "你是一个专业的翻译助手，只需输出翻译结果。"},
            {"role": "user", "content": f"### 将以下中文翻译成{目标语言}，中文:\n{原文}\n### {目标语言}:"},
            {"role": "assistant", "content": 译文}]