import torch import torch.nn as nn from loguru import logger as log from torch.distributions import Categorical from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessorList, TopKLogitsWarper, TopPLogitsWarper, \ TemperatureLogitsWarper class DiscreteActor(nn.Module): def __init__(self, state_dim, action_dim, hidden_dim=1024): super(DiscreteActor, self).__init__() self.state_dim = state_dim self.action_dim = action_dim self.ln = nn.LayerNorm(state_dim) self.linear1 = nn.Linear(self.state_dim, hidden_dim) self.linear2 = nn.Linear(hidden_dim, int(hidden_dim / 4)) self.output_linear = nn.Linear(int(hidden_dim / 4), self.action_dim) def forward(self, state): state = self.ln(state) x = torch.relu(self.linear1(state)) x = torch.relu(self.linear2(x)) output = torch.softmax(self.output_linear(x), dim=1) return output def sample(self, state): state = state.unsqueeze(0) prob = self.forward(state) distribution = Categorical(torch.Tensor(prob)) sample_action = distribution.sample().unsqueeze(-1).detach() z = (prob == 0.0).float() * 1e-8 logprob = torch.log(prob + z) greedy_action = torch.argmax(prob, dim=-1).unsqueeze(-1) # 1d tensor return sample_action, prob, logprob, greedy_action def select_action(self, state): state = state.unsqueeze(0) prob = self.forward(state) action = torch.argmax(prob, dim=-1).unsqueeze(-1) # 1d tensor action = action.squeeze().tolist() return action class MARAGenerator(): def __init__( self, agent_path, base_model_path, state_dim=4096, hidden_dim=1024, model_device="cuda:0", max_new_token=2048, topk=40, topp=0.95, temperature=0.8 ): # 文本生成模型初始化 self.base_model_path = base_model_path self.model_device = torch.device(model_device) self.base_model = AutoModelForCausalLM.from_pretrained(self.base_model_path).to( self.model_device).eval().requires_grad_(False) self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_path) self.tokenizer.pad_token_id = self.tokenizer.eos_token_id self.topk = topk self.topp = topp self.max_new_token = max_new_token self.temperature = temperature # instantiate logits processors and wraper self.logits_wraper = LogitsProcessorList( [TopKLogitsWarper(self.topk), TopPLogitsWarper(top_p=self.topp), TemperatureLogitsWarper(self.temperature)]) '''init agent''' self.mara_agent = self.get_mara_agent(agent_path, state_dim, hidden_dim).to(self.model_device) # log.info("mara_agent:{}".format(self.mara_agent)) self.instruction = "" # 输入指令 self.generate_ids = [] # 已生成的token self.new_token_cnt = 0 # 已生成token数 self.curr_input_ids = None # 生成new_token的输入input_id， self.curr_new_token_ids_list = None # 采样的new_token self.sum_logprobs = None self.curr_new_outputs_list = None # 以self.curr_input_ids+curr_new_token_id为输入的模型输出 self.next_new_outputs_list = None # 对应self.curr_new_outputs_list下一个token self.is_new_token = True self.chosen_indices = None self.cand_chosen_idx = 0 self.last_outputs = None self.attention_mask = None self.position_id = None # proxy_detail self.gen_info = {"gen_token_cnt": 0, # 每个样例生成的结果长度，len(gen_token_cnt_list)=sample_cnt "proxy_token_cnt": 0, # 需要经过agent进行决策的token长度，len(proxy_token_cnt_listt)=sample_cnt "cand_token_dict": {}, # 每个决策位的候选token数，1<=cand_token_cnt<=topK，0表示不需要决策 "accept_index_dict": {} # 每个决策位的选择第几个token， accept_idx } def get_mara_agent(self, agent_path, state_dim, hidden_dim): mara_agent = DiscreteActor(state_dim, 2, hidden_dim) log.info('Begin to load mara agent model from {}'.format(agent_path)) try: model_state_dict = torch.load(agent_path, map_location=self.model_device) mara_agent.load_state_dict(model_state_dict) except Exception as e: log.error("load mara_agent occur error: {}".format(str(e))) raise return mara_agent def get_input_text(self, instruction): messages = [{"role": "user", "content": instruction}] input_text = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) return input_text def get_raw_output(self, instruction, do_sample=True): if do_sample: generation_config = {"do_sample": True, "max_new_tokens": self.max_new_token, "top_k": self.topk, "top_p": self.topp, "temperature": self.temperature} else: generation_config = {"do_sample": False, "max_new_tokens": self.max_new_token} input_text = self.get_input_text(instruction) inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model_device) output_ids = self.base_model.generate(**inputs, **generation_config)[0] response = self.tokenizer.decode(output_ids[len(inputs.input_ids[0]):], skip_special_tokens=True) return {"answer": response} def get_proxy_output(self, instruction): input_text = self.get_input_text(instruction) model_inputs = self.tokenizer(input_text, return_tensors="pt").to(self.model_device) self.new_token_cnt = 0 self.generate_ids = [] self.gen_info = {"gen_token_cnt": 0, # 每个样例生成的结果长度，len(gen_token_cnt_list)=sample_cnt "mara_token_cnt": 0, # 需要经过agent进行决策的token长度，len(proxy_token_cnt_list)=sample_cnt "cand_token_dict": {}, # 每个决策位的候选token数，1<=cand_token_cnt<=topK，0表示不需要决策 "accept_index_dict": {} # 每个决策位的选择第几个token， accept_idx } input_ids = model_inputs.input_ids self.attention_mask = model_inputs.attention_mask self.curr_input_ids = input_ids self.last_outputs = self.base_model(input_ids=input_ids, attention_mask=self.attention_mask, output_hidden_states=True) # 一个token一个token地进行状态转移 self.is_new_token = True self.cand_chosen_idx = 0 self.position_id = None end_of_generate = False while not end_of_generate: end_of_generate, self.chosen_indices = self.rank_topk_ouputs_serial(self.curr_input_ids, self.last_outputs) self.is_new_token = False if end_of_generate: break accept_idx = 0 curr_new_outputs_list = [] for i in range(len(self.chosen_indices)): _, curr_new_outputs = self.rank_topk_ouputs_serial(self.curr_input_ids, self.last_outputs) curr_new_outputs_list.append(curr_new_outputs) curr_state = curr_new_outputs['hidden_states'][-1][0, -1].to(self.model_device) action = self.mara_agent.select_action(curr_state) if action == 1: accept_idx = i break self.is_new_token = True # log.info( # "new_token_cnt:{}/{}, accept_idx: {}".format(self.new_token_cnt, self.max_new_token, accept_idx)) accept_token_id = self.chosen_indices[accept_idx] self.generate_ids.append(accept_token_id) self.new_token_cnt += 1 self.gen_info["gen_token_cnt"] += 1 self.gen_info["mara_token_cnt"] += 1 if len(self.chosen_indices) not in self.gen_info["cand_token_dict"]: self.gen_info["cand_token_dict"][len(self.chosen_indices)] = 1 else: self.gen_info["cand_token_dict"][len(self.chosen_indices)] += 1 if accept_idx not in self.gen_info["accept_index_dict"]: self.gen_info["accept_index_dict"][accept_idx] = 1 else: self.gen_info["accept_index_dict"][accept_idx] += 1 self.curr_input_ids = torch.cat( (self.curr_input_ids, torch.LongTensor([[accept_token_id]]).to(self.model_device)), dim=-1) self.last_outputs = curr_new_outputs_list[accept_idx] if accept_token_id == self.tokenizer.eos_token_id or self.new_token_cnt >= self.max_new_token: end_of_generate = True completion = self.tokenizer.decode(self.generate_ids, skip_special_tokens=True) return {"answer": completion, "detail": self.gen_info} def one_step_transfer(self, pre_input_ids, past_key_values, new_token_id): attention_mask = torch.ones_like(pre_input_ids) attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) position_id = position_ids[:, -1:].to(self.model_device) new_token_id = new_token_id.to(self.model_device) new_outputs = self.base_model(input_ids=new_token_id, attention_mask=attention_mask.to(self.model_device), position_ids=position_id, past_key_values=past_key_values, output_hidden_states=True) new_input_ids = torch.cat((pre_input_ids, new_token_id), dim=-1) return new_input_ids, new_outputs def rank_topk_ouputs_serial(self, pre_input_ids, last_outputs): if self.is_new_token: end_of_generate = False next_token_logits = last_outputs.logits[:, -1, :].clone() # (batch_size, vocab_size) next_token_scores = nn.functional.log_softmax(next_token_logits, dim=-1) # (batch_size, vocab_size) next_token_scores = self.logits_wraper(pre_input_ids, next_token_scores) sorted_scores, sorted_indices = torch.sort(next_token_scores, descending=True) chosen_indices = torch.masked_select(sorted_indices, sorted_scores != -float("Inf")).tolist() # 当候选token数只有1个且不为终止eos_token_id时直接一步转移 while len(chosen_indices) == 1: self.new_token_cnt += 1 self.gen_info["gen_token_cnt"] += 1 self.generate_ids.append(chosen_indices[0]) if chosen_indices[0] == self.tokenizer.eos_token_id or self.new_token_cnt >= self.max_new_token: end_of_generate = True return end_of_generate, None else: pre_input_ids, last_outputs = self.one_step_transfer(pre_input_ids, past_key_values=last_outputs.past_key_values, new_token_id=torch.LongTensor( [[chosen_indices[0]]])) self.curr_input_ids = pre_input_ids next_token_logits = last_outputs.logits[:, -1, :].clone() # (batch_size, vocab_size) next_token_scores = nn.functional.log_softmax(next_token_logits, dim=-1) # (batch_size, vocab_size) next_token_scores = self.logits_wraper(pre_input_ids, next_token_scores) sorted_scores, sorted_indices = torch.sort(next_token_scores, descending=True) chosen_indices = torch.masked_select(sorted_indices, sorted_scores != -float("Inf")).tolist() attention_mask = torch.ones_like(pre_input_ids) attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) position_id = position_ids[:, -1:].to(self.model_device) # log.info("len(chosen_indices):{}".format(len(chosen_indices))) self.chosen_indices = chosen_indices self.cand_chosen_idx = 0 self.attention_mask = attention_mask self.position_id = position_id self.last_outputs = last_outputs return end_of_generate, chosen_indices else: new_token_id = torch.LongTensor([[self.chosen_indices[self.cand_chosen_idx]]]) curr_next_output = self.base_model(input_ids=new_token_id.to(self.model_device), attention_mask=self.attention_mask.to(self.model_device), position_ids=self.position_id.to(self.model_device), past_key_values=self.last_outputs.past_key_values, output_hidden_states=True) self.cand_chosen_idx += 1 return False, curr_next_output if __name__ == "__main__": agent_path = "../proxy_rlhf/train_result/multi_reward/mistral_v3_2_1/run2/trained_model/actor_11000.pth" base_model_path = "/mnt/public/model/huggingface/Mistral-7B-Instruct-v0.3" proxy_generator = MARAGenerator(agent_path, base_model_path) instruction = "Please introduce yourself." raw_result = proxy_generator.get_raw_output(instruction, do_sample=False) print("base model answer: ") print(raw_result["answer"]) proxy_result = proxy_generator.get_proxy_output(instruction) print("mara agent align answer: ") print(proxy_result["answer"])