Views
No views yet
1from openrlhf.models.model import get_llm_for_sequence_regression
2from transformers import AutoTokenizer
3from typing import List
4import torch
5import regex as re
6def strip_sequence(text, pad_token, eos_token):
7 pad_token_escaped = re.escape(pad_token)
8 eos_token_escaped = re.escape(eos_token)
9
10 pattern = f"^({eos_token_escaped}|{pad_token_escaped})+"
11 text = re.sub(pattern, "", text)
12
13 pattern = f"({eos_token_escaped}|{pad_token_escaped})+$"
14 text = re.sub(pattern, "", text)
15 return text
16
17class RewardModelProxy:
18 def __init__(
19 self,
20 reward_pretrain:str,
21 max_len:int,
22 batch_size:int,
23 normalize_reward:bool=False,
24 flash_attn:bool=True,
25 bf16:bool=True,
26 load_in_4bit:bool=False,
27 value_head_prefix:str="score",
28 disable_fast_tokenizer:bool=False,
29 ):
30
31 self.reward_model = get_llm_for_sequence_regression(
32 reward_pretrain,
33 "reward",
34 normalize_reward=normalize_reward,
35 use_flash_attention_2=flash_attn,
36 bf16=bf16,
37 load_in_4bit=load_in_4bit,
38 value_head_prefix=value_head_prefix,
39 device_map="cuda:5",
40 )
41 self.reward_model.eval()
42
43 self.tokenizer = AutoTokenizer.from_pretrained(reward_pretrain, trust_remote_code=True, use_fast=not disable_fast_tokenizer)
44 self.max_length = max_len
45 self.batch_size = batch_size
46
47 def get_reward(self, conversations:List[List[dict]]):
48 if self.batch_size is None:
49 batch_size = len(conversations)
50 else:
51 batch_size = self.batch_size
52
53 queries = []
54 for conversation in conversations:
55 query = self.tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
56 queries.append(query)
57
58 # remove pad_token
59 for i in range(len(queries)):
60 queries[i] = (
61 strip_sequence(queries[i], self.tokenizer.pad_token, self.tokenizer.eos_token)
62 + self.tokenizer.eos_token
63 )
64
65 scores = []
66 # batch
67 with torch.no_grad():
68 for i in range(0, len(queries), batch_size):
69 inputs = self.tokenize_fn(
70 queries[i : min(len(queries), i + batch_size)], device=self.reward_model.device
71 )
72 r = self.reward_model(inputs["input_ids"], inputs["attention_mask"])
73 r = r.tolist()
74 scores.extend(r)
75 return scores
76
77 def tokenize_fn(self, texts, device):
78 batch = self.tokenizer(
79 texts,
80 return_tensors="pt",
81 add_special_tokens=False,
82 max_length=self.max_length,
83 padding=True,
84 truncation=True,
85 )
86 return {k: v.to(device) for k, v in batch.items()}
87
88 def __call__(self, conversations:List[List[dict]]):
89 return self.get_reward(conversations)
90
91RM = RewardModelProxy(
92 "CodeDPO/Qwen2.5-Coder-7B_with_margin_scalebt",
93 max_len=2048,
94 batch_size=8,
95)
96conversations = [
97 [
98 {"role": "system", "content": "Hello, how can I help you today?"},
99 {"role": "user", "content": "I want to book a flight."},
100 ],
101]
102
103scores = RM(conversations)
104print(scores)