Views
No views yet
@article{zhang2024general,
title={General Preference Modeling with Preference Representations for Aligning Language Models},
author={Zhang, Yifan and Zhang, Ge and Wu, Yue and Xu, Kangping and Gu, Quanquan},
journal={arXiv preprint arXiv:2410.02197},
year={2024}
}1from typing import Optional, List, Dict
2import torch
3import torch.nn as nn
4from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
5import torch.nn.functional as F
6from transformers import AutoTokenizer
7
8def get_tokenizer(pretrain, model, padding_side="left", use_fast=True):
9 tokenizer = AutoTokenizer.from_pretrained(pretrain, trust_remote_code=True, use_fast=use_fast)
10 tokenizer.padding_side = padding_side
11 if tokenizer.pad_token is None:
12 tokenizer.pad_token = tokenizer.eos_token
13 tokenizer.pad_token_id = tokenizer.eos_token_id
14 model.config.pad_token_id = tokenizer.pad_token_id
15 return tokenizer
16
17def get_reward_model(base_causal_model, base_llm_model, is_general_preference: bool=False, add_prompt_head: bool=False, value_head_dim: int=2):
18 class CustomRewardModel(base_causal_model):
19
20 def __init__(self, config: AutoConfig):
21 super().__init__(config)
22 setattr(self, self.base_model_prefix, base_llm_model(config))
23 if not is_general_preference:
24 self.value_head = nn.Linear(config.hidden_size, 1, bias=False)
25 else:
26 self.value_head = nn.Linear(config.hidden_size, value_head_dim, bias=False)
27 if add_prompt_head:
28 self.prompt_head = nn.Linear(config.hidden_size, value_head_dim // 2, bias=False)
29
30 self.is_general_preference = is_general_preference
31
32 self.post_init()
33
34 def custom_forward(
35 self,
36 input_ids: torch.LongTensor = None,
37 attention_mask: Optional[torch.Tensor] = None,
38 return_output=False,
39 ) -> torch.Tensor:
40 position_ids = attention_mask.long().cumsum(-1) - 1
41 position_ids.masked_fill_(attention_mask == 0, 1)
42 outputs = getattr(self, self.base_model_prefix)(
43 input_ids, attention_mask=attention_mask, position_ids=position_ids
44 )
45 last_hidden_states = outputs["last_hidden_state"]
46
47 if not self.is_general_preference:
48 values = self.value_head(last_hidden_states).squeeze(-1)
49 # left padding in training mode
50 if self.training:
51 reward = values[:, -1]
52 else:
53 eos_indices = attention_mask.size(1) - 1 - attention_mask.long().fliplr().argmax(dim=1, keepdim=True)
54 reward = values.gather(dim=1, index=eos_indices).squeeze(1)
55 if return_output:
56 return reward, outputs
57 else:
58 return reward, None
59 else:
60 values = self.value_head(last_hidden_states)
61 # left padding in training mode
62 if self.training:
63 reward = values[:, -1, :]
64 reward = F.normalize(reward, p=2, dim=-1) # Shape will be [batch_size, value_head_dim]
65 else:
66 eos_indices = attention_mask.size(1) - 1 - attention_mask.long().fliplr().argmax(dim=1)
67 eos_indices = eos_indices.unsqueeze(1) # Change shape to [batch_size, 1]
68 reward_list = []
69 for dim in range(value_head_dim):
70 reward_list.append(values[:,:,dim].gather(dim=1, index=eos_indices))
71 reward = torch.cat(reward_list, dim=1)
72 reward = F.normalize(reward, p=2, dim=-1) # Shape will be [batch_size, value_head_dim]
73 if return_output:
74 return reward, outputs
75 else:
76 return reward, None
77
78 def create_skew_symmetric_block_matrix(self, dim, device, dtype, prompt_hidden_states):
79 """
80 Create a batch of skew-symmetric block matrices where each matrix is data-dependent on
81 the corresponding prompt_hidden_states. Only the relevant block diagonal parts are generated.
82
83 Args:
84 - dim: Dimension of the square matrix (must be even).
85 - prompt_hidden_states: Tensor of shape [batch_size, hidden_dim].
86
87 Returns:
88 - batch_R_matrices: Tensor of shape [batch_size, dim, dim], with skew-symmetric block entries.
89 """
90 if hasattr(self, 'prompt_head'):
91 batch_size = prompt_hidden_states.shape[0]
92
93 # Ensure that dim is even, as we're creating blocks of size 2x2
94 assert dim % 2 == 0, "dim must be even for skew-symmetric block generation"
95
96 # Pass through the linear layer to get the block diagonal entries (half of the matrix's off-diagonal blocks)
97 block_values = self.prompt_head(prompt_hidden_states).view(batch_size, dim // 2)
98 block_values = torch.softmax(block_values, dim=-1)
99
100 # Create a batch of zero matrices [batch_size, dim, dim]
101 batch_R_matrices = torch.zeros((batch_size, dim, dim), device=device, dtype=dtype)
102
103 # Fill only the block diagonal entries with the learned values
104 for i in range(0, dim, 2):
105 batch_R_matrices[:, i, i + 1] = -block_values[:, i // 2]
106 batch_R_matrices[:, i + 1, i] = block_values[:, i // 2] # Skew-symmetric condition
107 else:
108 raise AttributeError("prompt_head is not defined. Ensure 'add_prompt_head' is set to True during initialization.")
109
110 return batch_R_matrices
111
112 return CustomRewardModel
113
114def generate_high_dim_result_with_prompt(model, value_head_dim, chosen_reward, rejected_reward, prompt_hidden_states):
115 R_matrix = model.create_skew_symmetric_block_matrix(value_head_dim, chosen_reward.device, chosen_reward.dtype, prompt_hidden_states)
116 if chosen_reward.device == rejected_reward.device == R_matrix.device:
117 transformed_chosen = torch.bmm(chosen_reward.view(chosen_reward.shape[0], 1, value_head_dim), R_matrix.transpose(1, 2))
118 result = torch.bmm(transformed_chosen, rejected_reward.view(rejected_reward.shape[0], value_head_dim, 1))
119 result = result.view(chosen_reward.shape[0])
120 return result
121
122class GPMPipeline:
123 def __init__(self, model_name_or_path, device=torch.device("cuda:0"), is_general_preference: bool=True, add_prompt_head: bool=True, value_head_dim: int=2, bf16: bool=True, truncation: bool=True, max_length: int=4096, padding: bool=True, tau: float=0.1):
124 self.device = device
125 self.is_general_preference = is_general_preference
126 self.add_prompt_head = add_prompt_head
127 self.value_head_dim = value_head_dim
128 self.truncation = truncation
129 self.max_length = max_length
130 self.padding = padding
131 self.tau = tau
132
133 config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
134 config._attn_implementation = "flash_attention_2"
135 base_class = AutoModel._model_mapping[type(config)]
136 base_causal_class = AutoModelForCausalLM._model_mapping.get(type(config), None)
137 cls_class = get_reward_model(base_causal_class, base_class, is_general_preference, add_prompt_head, value_head_dim)
138
139 # configure model
140 self.model = cls_class.from_pretrained(
141 model_name_or_path,
142 config=config,
143 trust_remote_code=True,
144 torch_dtype=torch.bfloat16 if bf16 else "auto",
145 )
146 # configure tokenizer
147 self.tokenizer = get_tokenizer(model_name_or_path, self.model, "left", use_fast=True)
148 self.tokenizer.truncation_side = "right"
149
150 # prepare model
151 self.model.to(device)
152 self.model.eval()
153
154 def __call__(self, samples: List[List[Dict[str, str]]], return_prompt=False):
155 input_texts = [self.tokenizer.apply_chat_template(sample, tokenize=False) for sample in samples]
156
157 inputs = self.tokenizer(
158 input_texts,
159 truncation=True,
160 max_length=self.max_length,
161 padding=True,
162 return_tensors="pt",
163 ).to(self.device)
164
165 inputs["input_ids"][:, -1] = self.tokenizer.eos_token_id
166 inputs["attention_mask"][:, -1] = 1
167
168 with torch.no_grad():
169 rewards, outputs = self.model.custom_forward(**inputs, return_output=return_prompt)
170
171 chosen_response_len_list = []
172 if return_prompt:
173 prompt_texts = [self.tokenizer.apply_chat_template([sample[0]], tokenize=False) for sample in samples]
174 for i in range(len(input_texts)):
175 prompt_token = self.tokenizer(
176 prompt_texts[i],
177 max_length=self.max_length,
178 padding=False,
179 truncation=True,
180 return_tensors="pt",
181 )
182 chosen_token = self.tokenizer(
183 input_texts[i],
184 max_length=self.max_length,
185 padding=False,
186 truncation=True,
187 return_tensors="pt",
188 )
189 chosen_response_len = chosen_token["attention_mask"].sum() - prompt_token["attention_mask"].sum()
190 chosen_response_len_list.append(chosen_response_len)
191 chosen_response_len = torch.tensor(chosen_response_len_list).view(-1, 1).to(self.device)
192 if return_prompt:
193 chosen_last_hidden_states = outputs["last_hidden_state"]
194 prompt_end_index = chosen_last_hidden_states.size(1) - chosen_response_len - 1
195 prompt_end_index_expanded = prompt_end_index.unsqueeze(-1).expand(-1, -1, chosen_last_hidden_states.size(-1))
196 prompt_hidden_state = torch.gather(chosen_last_hidden_states, dim=1, index=prompt_end_index_expanded).squeeze(1)
197 return rewards, prompt_hidden_state
198 else:
199 return rewards
200
201
202prompt_text = "Describe the importance of reading books in today's digital age."
203response1 = "Books remain crucial in the digital era, offering in-depth knowledge and fostering critical thinking. They provide a unique, immersive experience that digital media can't replicate, contributing significantly to personal and intellectual growth."
204response2 = "Books are still useful for learning new things. They help you relax and can be a good break from screens."
205
206context1 = [
207 {"role": "user", "content": prompt_text},
208 {"role": "assistant", "content": response1}
209]
210
211context2 = [
212 {"role": "user", "content": prompt_text},
213 {"role": "assistant", "content": response2}
214]
215
216rm = GPMPipeline("general-preference/GPM-Llama-3.1-8B", value_head_dim=6)
217
218reward1, prompt_hidden_state = rm([context1], return_prompt=True)
219reward2 = rm([context2])
220
221result = generate_high_dim_result_with_prompt(rm.model, rm.value_head_dim, reward1, reward2, prompt_hidden_state)
222
223result_batch = result.float().cpu().detach().numpy().tolist()
224
225results = []
226[
227 results.append(1) if result > 0 else results.append(0)
228 for result in result_batch
229]
230
231print(result_batch)
232
233