File size: 2,429 Bytes
9accbc6 06008f2 4a941c2 9accbc6 4a941c2 9accbc6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import logging
from typing import List
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoConfig, MistralForCausalLM
from safetensors.torch import load_file
class RewardModel:
def __init__(self, model_dir) -> None:
config = AutoConfig.from_pretrained(model_dir)
# config._attn_implementation = "flash_attention_2"
self.device = torch.device('cuda')
self.model = MistralForCausalLM(config)
self.model.lm_head = nn.Linear(config.hidden_size, 1, bias=False)
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
state_dict = load_file(f"{model_dir}/model.safetensors")
self.model.load_state_dict(state_dict, strict=False)
self.model.to(dtype=torch.bfloat16)
self.model.to(device=self.device)
self.model.eval()
logging.info("Load model completed.")
@torch.no_grad()
def score(self, prompts, chosens) -> List[float]:
# Concat prompt and chosen, append eos_id
input_ids_list = [self.tokenizer.encode(prompt) + self.tokenizer.encode(chosen) + [self.tokenizer.eos_token_id] for prompt, chosen in zip(prompts, chosens)]
# Pad sequences to the maximum length
max_length = max(len(ids) for ids in input_ids_list)
padded_input_ids = [ids + [self.tokenizer.pad_token_id or self.tokenizer.eos_token_id] * (max_length - len(ids)) for ids in input_ids_list]
# Forward pass
input_ids = torch.tensor(padded_input_ids).to(device=self.device)
logits = self.model(input_ids).logits
# Extract logits corresponding to eos_token_id positions
scores = []
for i, input_ids in enumerate(input_ids_list):
eos_position = input_ids.index(self.tokenizer.eos_token_id)
eos_logit = logits[i, eos_position, :].squeeze().item()
scores.append(eos_logit)
return scores
if __name__ == '__main__':
local_model_dir = "Your local model dir"
model_dir = f"{local_model_dir}/Seed-X-RM-7B"
prompt = ["Translate the following English sentence into Chinese:\nMay the force be with you <zh>", "Translate the following English sentence into Chinese:\nMay the force be with you <zh>"]
candidate = ["愿原力与你同在","愿力量与你同在"]
model = RewardModel(model_dir)
scores = model.score(prompt, candidate) # output [1.46875, -0.376953125]
print(scores) |