wehub-resource-sync
460 lines
18 KiB
Python
460 lines
18 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, List, Optional, Tuple, Union
|
|
|
|
import paddle
|
|
import paddle.nn.functional as F
|
|
from paddle import nn
|
|
|
|
try:
|
|
from paddle.distributed.fleet.utils.sequence_parallel_utils import GatherOp
|
|
except:
|
|
pass
|
|
|
|
import paddlenlp
|
|
from paddlenlp.transformers import (
|
|
LlamaConfig,
|
|
LlamaLMHead,
|
|
LlamaModel,
|
|
LlamaPretrainedModel,
|
|
LlamaPretrainingCriterion,
|
|
MistralLMHead,
|
|
MistralModel,
|
|
MistralPreTrainedModel,
|
|
MistralPretrainingCriterion,
|
|
PretrainedConfig,
|
|
PretrainedModel,
|
|
)
|
|
from paddlenlp.transformers.conversion_utils import (
|
|
StateDictNameMapping,
|
|
init_name_mappings,
|
|
)
|
|
from paddlenlp.transformers.model_outputs import (
|
|
CausalLMOutputWithCrossAttentions,
|
|
CausalLMOutputWithPast,
|
|
)
|
|
from paddlenlp.utils.log import logger
|
|
|
|
|
|
class LlamaModelForScore(LlamaPretrainedModel):
|
|
_keys_to_ignore_on_load_missing = ["lm_head.weight"]
|
|
|
|
def __init__(self, config: PretrainedConfig, **kwargs: Any) -> None:
|
|
super().__init__(config)
|
|
self.llama = LlamaModel(config)
|
|
self.score_head = nn.Linear(config.hidden_size, 1, bias_attr=False)
|
|
|
|
def get_input_embeddings(self) -> nn.Embedding:
|
|
return self.llama.embed_tokens
|
|
|
|
def set_input_embeddings(self, value: nn.Embedding) -> None:
|
|
self.llama.embed_tokens = value
|
|
|
|
def get_decoder(self) -> PretrainedModel:
|
|
return self.llama
|
|
|
|
def set_decoder(self, decoder: PretrainedModel) -> None:
|
|
self.llama = decoder
|
|
|
|
def forward( # pylint: disable=too-many-arguments
|
|
self,
|
|
input_ids: paddle.Tensor,
|
|
attention_mask: paddle.Tensor | None = None,
|
|
position_ids: paddle.Tensor | None = None,
|
|
past_key_values: list[paddle.Tensor] | None = None,
|
|
inputs_embeds: paddle.Tensor | None = None,
|
|
use_cache: bool | None = None,
|
|
output_attentions: bool | None = None,
|
|
output_hidden_states: bool | None = None,
|
|
return_dict: bool | None = None,
|
|
attn_mask_startend_row_indices: paddle.Tensor | None = None,
|
|
response_indexs: paddle.Tensor | None = None,
|
|
):
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
outputs = self.llama(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
attn_mask_startend_row_indices=attn_mask_startend_row_indices,
|
|
)
|
|
hidden_states = outputs[0] # size = (B, L, E)
|
|
if self.config.sequence_parallel:
|
|
hidden_states = GatherOp.apply(hidden_states)
|
|
hidden_states = paddle.reshape_(hidden_states, [-1, self.config.seq_length, self.config.hidden_size])
|
|
chosen_indexes = paddle.to_tensor(
|
|
[[response_index[0], response_index[1]] for response_index in response_indexs]
|
|
)
|
|
rejected_indexes = paddle.to_tensor(
|
|
[[response_index[0], response_index[2]] for response_index in response_indexs]
|
|
)
|
|
chosen_hidden_states = hidden_states.gather_nd(chosen_indexes)
|
|
rejected_hidden_states = hidden_states.gather_nd(rejected_indexes)
|
|
|
|
chosen_scores = self.score_head(chosen_hidden_states)
|
|
rejected_scores = self.score_head(rejected_hidden_states)
|
|
loss = -F.log_sigmoid(chosen_scores - rejected_scores).mean()
|
|
return loss, chosen_scores, rejected_scores
|
|
|
|
@classmethod
|
|
def _get_name_mappings(cls, config: LlamaConfig) -> list[StateDictNameMapping]:
|
|
mappings: list[StateDictNameMapping] = []
|
|
model_mappings = [
|
|
["embed_tokens.weight"],
|
|
["norm.weight"],
|
|
]
|
|
for layer_index in range(config.num_hidden_layers):
|
|
layer_mappings = [
|
|
[f"layers.{layer_index}.self_attn.q_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.self_attn.k_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.self_attn.v_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.self_attn.o_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.self_attn.rotary_emb.inv_freq"],
|
|
[f"layers.{layer_index}.mlp.gate_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.mlp.down_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.mlp.up_proj.weight", None, "transpose"],
|
|
[f"layers.{layer_index}.input_layernorm.weight"],
|
|
[f"layers.{layer_index}.post_attention_layernorm.weight"],
|
|
]
|
|
model_mappings.extend(layer_mappings)
|
|
|
|
init_name_mappings(mappings=model_mappings)
|
|
# base-model prefix "LlamaModel"
|
|
if "LlamaModel" not in config.architectures:
|
|
for mapping in model_mappings:
|
|
mapping[0] = "model." + mapping[0]
|
|
mapping[1] = "llama." + mapping[1]
|
|
model_mappings.append(["lm_head.weight", "lm_head.weight", "transpose"])
|
|
model_mappings.extend(
|
|
[
|
|
["score_head.weight", "score_head.weight", "transpose"],
|
|
["normalizer.var", "normalizer.var"],
|
|
["normalizer.mean", "normalizer.mean"],
|
|
["normalizer.count", "normalizer.count"],
|
|
]
|
|
)
|
|
|
|
mappings = [StateDictNameMapping(*mapping, index=index) for index, mapping in enumerate(model_mappings)]
|
|
return mappings
|
|
|
|
|
|
class LlamaModelForPRM(LlamaPretrainedModel):
|
|
_keys_to_ignore_on_load_missing = ["lm_head.weight"]
|
|
|
|
def __init__(
|
|
self, config: PretrainedConfig, placeholder_token_id: int, reward_token_ids: List, **kwargs: Any
|
|
) -> None:
|
|
super().__init__(config)
|
|
self.config = config
|
|
self.llama = LlamaModel(config)
|
|
if config.tie_word_embeddings:
|
|
self.lm_head = LlamaLMHead(config, embedding_weights=self.llama.embed_tokens.weight, transpose_y=True)
|
|
self.tie_weights()
|
|
else:
|
|
self.lm_head = LlamaLMHead(config)
|
|
self.criterion = LlamaPretrainingCriterion(config)
|
|
self.IGNORE_INDEX = -100
|
|
self.loss = nn.CrossEntropyLoss(ignore_index=self.IGNORE_INDEX)
|
|
self.placeholder_token_id = placeholder_token_id
|
|
self.reward_token_ids = reward_token_ids
|
|
|
|
def get_input_embeddings(self):
|
|
return self.llama.embed_tokens
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.llama.embed_tokens = value
|
|
|
|
def get_output_embeddings(self):
|
|
return self.lm_head
|
|
|
|
def set_output_embeddings(self, new_embeddings):
|
|
self.lm_head = new_embeddings
|
|
|
|
def set_decoder(self, decoder):
|
|
self.llama = decoder
|
|
|
|
def get_decoder(self):
|
|
return self.llama
|
|
|
|
def prepare_inputs_for_generation(
|
|
self, input_ids, use_cache=False, past_key_values=None, inputs_embeds=None, **kwargs
|
|
):
|
|
batch_size, seq_length = input_ids.shape
|
|
position_ids = kwargs.get("position_ids", paddle.arange(seq_length).expand((batch_size, seq_length)))
|
|
attention_mask = kwargs.get("attention_mask", None)
|
|
if past_key_values:
|
|
input_ids = input_ids[:, -1].unsqueeze(axis=-1)
|
|
position_ids = position_ids[:, -1].unsqueeze(-1)
|
|
|
|
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
|
if inputs_embeds is not None and past_key_values is None:
|
|
model_inputs = {"inputs_embeds": inputs_embeds}
|
|
else:
|
|
model_inputs = {"input_ids": input_ids}
|
|
|
|
model_inputs.update(
|
|
{
|
|
"position_ids": position_ids,
|
|
"past_key_values": past_key_values,
|
|
"use_cache": use_cache,
|
|
"attention_mask": attention_mask,
|
|
}
|
|
)
|
|
return model_inputs
|
|
|
|
def _get_model_inputs_spec(self, dtype: str):
|
|
return {
|
|
"input_ids": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
|
|
"attention_mask": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
|
|
"position_ids": paddle.static.InputSpec(shape=[None, None], dtype="int64"),
|
|
}
|
|
|
|
@staticmethod
|
|
def update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder=False):
|
|
# update cache
|
|
if isinstance(outputs, tuple) and len(outputs) > 1 and not isinstance(outputs[1], paddle.Tensor):
|
|
model_kwargs["past_key_values"] = outputs[1]
|
|
|
|
if isinstance(outputs, CausalLMOutputWithCrossAttentions) and "past_key_values" in outputs:
|
|
model_kwargs["past_key_values"] = outputs.past_key_values
|
|
|
|
# update position_ids
|
|
if "position_ids" in model_kwargs and model_kwargs["position_ids"] is not None:
|
|
position_ids = model_kwargs["position_ids"]
|
|
model_kwargs["position_ids"] = paddle.concat([position_ids, position_ids[..., -1:] + 1], axis=-1)
|
|
|
|
if not is_encoder_decoder and "attention_mask" in model_kwargs and model_kwargs["attention_mask"] is not None:
|
|
attention_mask = model_kwargs["attention_mask"]
|
|
model_kwargs["attention_mask"] = paddle.concat(
|
|
[attention_mask, paddle.ones([attention_mask.shape[0], 1], dtype=attention_mask.dtype)], axis=-1
|
|
)
|
|
|
|
return model_kwargs
|
|
|
|
def forward(
|
|
self,
|
|
input_ids=None,
|
|
position_ids=None,
|
|
attention_mask=None,
|
|
inputs_embeds=None,
|
|
labels=None,
|
|
use_cache=False,
|
|
past_key_values=None,
|
|
output_attentions=None,
|
|
output_hidden_states=None,
|
|
return_dict=None,
|
|
attn_mask_startend_row_indices=None,
|
|
):
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
if attn_mask_startend_row_indices is not None and attention_mask is not None:
|
|
logger.warning(
|
|
"You have provided both attn_mask_startend_row_indices and attention_mask. "
|
|
"The attn_mask_startend_row_indices will be used."
|
|
)
|
|
attention_mask = None
|
|
|
|
outputs = self.llama(
|
|
input_ids, # [bs, seq_len]
|
|
position_ids=position_ids,
|
|
attention_mask=attention_mask,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
past_key_values=past_key_values,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
attn_mask_startend_row_indices=attn_mask_startend_row_indices,
|
|
)
|
|
|
|
hidden_states = outputs[0] # [bs, seq_len, dim]
|
|
|
|
logits = self.lm_head(hidden_states)
|
|
|
|
placeholder_mask = input_ids == self.placeholder_token_id
|
|
logits = logits[placeholder_mask]
|
|
labels = labels[placeholder_mask]
|
|
|
|
logits = logits[..., self.reward_token_ids]
|
|
for idx, token in enumerate(self.reward_token_ids):
|
|
labels = paddle.where(labels == token, idx, labels)
|
|
|
|
loss = self.loss(logits, labels)
|
|
|
|
if labels.dtype == logits.dtype:
|
|
labels = labels.argmax(dim=-1)
|
|
acc = (logits.argmax(axis=-1) == labels).astype("float32").mean()
|
|
|
|
return loss, acc
|
|
|
|
|
|
class MistralModelForPRM(MistralPreTrainedModel):
|
|
_keys_to_ignore_on_load_missing = ["lm_head.weight"]
|
|
_tied_weights_keys = ["lm_head.weight"]
|
|
|
|
def __init__(self, config, placeholder_token_id: int, reward_token_ids: List, **kwargs: Any) -> None:
|
|
super().__init__(config)
|
|
self.mistral = MistralModel(config)
|
|
self.vocab_size = config.vocab_size
|
|
self.lm_head = MistralLMHead(config)
|
|
self.criterion = MistralPretrainingCriterion(config)
|
|
self.IGNORE_INDEX = -100
|
|
self.loss = nn.CrossEntropyLoss(ignore_index=self.IGNORE_INDEX)
|
|
self.placeholder_token_id = placeholder_token_id
|
|
self.reward_token_ids = reward_token_ids
|
|
|
|
def get_input_embeddings(self):
|
|
return self.mistral.embed_tokens
|
|
|
|
def set_input_embeddings(self, value):
|
|
self.mistral.embed_tokens = value
|
|
|
|
def get_output_embeddings(self):
|
|
return self.lm_head
|
|
|
|
def set_output_embeddings(self, new_embeddings):
|
|
self.lm_head = new_embeddings
|
|
|
|
def set_decoder(self, decoder):
|
|
self.mistral = decoder
|
|
|
|
def get_decoder(self):
|
|
return self.mistral
|
|
|
|
def prepare_inputs_for_generation(
|
|
self, input_ids, use_cache=False, past_key_values=None, inputs_embeds=None, **kwargs
|
|
):
|
|
batch_size, seq_length = input_ids.shape
|
|
position_ids = kwargs.get("position_ids", paddle.arange(seq_length).expand((batch_size, seq_length)))
|
|
attention_mask = kwargs.get("attention_mask", None)
|
|
if past_key_values:
|
|
input_ids = input_ids[:, -1].unsqueeze(axis=-1)
|
|
position_ids = position_ids[:, -1].unsqueeze(-1)
|
|
|
|
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
|
|
if inputs_embeds is not None and past_key_values is None:
|
|
model_inputs = {"inputs_embeds": inputs_embeds}
|
|
else:
|
|
model_inputs = {"input_ids": input_ids}
|
|
|
|
model_inputs.update(
|
|
{
|
|
"position_ids": position_ids,
|
|
"past_key_values": past_key_values,
|
|
"use_cache": use_cache,
|
|
"attention_mask": attention_mask,
|
|
}
|
|
)
|
|
return model_inputs
|
|
|
|
@staticmethod
|
|
def update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder=False):
|
|
# update cache
|
|
if isinstance(outputs, tuple) and len(outputs) > 1 and not isinstance(outputs[1], paddle.Tensor):
|
|
model_kwargs["past_key_values"] = outputs[1]
|
|
|
|
if isinstance(outputs, CausalLMOutputWithCrossAttentions) and "past_key_values" in outputs:
|
|
model_kwargs["past_key_values"] = outputs.past_key_values
|
|
|
|
# update position_ids
|
|
if "position_ids" in model_kwargs and model_kwargs["position_ids"] is not None:
|
|
position_ids = model_kwargs["position_ids"]
|
|
model_kwargs["position_ids"] = paddle.concat([position_ids, position_ids[..., -1:] + 1], axis=-1)
|
|
|
|
if not is_encoder_decoder and "attention_mask" in model_kwargs:
|
|
attention_mask = model_kwargs.pop("attention_mask", None)
|
|
|
|
if attention_mask is not None and len(attention_mask.shape) == 2:
|
|
model_kwargs["attention_mask"] = paddle.concat(
|
|
[attention_mask, paddle.ones([attention_mask.shape[0], 1], dtype=attention_mask.dtype)], axis=-1
|
|
)
|
|
|
|
return model_kwargs
|
|
|
|
def forward(
|
|
self,
|
|
input_ids: paddle.Tensor = None,
|
|
attention_mask: Optional[paddle.Tensor] = None,
|
|
position_ids: Optional[paddle.Tensor] = None,
|
|
past_key_values: Optional[List[paddle.Tensor]] = None,
|
|
inputs_embeds: Optional[paddle.Tensor] = None,
|
|
labels: Optional[paddle.Tensor] = None,
|
|
use_cache: Optional[bool] = None,
|
|
output_attentions: Optional[bool] = None,
|
|
output_hidden_states: Optional[bool] = None,
|
|
return_dict: Optional[bool] = None,
|
|
) -> Union[Tuple, CausalLMOutputWithPast]:
|
|
|
|
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
|
output_hidden_states = (
|
|
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
|
)
|
|
output_hidden_states = True # TODO: for align
|
|
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
|
|
|
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
|
outputs = self.mistral(
|
|
input_ids=input_ids,
|
|
attention_mask=attention_mask,
|
|
position_ids=position_ids,
|
|
past_key_values=past_key_values,
|
|
inputs_embeds=inputs_embeds,
|
|
use_cache=use_cache,
|
|
output_attentions=output_attentions,
|
|
output_hidden_states=output_hidden_states,
|
|
return_dict=return_dict,
|
|
)
|
|
|
|
hidden_states = outputs[0] # [bs, seq_len, dim]
|
|
|
|
logits = self.lm_head(hidden_states)
|
|
|
|
placeholder_mask = input_ids == self.placeholder_token_id
|
|
logits = logits[placeholder_mask]
|
|
labels = labels[placeholder_mask]
|
|
|
|
logits = logits[..., self.reward_token_ids]
|
|
for idx, token in enumerate(self.reward_token_ids):
|
|
labels = paddle.where(labels == token, idx, labels)
|
|
|
|
loss = self.loss(logits, labels)
|
|
|
|
if labels.dtype == logits.dtype:
|
|
labels = labels.argmax(dim=-1)
|
|
acc = (logits.argmax(axis=-1) == labels).astype("float32").mean()
|
|
|
|
return loss, acc
|
|
|
|
|
|
paddlenlp.transformers.LlamaModelForScore = LlamaModelForScore
|
|
paddlenlp.transformers.LlamaModelForPRM = LlamaModelForPRM
|
|
paddlenlp.transformers.MistralModelForPRM = MistralModelForPRM
|