Math-Shepherd中使用的过程奖励模型（mistral-7b）。

输入：问题 + 带有特殊步骤标记ки的逐步解答，例如：

珍妮特的鸭子每天产16个蛋。她每天早上吃三个当早餐，还烤......？步骤1：珍妮特的鸭子每天产16个蛋。ки\n步骤2：她每天早上吃三个当早餐，所以剩下16 - 3 = 13个蛋。ки\n步骤3：她每天用四个蛋为朋友烤松饼，所以剩下13 - 4 = 9个蛋。ки\n步骤4：她每天将剩余的蛋以每个2美元的价格在农贸市场出售，因此她每天在农贸市场赚取9 * 2 = 18美元。答案是：18 ки

输出：逻辑值。需后处理以获得每个步骤的评分。

from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
import torch

good_token = '+'
bad_token = '-'
step_tag = 'ки'

tokenizer = AutoTokenizer.from_pretrained('peiyi9979/math-shepherd-mistral-7b-prm')
candidate_tokens = tokenizer.encode(f"{good_token} {bad_token}")[1:] # [648, 387]
step_tag_id = tokenizer.encode(f"{step_tag}")[-1] # 12902
model = AutoModelForCausalLM.from_pretrained('peiyi9979/math-shepherd-mistral-7b-prm').eval()

question = """珍妮特的鸭子每天产16个蛋。她每天早上吃三个当早餐，还每天用四个蛋为朋友烤松饼。她每天将剩余的蛋以每个2美元的价格在农贸市场出售。她每天在农贸市场赚多少钱？"""
output1 = """步骤1：珍妮特的鸭子每天产16个蛋。ки\n步骤2：她每天早上吃三个当早餐，所以剩下16 - 3 = 13个蛋。ки\n步骤3：她每天用四个蛋为朋友烤松饼，所以剩下13 - 4 = 9个蛋。ки\n步骤4：她每天将剩余的蛋以每个2美元的价格在农贸市场出售，因此她每天赚取9 * 2 = 18美元。答案是：18 ки""" # 18正确
output2 = """步骤1：珍妮特的鸭子每天产16个蛋。ки\n步骤2：她每天早上吃三个当早餐，所以剩下16 - 3 = 13个蛋。ки\n步骤3：她每天用四个蛋为朋友烤松饼，所以剩下13 - 4 = 9个蛋。ки\n步骤4：她每天将剩余的蛋以每个2美元的价格在农贸市场出售，因此她每天赚取9 * 2 = 17美元。答案是：17 ки""" # 17错误

for output in [output1, output2]:
    input_for_prm = f"{question} {output}"
    input_id = torch.tensor([tokenizer.encode(input_for_prm)])

    with torch.no_grad():
        logits = model(input_id).logits[:,:,candidate_tokens]
        scores = logits.softmax(dim=-1)[:,:,0] 
        step_scores = scores[input_id == step_tag_id]
        print(step_scores)
        
# tensor([0.9955, 0.9958, 0.9983, 0.9957])
# tensor([0.9955, 0.9958, 0.9983, 0.0240])