mara_generator.py

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"])