chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,18 @@
|
||||
from FlagEmbedding.abc.finetune.embedder import (
|
||||
AbsEmbedderDataArguments as DecoderOnlyEmbedderDataArguments,
|
||||
AbsEmbedderTrainingArguments as DecoderOnlyEmbedderTrainingArguments,
|
||||
)
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderModelArguments
|
||||
from .modeling import BiDecoderOnlyEmbedderModel
|
||||
from .trainer import DecoderOnlyEmbedderTrainer
|
||||
from .runner import DecoderOnlyEmbedderRunner
|
||||
|
||||
__all__ = [
|
||||
'DecoderOnlyEmbedderDataArguments',
|
||||
'DecoderOnlyEmbedderTrainingArguments',
|
||||
'DecoderOnlyEmbedderModelArguments',
|
||||
'BiDecoderOnlyEmbedderModel',
|
||||
'DecoderOnlyEmbedderTrainer',
|
||||
'DecoderOnlyEmbedderRunner',
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.finetune.embedder.decoder_only.base import (
|
||||
DecoderOnlyEmbedderDataArguments,
|
||||
DecoderOnlyEmbedderTrainingArguments,
|
||||
DecoderOnlyEmbedderModelArguments,
|
||||
DecoderOnlyEmbedderRunner,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((
|
||||
DecoderOnlyEmbedderModelArguments,
|
||||
DecoderOnlyEmbedderDataArguments,
|
||||
DecoderOnlyEmbedderTrainingArguments
|
||||
))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: DecoderOnlyEmbedderModelArguments
|
||||
data_args: DecoderOnlyEmbedderDataArguments
|
||||
training_args: DecoderOnlyEmbedderTrainingArguments
|
||||
|
||||
runner = DecoderOnlyEmbedderRunner(
|
||||
model_args=model_args,
|
||||
data_args=data_args,
|
||||
training_args=training_args
|
||||
)
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,76 @@
|
||||
from typing import Optional, List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderModelArguments
|
||||
|
||||
|
||||
def default_target_modules() -> List[int]:
|
||||
return ['v_proj', 'q_proj', 'k_proj', 'gate_proj', 'down_proj', 'o_proj', 'up_proj']
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecoderOnlyEmbedderModelArguments(AbsEmbedderModelArguments):
|
||||
"""
|
||||
Model argument class for decoder only base model.
|
||||
"""
|
||||
peft_model_path: str = field(
|
||||
default='', metadata={"help": "The peft model checkpoint for initialization."}
|
||||
)
|
||||
use_lora: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use LORA (low-rank parameter-efficient training) to train the model."}
|
||||
)
|
||||
lora_rank: int = field(
|
||||
default=64,
|
||||
metadata={"help": "The rank of lora."}
|
||||
)
|
||||
lora_alpha: float = field(
|
||||
default=16,
|
||||
metadata={"help": "The alpha parameter of lora."}
|
||||
)
|
||||
lora_dropout: float = field(
|
||||
default=0.1,
|
||||
metadata={"help": "The dropout rate of lora modules."}
|
||||
)
|
||||
target_modules: List[str] = field(
|
||||
default_factory=default_target_modules,
|
||||
metadata={"help": "The target modules to apply LORA."}
|
||||
)
|
||||
use_flash_attn: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use flash attention to train the model."}
|
||||
)
|
||||
use_slow_tokenizer: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library)."}
|
||||
)
|
||||
# low_cpu_mem_usage: bool = field(
|
||||
# default=False,
|
||||
# metadata={"help": "It is an option to create the model as an empty shell,"
|
||||
# "then only materialize its parameters when the pretrained weights are loaded."
|
||||
# "If passed, LLM loading time and RAM consumption will be benefited."}
|
||||
# )
|
||||
from_peft: str = field(
|
||||
default=None
|
||||
)
|
||||
modules_to_save: List[str] = field(
|
||||
default=None,
|
||||
)
|
||||
raw_peft: str = field(
|
||||
default=None
|
||||
)
|
||||
|
||||
additional_special_tokens: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "additional special tokens", "nargs": "+"}
|
||||
)
|
||||
|
||||
save_merged_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will merge the lora modules and save the entire model."}
|
||||
)
|
||||
|
||||
only_merge_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will only merge the lora modules and save the entire model."}
|
||||
)
|
||||
@@ -0,0 +1,185 @@
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
import logging
|
||||
from transformers import AutoConfig, AutoModel, AutoTokenizer
|
||||
from peft import LoraConfig, TaskType, get_peft_model, PeftModel
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderModelArguments
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def find_largest_checkpoint(checkpoint_dir):
|
||||
"""Find the largest checkpoint from directory.
|
||||
|
||||
Args:
|
||||
checkpoint_dir (str): Directory to the checkpoint.
|
||||
|
||||
Returns:
|
||||
str: Directory to the checkpoint, None no matching found.
|
||||
"""
|
||||
checkpoint_pattern = re.compile(r'checkpoint-(\d+)')
|
||||
max_number = -1
|
||||
max_checkpoint_file = None
|
||||
for file in os.listdir(checkpoint_dir):
|
||||
match = checkpoint_pattern.search(file)
|
||||
if match:
|
||||
number = int(match.group(1))
|
||||
if number > max_number:
|
||||
max_number = number
|
||||
max_checkpoint_file = file
|
||||
if max_checkpoint_file:
|
||||
return os.path.join(checkpoint_dir, max_checkpoint_file)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str, resize: bool, resize_tokens: int):
|
||||
"""Get the model.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
|
||||
output_dir (str): Directory to save the model.
|
||||
resize (bool): Whether to resize the number of tokens.
|
||||
resize_tokens (int): The new token size.
|
||||
|
||||
Returns:
|
||||
transformers.PreTrainedModel or PeftModel: The loaded model.
|
||||
"""
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = AutoModel.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
logger.info("Training new model from scratch")
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.raw_peft is not None:
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.raw_peft, 'embedding', 'emb.pth')))
|
||||
model = PeftModel.from_pretrained(model, model_args.raw_peft)
|
||||
model = model.merge_and_unload()
|
||||
|
||||
if resize:
|
||||
model.resize_token_embeddings(resize_tokens)
|
||||
os.makedirs(os.path.join(output_dir, 'embedding'), exist_ok=True)
|
||||
torch.save(model.embed_tokens, os.path.join(output_dir, 'embedding', 'emb.pth'))
|
||||
target_modules = model_args.target_modules
|
||||
else:
|
||||
target_modules = model_args.target_modules
|
||||
if 'embed_tokens' in target_modules:
|
||||
target_modules.remove('embed_tokens')
|
||||
|
||||
if model_args.from_peft is not None:
|
||||
if os.path.exists(os.path.join(model_args.from_peft, 'embedding')):
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.from_peft, 'embedding', 'emb.pth')))
|
||||
torch.save(model.embed_tokens, os.path.join(output_dir, 'embedding', 'emb.pth'))
|
||||
model = PeftModel.from_pretrained(model, model_args.from_peft, is_trainable=True)
|
||||
model.print_trainable_parameters()
|
||||
else:
|
||||
if model_args.use_lora:
|
||||
peft_config = LoraConfig(
|
||||
task_type=TaskType.FEATURE_EXTRACTION,
|
||||
inference_mode=False,
|
||||
r=model_args.lora_rank,
|
||||
target_modules=target_modules,
|
||||
modules_to_save=model_args.modules_to_save,
|
||||
lora_alpha=model_args.lora_alpha,
|
||||
lora_dropout=model_args.lora_dropout
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def save_merged_model(model_args: DecoderOnlyEmbedderModelArguments, output_dir: str):
|
||||
"""
|
||||
Loads a model with specified configurations, merges it with PEFT layers if available.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
|
||||
output_dir (str): Directory to save the model.
|
||||
"""
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = AutoModel.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.raw_peft is not None:
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.raw_peft, 'embedding', 'emb.pth')))
|
||||
model = PeftModel.from_pretrained(model, model_args.raw_peft)
|
||||
model = model.merge_and_unload()
|
||||
|
||||
if os.path.exists(os.path.join(output_dir, 'embedding', 'emb.pth')):
|
||||
model.set_input_embeddings(torch.load(os.path.join(output_dir, 'embedding', 'emb.pth')))
|
||||
# modify the vocab size in the model configuration
|
||||
model.config.vocab_size = len(tokenizer)
|
||||
|
||||
try:
|
||||
model = PeftModel.from_pretrained(model, output_dir)
|
||||
model = model.merge_and_unload()
|
||||
except:
|
||||
model = PeftModel.from_pretrained(model, find_largest_checkpoint(output_dir))
|
||||
model = model.merge_and_unload()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=model_args.trust_remote_code)
|
||||
tokenizer.save_pretrained(os.path.join(output_dir, 'merged_model'))
|
||||
|
||||
model.save_pretrained(os.path.join(output_dir, 'merged_model'))
|
||||
@@ -0,0 +1,207 @@
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from transformers import AutoModel, PreTrainedModel, PreTrainedTokenizer
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiDecoderOnlyEmbedderModel(AbsEmbedderModel):
|
||||
"""Embedder model class for decoder only model.
|
||||
|
||||
Args:
|
||||
base_model (PreTrainedModel): The base model to train on.
|
||||
tokenizer (PreTrainedTokenizer, optional): The tokenizer to use. Defaults to ``None``.
|
||||
negatives_cross_device (bool, optional): If True, will compute cross devices negative loss. Defaults to ``False``.
|
||||
temperature (float, optional): Temperature to control the scale of scores. Defaults to ``1.0``.
|
||||
sub_batch_size (int, optional): Sub-batch size during encoding. If negative, will not split to sub-batch.
|
||||
Defaults to ``-1``.
|
||||
kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``'kl_div'``.
|
||||
use_mrl (bool, optional): Whether to use MRL for training. Defaults to ``False``.
|
||||
mrl_dims (List[int], optional): The dimensions of MRL layers. Defaults to ``[]``.
|
||||
sentence_pooling_method (str, optional): Pooling method to get sentence embedding. Defaults to ``'last_token'``.
|
||||
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to ``False``.
|
||||
"""
|
||||
TRANSFORMER_CLS = AutoModel
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_model: PreTrainedModel,
|
||||
tokenizer: PreTrainedTokenizer = None,
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1.0,
|
||||
sub_batch_size: int = -1,
|
||||
kd_loss_type: str = 'kl_div',
|
||||
use_mrl: bool = False,
|
||||
mrl_dims: List[int] = [],
|
||||
sentence_pooling_method: str = 'last_token',
|
||||
normalize_embeddings: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
base_model,
|
||||
tokenizer=tokenizer,
|
||||
negatives_cross_device=negatives_cross_device,
|
||||
temperature=temperature,
|
||||
sub_batch_size=sub_batch_size,
|
||||
kd_loss_type=kd_loss_type,
|
||||
use_mrl=use_mrl,
|
||||
mrl_dims=mrl_dims,
|
||||
)
|
||||
self.sentence_pooling_method = sentence_pooling_method
|
||||
self.normalize_embeddings = normalize_embeddings
|
||||
self.cross_entropy = torch.nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
def encode(self, features):
|
||||
"""
|
||||
Encode and get the embedding.
|
||||
|
||||
Args:
|
||||
features (Union[list, dict]): Features feed to the model.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, List[torch.Tensor]]: The embedding vectors or a list of embedding vectors if MRL is used.
|
||||
"""
|
||||
if features is None:
|
||||
return None
|
||||
if not isinstance(features, list):
|
||||
if self.sub_batch_size is not None and self.sub_batch_size > 0:
|
||||
all_p_reps = []
|
||||
for i in range(0, len(features['attention_mask']), self.sub_batch_size):
|
||||
end_inx = min(i + self.sub_batch_size, len(features['attention_mask']))
|
||||
sub_features = {}
|
||||
for k, v in features.items():
|
||||
sub_features[k] = v[i:end_inx]
|
||||
last_hidden_state = self.model(**sub_features, return_dict=True).last_hidden_state
|
||||
p_reps = self._sentence_embedding(last_hidden_state, sub_features['attention_mask'])
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
else:
|
||||
last_hidden_state = self.model(**features, return_dict=True).last_hidden_state
|
||||
all_p_reps = self._sentence_embedding(last_hidden_state, features['attention_mask'])
|
||||
else:
|
||||
all_p_reps = []
|
||||
for sub_features in features:
|
||||
last_hidden_state = self.model(**sub_features, return_dict=True).last_hidden_state
|
||||
p_reps = self._sentence_embedding(last_hidden_state, sub_features['attention_mask'])
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
|
||||
if self.use_mrl:
|
||||
p_reps_list = []
|
||||
ori_dim = all_p_reps.size(-1)
|
||||
for dim in self.mrl_dims:
|
||||
if dim > ori_dim:
|
||||
logger.warning(f"MRL dim {dim} is larger than original dimension {ori_dim}, using original dimension instead.")
|
||||
dim = min(dim, ori_dim)
|
||||
dim_p_reps = all_p_reps[:, :dim]
|
||||
if self.normalize_embeddings:
|
||||
dim_p_reps = torch.nn.functional.normalize(dim_p_reps, dim=-1)
|
||||
p_reps_list.append(dim_p_reps.contiguous())
|
||||
return p_reps_list
|
||||
else:
|
||||
if self.normalize_embeddings:
|
||||
all_p_reps = torch.nn.functional.normalize(all_p_reps, dim=-1)
|
||||
return all_p_reps.contiguous()
|
||||
|
||||
def _sentence_embedding(self, last_hidden_state, attention_mask):
|
||||
"""Use the pooling method to get the sentence embedding.
|
||||
|
||||
Args:
|
||||
last_hidden_state (torch.Tensor): The model output's last hidden state.
|
||||
attention_mask (torch.Tensor): Mask out padding tokens during pooling.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Specified pooling method not implemented.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The sentence embeddings.
|
||||
"""
|
||||
if self.sentence_pooling_method == "cls":
|
||||
return last_hidden_state[:, 0]
|
||||
elif self.sentence_pooling_method == "mean":
|
||||
s = torch.sum(
|
||||
last_hidden_state * attention_mask.unsqueeze(-1).float(), dim=1
|
||||
)
|
||||
d = attention_mask.sum(dim=1, keepdim=True).float()
|
||||
return s / d
|
||||
elif self.sentence_pooling_method == "last_token":
|
||||
left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
|
||||
if left_padding:
|
||||
return last_hidden_state[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = last_hidden_state.shape[0]
|
||||
return last_hidden_state[
|
||||
torch.arange(batch_size, device=last_hidden_state.device),
|
||||
sequence_lengths,
|
||||
]
|
||||
else:
|
||||
raise NotImplementedError(f"pooling method {self.sentence_pooling_method} not implemented")
|
||||
|
||||
def compute_score(self, q_reps, p_reps):
|
||||
"""Computes the scores between query and passage representations.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed scores, adjusted by temperature.
|
||||
"""
|
||||
scores = self._compute_similarity(q_reps, p_reps) / self.temperature
|
||||
scores = scores.view(q_reps.size(0), -1)
|
||||
return scores
|
||||
|
||||
def _compute_similarity(self, q_reps, p_reps):
|
||||
"""Computes the similarity between query and passage representations using inner product.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed similarity matrix.
|
||||
"""
|
||||
if len(p_reps.size()) == 2:
|
||||
return torch.matmul(q_reps, p_reps.transpose(0, 1))
|
||||
return torch.matmul(q_reps, p_reps.transpose(-2, -1))
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
"""Compute the loss using cross entropy.
|
||||
|
||||
Args:
|
||||
scores (torch.Tensor): Computed score.
|
||||
target (torch.Tensor): The target value.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed cross entropy loss.
|
||||
"""
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
"""
|
||||
Activates gradient checkpointing for the current model.
|
||||
"""
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def enable_input_require_grads(self, **kwargs):
|
||||
"""
|
||||
Enables the gradients for the input embeddings.
|
||||
"""
|
||||
self.model.enable_input_require_grads(**kwargs)
|
||||
|
||||
def save(self, output_dir: str):
|
||||
"""Save the model to the directory.
|
||||
|
||||
Args:
|
||||
output_dir (str): Directory for saving the model.
|
||||
"""
|
||||
state_dict = self.model.state_dict()
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
self.model.save_pretrained(output_dir, state_dict=state_dict)
|
||||
@@ -0,0 +1,134 @@
|
||||
import logging
|
||||
from typing import Tuple
|
||||
from pathlib import Path
|
||||
from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizer
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder.AbsArguments import AbsEmbedderDataArguments, AbsEmbedderTrainingArguments
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderRunner, AbsEmbedderModel, EmbedderTrainerCallbackForDataRefresh
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderModelArguments
|
||||
from .trainer import DecoderOnlyEmbedderTrainer
|
||||
from .modeling import BiDecoderOnlyEmbedderModel
|
||||
from .load_model import get_model, save_merged_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DecoderOnlyEmbedderRunner(AbsEmbedderRunner):
|
||||
"""Runner class for decoder only embedding model.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
|
||||
data_args (AbsEmbedderDataArguments): Data arguments instance.
|
||||
training_args (AbsEmbedderTrainingArguments): Trainer arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model_args: DecoderOnlyEmbedderModelArguments,
|
||||
data_args: AbsEmbedderDataArguments,
|
||||
training_args: AbsEmbedderTrainingArguments
|
||||
):
|
||||
super().__init__(model_args, data_args, training_args)
|
||||
self.model_args: DecoderOnlyEmbedderModelArguments
|
||||
self.data_args: AbsEmbedderDataArguments
|
||||
self.training_args: AbsEmbedderTrainingArguments
|
||||
|
||||
def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
|
||||
"""Load tokenizer and model.
|
||||
|
||||
Returns:
|
||||
Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Tokenizer and model instances.
|
||||
"""
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.model_args.tokenizer_name if self.model_args.tokenizer_name else self.model_args.model_name_or_path,
|
||||
token=self.model_args.token,
|
||||
cache_dir=self.model_args.cache_dir,
|
||||
use_fast=self.model_args.use_fast_tokenizer,
|
||||
add_eos_token=True,
|
||||
trust_remote_code=self.model_args.trust_remote_code,
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
if tokenizer.unk_token is not None:
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
tokenizer.pad_token_id = tokenizer.unk_token_id
|
||||
else:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = 'left'
|
||||
|
||||
resize = False
|
||||
if self.model_args.additional_special_tokens is not None:
|
||||
special_tokens_dict = {'additional_special_tokens': self.model_args.additional_special_tokens}
|
||||
add_num = tokenizer.add_special_tokens(special_tokens_dict)
|
||||
if add_num > 0:
|
||||
resize = True
|
||||
logger.info(f"Add {add_num} special tokens to the tokenizer. Special tokens: {self.model_args.additional_special_tokens}")
|
||||
else:
|
||||
logger.warning(f"Special tokens {self.model_args.additional_special_tokens} already exists in the tokenizer.")
|
||||
base_model = get_model(self.model_args, self.training_args.output_dir, resize, len(tokenizer))
|
||||
|
||||
num_labels = 1
|
||||
config = AutoConfig.from_pretrained(
|
||||
self.model_args.config_name if self.model_args.config_name else self.model_args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
cache_dir=self.model_args.cache_dir,
|
||||
token=self.model_args.token,
|
||||
trust_remote_code=self.model_args.trust_remote_code,
|
||||
)
|
||||
logger.info('Config: %s', config)
|
||||
|
||||
model = BiDecoderOnlyEmbedderModel(
|
||||
base_model,
|
||||
tokenizer=tokenizer,
|
||||
negatives_cross_device=self.training_args.negatives_cross_device,
|
||||
temperature=self.training_args.temperature,
|
||||
sub_batch_size=self.training_args.sub_batch_size,
|
||||
kd_loss_type=self.training_args.kd_loss_type,
|
||||
use_mrl=self.training_args.use_mrl,
|
||||
mrl_dims=self.training_args.mrl_dims,
|
||||
sentence_pooling_method=self.training_args.sentence_pooling_method,
|
||||
normalize_embeddings=self.training_args.normalize_embeddings
|
||||
)
|
||||
|
||||
if self.training_args.gradient_checkpointing:
|
||||
model.enable_input_require_grads()
|
||||
|
||||
if self.training_args.fix_position_embedding:
|
||||
for k, v in model.named_parameters():
|
||||
if "position_embeddings" in k:
|
||||
logging.info(f"Freeze the parameters for {k}")
|
||||
v.requires_grad = False
|
||||
return tokenizer, model
|
||||
|
||||
def load_trainer(self) -> DecoderOnlyEmbedderTrainer:
|
||||
"""Load the trainer.
|
||||
|
||||
Returns:
|
||||
DecoderOnlyEmbedderTrainer: Loaded trainer instance.
|
||||
"""
|
||||
trainer = DecoderOnlyEmbedderTrainer(
|
||||
model=self.model,
|
||||
args=self.training_args,
|
||||
train_dataset=self.train_dataset,
|
||||
data_collator=self.data_collator,
|
||||
processing_class=self.tokenizer
|
||||
)
|
||||
if self.data_args.same_dataset_within_batch:
|
||||
trainer.add_callback(EmbedderTrainerCallbackForDataRefresh(self.train_dataset))
|
||||
return trainer
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Run the finetune.
|
||||
"""
|
||||
if not self.model_args.only_merge_lora_model:
|
||||
Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
|
||||
self.trainer.save_model()
|
||||
|
||||
# save merged model
|
||||
if self.model_args.save_merged_lora_model and self.training_args.process_index == 0:
|
||||
save_merged_model(self.model_args, self.training_args.output_dir)
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DecoderOnlyEmbedderTrainer(AbsEmbedderTrainer):
|
||||
"""
|
||||
Trainer class for base encoder models.
|
||||
"""
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
"""Save the model to directory.
|
||||
|
||||
Args:
|
||||
output_dir (Optional[str], optional): Output directory to save the model. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError
|
||||
"""
|
||||
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
logger.info("Saving model checkpoint to %s", output_dir)
|
||||
# Save a trained model and configuration using `save_pretrained()`.
|
||||
# They can then be reloaded using `from_pretrained()`
|
||||
if not hasattr(self.model, 'save'):
|
||||
raise NotImplementedError(
|
||||
f'MODEL {self.model.__class__.__name__} '
|
||||
f'does not support save interface')
|
||||
else:
|
||||
self.model.save(output_dir)
|
||||
|
||||
if self.tokenizer is not None and self.is_world_process_zero():
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
|
||||
# save the checkpoint for sentence-transformers library
|
||||
# if self.is_world_process_zero():
|
||||
# save_ckpt_for_sentence_transformers(output_dir,
|
||||
# pooling_mode=self.args.sentence_pooling_method,
|
||||
# normlized=self.args.normlized)
|
||||
@@ -0,0 +1,24 @@
|
||||
from FlagEmbedding.abc.finetune.embedder import (
|
||||
AbsEmbedderTrainingArguments as DecoderOnlyEmbedderICLTrainingArguments,
|
||||
)
|
||||
|
||||
from .arguments import (
|
||||
DecoderOnlyEmbedderICLModelArguments,
|
||||
DecoderOnlyEmbedderICLDataArguments
|
||||
)
|
||||
from .dataset import (
|
||||
DecoderOnlyEmbedderICLSameDatasetTrainDataset,
|
||||
AbsEmbedderSameDatasetCollator
|
||||
)
|
||||
from .modeling import BiDecoderOnlyEmbedderICLModel
|
||||
from .trainer import DecoderOnlyEmbedderICLTrainer
|
||||
from .runner import DecoderOnlyEmbedderICLRunner
|
||||
|
||||
__all__ = [
|
||||
'DecoderOnlyEmbedderICLModelArguments',
|
||||
'DecoderOnlyEmbedderICLDataArguments',
|
||||
'DecoderOnlyEmbedderICLTrainingArguments',
|
||||
'BiDecoderOnlyEmbedderICLModel',
|
||||
'DecoderOnlyEmbedderICLTrainer',
|
||||
'DecoderOnlyEmbedderICLRunner',
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.finetune.embedder.decoder_only.icl import (
|
||||
DecoderOnlyEmbedderICLDataArguments,
|
||||
DecoderOnlyEmbedderICLTrainingArguments,
|
||||
DecoderOnlyEmbedderICLModelArguments,
|
||||
DecoderOnlyEmbedderICLRunner,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((
|
||||
DecoderOnlyEmbedderICLModelArguments,
|
||||
DecoderOnlyEmbedderICLDataArguments,
|
||||
DecoderOnlyEmbedderICLTrainingArguments
|
||||
))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: DecoderOnlyEmbedderICLModelArguments
|
||||
data_args: DecoderOnlyEmbedderICLDataArguments
|
||||
training_args: DecoderOnlyEmbedderICLTrainingArguments
|
||||
|
||||
runner = DecoderOnlyEmbedderICLRunner(
|
||||
model_args=model_args,
|
||||
data_args=data_args,
|
||||
training_args=training_args
|
||||
)
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,102 @@
|
||||
from typing import Optional, List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import (
|
||||
AbsEmbedderModelArguments,
|
||||
AbsEmbedderDataArguments,
|
||||
)
|
||||
|
||||
|
||||
def default_target_modules() -> List[int]:
|
||||
return ['v_proj', 'q_proj', 'k_proj', 'gate_proj', 'down_proj', 'o_proj', 'up_proj']
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecoderOnlyEmbedderICLModelArguments(AbsEmbedderModelArguments):
|
||||
"""
|
||||
Model argument class for decoder only icl model.
|
||||
"""
|
||||
peft_model_path: str = field(
|
||||
default='', metadata={"help": "The peft model checkpoint for initialization."}
|
||||
)
|
||||
use_lora: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use LORA (low-rank parameter-efficient training) to train the model."}
|
||||
)
|
||||
lora_rank: int = field(
|
||||
default=64,
|
||||
metadata={"help": "The rank of lora."}
|
||||
)
|
||||
lora_alpha: float = field(
|
||||
default=16,
|
||||
metadata={"help": "The alpha parameter of lora."}
|
||||
)
|
||||
lora_dropout: float = field(
|
||||
default=0.1,
|
||||
metadata={"help": "The dropout rate of lora modules."}
|
||||
)
|
||||
target_modules: List[str] = field(
|
||||
default_factory=default_target_modules,
|
||||
metadata={"help": "The target modules to apply LORA."}
|
||||
)
|
||||
use_flash_attn: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use flash attention to train the model."}
|
||||
)
|
||||
use_slow_tokenizer: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library)."}
|
||||
)
|
||||
# low_cpu_mem_usage: bool = field(
|
||||
# default=False,
|
||||
# metadata={"help": "It is an option to create the model as an empty shell,"
|
||||
# "then only materialize its parameters when the pretrained weights are loaded."
|
||||
# "If passed, LLM loading time and RAM consumption will be benefited."}
|
||||
# )
|
||||
from_peft: str = field(
|
||||
default=None
|
||||
)
|
||||
modules_to_save: List[str] = field(
|
||||
default=None,
|
||||
)
|
||||
raw_peft: str = field(
|
||||
default=None
|
||||
)
|
||||
|
||||
additional_special_tokens: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={"help": "additional special tokens", "nargs": "+"}
|
||||
)
|
||||
|
||||
save_merged_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will merge the lora modules and save the entire model."}
|
||||
)
|
||||
|
||||
only_merge_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will only merge the lora modules and save the entire model."}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DecoderOnlyEmbedderICLDataArguments(AbsEmbedderDataArguments):
|
||||
"""
|
||||
Data argument class for decoder only icl model.
|
||||
"""
|
||||
example_query_max_len: int = field(
|
||||
default=64,
|
||||
metadata={"help": "The max length of example query."}
|
||||
)
|
||||
example_passage_max_len: int = field(
|
||||
default=96,
|
||||
metadata={"help": "The max length of example passage."}
|
||||
)
|
||||
retrieval_use_examples: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use examples for retrieval."}
|
||||
)
|
||||
icl_suffix_str: str = field(
|
||||
default='\nResponse:',
|
||||
metadata={"help": "The suffix string for ICL dataset."}
|
||||
)
|
||||
@@ -0,0 +1,291 @@
|
||||
import math
|
||||
import random
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from transformers import (
|
||||
PreTrainedTokenizer,
|
||||
DataCollatorWithPadding,
|
||||
)
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderSameDatasetTrainDataset
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderICLDataArguments
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DecoderOnlyEmbedderICLSameDatasetTrainDataset(AbsEmbedderSameDatasetTrainDataset):
|
||||
"""Dataset class for icl model.
|
||||
|
||||
Args:
|
||||
args (DecoderOnlyEmbedderICLDataArguments): Data argument class for icl model.
|
||||
default_batch_size (int): The default batch size.
|
||||
seed (int): Random seed to use.
|
||||
tokenizer (PreTrainedTokenizer): Tokenzier.
|
||||
process_index (int, optional): Current process index. Defaults to 0.
|
||||
num_processes (int, optional): Total number of processes. Defaults to 1.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
args: DecoderOnlyEmbedderICLDataArguments,
|
||||
default_batch_size: int,
|
||||
seed: int,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
process_index: int=0,
|
||||
num_processes: int=1
|
||||
):
|
||||
super().__init__(
|
||||
args=args,
|
||||
default_batch_size=default_batch_size,
|
||||
seed=seed,
|
||||
tokenizer=tokenizer,
|
||||
process_index=process_index,
|
||||
num_processes=num_processes
|
||||
)
|
||||
self.args: DecoderOnlyEmbedderICLDataArguments
|
||||
|
||||
self.suffix = self.tokenizer(f"{self.args.icl_suffix_str}{self.tokenizer.eos_token}", add_special_tokens=False)['input_ids']
|
||||
|
||||
self.prefix = self.tokenizer(f"{self.tokenizer.bos_token}", add_special_tokens=False)['input_ids']
|
||||
|
||||
def _create_batch_data(self, batch_raw_data):
|
||||
"""Create a comple batch of data with queries, documents and teacher scores.
|
||||
|
||||
Args:
|
||||
batch_raw_data (datasets.Dataset): One batch of raw data.
|
||||
|
||||
Returns:
|
||||
List[str]: Queries with instruction format.
|
||||
List[str]: Documents with instruction format.
|
||||
List[float]: Teacher scores for model distillation.
|
||||
"""
|
||||
queries, passages, teacher_scores = [], [], []
|
||||
|
||||
train_group_size, data_type = self._get_train_group_size(batch_raw_data)
|
||||
|
||||
icl_pairs = []
|
||||
|
||||
for i in range(len(batch_raw_data['query'])):
|
||||
if data_type is not None:
|
||||
assert batch_raw_data['type'][i] == data_type, f"Data type is not consistent in the same batch"
|
||||
|
||||
queries.append(
|
||||
self.args.query_instruction_format.format(
|
||||
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
||||
batch_raw_data['query'][i]
|
||||
)
|
||||
)
|
||||
tmp_passages = []
|
||||
pos_idx = random.choice(list(range(len(batch_raw_data['pos'][i]))))
|
||||
pos = self._shuffle_text(batch_raw_data['pos'][i][pos_idx])
|
||||
tmp_passages.append(pos)
|
||||
|
||||
neg_all_idx = list(range(len(batch_raw_data['neg'][i])))
|
||||
if len(batch_raw_data['neg'][i]) < train_group_size - 1:
|
||||
num = math.ceil((train_group_size - 1) / len(batch_raw_data['neg'][i]))
|
||||
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
||||
else:
|
||||
neg_idxs = random.sample(neg_all_idx, train_group_size - 1)
|
||||
for neg_idx in neg_idxs:
|
||||
tmp_passages.append(batch_raw_data['neg'][i][neg_idx])
|
||||
|
||||
if self.args.knowledge_distillation:
|
||||
if 'pos_scores' in batch_raw_data and batch_raw_data['pos_scores'][i] is not None:
|
||||
teacher_scores.append(batch_raw_data['pos_scores'][i][pos_idx])
|
||||
for neg_idx in neg_idxs:
|
||||
if 'neg_scores' in batch_raw_data and batch_raw_data['neg_scores'][i] is not None:
|
||||
teacher_scores.append(batch_raw_data['neg_scores'][i][neg_idx])
|
||||
else:
|
||||
teacher_scores = None
|
||||
|
||||
if data_type is not None and data_type in ['symmetric_sts', 'symmetric_clustering']:
|
||||
tmp_passages = [
|
||||
self.args.query_instruction_format.format(
|
||||
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
||||
p
|
||||
) for p in tmp_passages
|
||||
]
|
||||
tmp_passages = self.tokenizer.batch_decode(
|
||||
self.tokenizer(
|
||||
tmp_passages,
|
||||
max_length=self.args.passage_max_len - 1 - len(self.suffix),
|
||||
truncation=True,
|
||||
add_special_tokens=False,
|
||||
)['input_ids']
|
||||
)
|
||||
for j in range(len(tmp_passages)):
|
||||
tmp_passages[j] += self.args.icl_suffix_str
|
||||
else:
|
||||
if self.args.passage_instruction_for_retrieval is not None:
|
||||
tmp_passages = [
|
||||
self.args.passage_instruction_format.format(
|
||||
self.args.passage_instruction_for_retrieval, p
|
||||
) for p in tmp_passages
|
||||
]
|
||||
|
||||
passages.extend(tmp_passages)
|
||||
|
||||
if teacher_scores is not None:
|
||||
if len(teacher_scores) > 0 and len(passages) > 0:
|
||||
assert len(teacher_scores) == len(passages)
|
||||
|
||||
# add icl pairs
|
||||
if self.args.retrieval_use_examples or (
|
||||
data_type in ['symmetric_sts', 'symmetric_clustering', 'symmetric_class']
|
||||
):
|
||||
if data_type == 'symmetric_clustering':
|
||||
icl_pairs.append((
|
||||
self.tokenizer.decode(
|
||||
self.tokenizer(
|
||||
queries[-1],
|
||||
add_special_tokens=False
|
||||
)['input_ids'][:self.args.example_query_max_len]
|
||||
),
|
||||
self.tokenizer.decode(
|
||||
self.tokenizer(
|
||||
batch_raw_data['category'][i], # use category as example
|
||||
add_special_tokens=False
|
||||
)['input_ids'][:self.args.example_passage_max_len]
|
||||
)
|
||||
))
|
||||
else:
|
||||
icl_pairs.append((
|
||||
self.tokenizer.decode(
|
||||
self.tokenizer(
|
||||
queries[-1],
|
||||
add_special_tokens=False
|
||||
)['input_ids'][:self.args.example_query_max_len]
|
||||
),
|
||||
self.tokenizer.decode(
|
||||
self.tokenizer(
|
||||
pos,
|
||||
add_special_tokens=False
|
||||
)['input_ids'][:self.args.example_passage_max_len]
|
||||
)
|
||||
))
|
||||
else:
|
||||
icl_pairs = []
|
||||
|
||||
# handle queries
|
||||
for i in range(len(queries)):
|
||||
choices = random.choice([0, 1, 2, 3, 4, 5])
|
||||
if choices > 0 and len(icl_pairs) > 0:
|
||||
prefix_ids = random.sample(list(range(len(icl_pairs))), min(choices + 1, len(icl_pairs)))
|
||||
if i in prefix_ids:
|
||||
prefix_ids.remove(i)
|
||||
prefix_ids = prefix_ids[:choices]
|
||||
|
||||
prefix = ''
|
||||
for idx in prefix_ids:
|
||||
tmp = prefix + self.args.icl_suffix_str.join(icl_pairs[idx]) + '\n\n'
|
||||
if len(self.tokenizer(tmp)['input_ids']) > self.args.query_max_len - 512:
|
||||
break
|
||||
prefix = tmp
|
||||
else:
|
||||
prefix = ''
|
||||
|
||||
queries[i] = prefix + queries[i]
|
||||
queries[i] = self.tokenizer.decode(
|
||||
self.tokenizer(
|
||||
queries[i],
|
||||
max_length=self.args.query_max_len - len(self.prefix) - len(self.suffix),
|
||||
truncation=True,
|
||||
add_special_tokens=False
|
||||
)['input_ids']
|
||||
) + self.args.icl_suffix_str
|
||||
|
||||
return queries, passages, teacher_scores
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderSameDatasetCollator(DataCollatorWithPadding):
|
||||
"""
|
||||
EmbedCollator for SameDataset.
|
||||
Note that after using this collator, the training_args should be set as:
|
||||
|
||||
``training_args.per_device_train_batch_size = 1``
|
||||
|
||||
``training_args.dataloader_num_workers = 0 # avoid multi-processing``
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
sub_batch_size: int = -1
|
||||
|
||||
def __call__(self, features):
|
||||
queries = features[0][0]
|
||||
passages = features[0][1]
|
||||
teacher_scores = features[0][2]
|
||||
no_in_batch_neg_flag = features[0][3]
|
||||
|
||||
queries_inputs = self.tokenizer(
|
||||
queries,
|
||||
truncation=True,
|
||||
max_length=self.query_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
passages_inputs = self.tokenizer(
|
||||
passages,
|
||||
truncation=True,
|
||||
max_length=self.passage_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
|
||||
if self.sub_batch_size is None or self.sub_batch_size <= 0:
|
||||
q_collated = self.tokenizer.pad(
|
||||
queries_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
)
|
||||
|
||||
d_collated = self.tokenizer.pad(
|
||||
passages_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
)
|
||||
else:
|
||||
batch_size = self.sub_batch_size
|
||||
|
||||
q_collated = []
|
||||
for i in range(0, len(queries_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(queries_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
for k, v in queries_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
q_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
))
|
||||
|
||||
d_collated = []
|
||||
for i in range(0, len(passages_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(passages_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
|
||||
for k, v in passages_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
d_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
))
|
||||
|
||||
if isinstance(teacher_scores, list) and len(teacher_scores) == 0:
|
||||
teacher_scores = None
|
||||
|
||||
return {
|
||||
"queries": q_collated,
|
||||
"passages": d_collated,
|
||||
"teacher_scores": teacher_scores,
|
||||
"no_in_batch_neg_flag": no_in_batch_neg_flag
|
||||
}
|
||||
@@ -0,0 +1,183 @@
|
||||
import os
|
||||
import re
|
||||
import torch
|
||||
import logging
|
||||
from transformers import AutoConfig, AutoModel, AutoTokenizer
|
||||
from peft import LoraConfig, TaskType, get_peft_model, PeftModel
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderICLModelArguments
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def find_largest_checkpoint(checkpoint_dir):
|
||||
"""Find the largest checkpoint from directory.
|
||||
|
||||
Args:
|
||||
checkpoint_dir (str): Directory to the checkpoint.
|
||||
|
||||
Returns:
|
||||
str: Directory to the checkpoint, None no matching found.
|
||||
"""
|
||||
checkpoint_pattern = re.compile(r'checkpoint-(\d+)')
|
||||
max_number = -1
|
||||
max_checkpoint_file = None
|
||||
for file in os.listdir(checkpoint_dir):
|
||||
match = checkpoint_pattern.search(file)
|
||||
if match:
|
||||
number = int(match.group(1))
|
||||
if number > max_number:
|
||||
max_number = number
|
||||
max_checkpoint_file = file
|
||||
if max_checkpoint_file:
|
||||
return os.path.join(checkpoint_dir, max_checkpoint_file)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def get_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str, resize: bool, resize_tokens: int):
|
||||
"""Get the model.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
|
||||
output_dir (str): Directory to save the model.
|
||||
resize (bool): Whether to resize the number of tokens.
|
||||
resize_tokens (int): The new token size.
|
||||
|
||||
Returns:
|
||||
transformers.PreTrainedModel or PeftModel: The loaded model.
|
||||
"""
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = AutoModel.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
logger.info("Training new model from scratch")
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.raw_peft is not None:
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.raw_peft, 'embedding', 'emb.pth')))
|
||||
model = PeftModel.from_pretrained(model, model_args.raw_peft)
|
||||
model = model.merge_and_unload()
|
||||
|
||||
if resize:
|
||||
model.resize_token_embeddings(resize_tokens)
|
||||
os.makedirs(os.path.join(output_dir, 'embedding'), exist_ok=True)
|
||||
torch.save(model.embed_tokens, os.path.join(output_dir, 'embedding', 'emb.pth'))
|
||||
target_modules = model_args.target_modules
|
||||
else:
|
||||
target_modules = model_args.target_modules
|
||||
if 'embed_tokens' in target_modules:
|
||||
target_modules.remove('embed_tokens')
|
||||
|
||||
if model_args.from_peft is not None:
|
||||
if os.path.exists(os.path.join(model_args.from_peft, 'embedding')):
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.from_peft, 'embedding', 'emb.pth')))
|
||||
torch.save(model.embed_tokens, os.path.join(output_dir, 'embedding', 'emb.pth'))
|
||||
model = PeftModel.from_pretrained(model, model_args.from_peft, is_trainable=True)
|
||||
model.print_trainable_parameters()
|
||||
else:
|
||||
if model_args.use_lora:
|
||||
peft_config = LoraConfig(
|
||||
task_type=TaskType.FEATURE_EXTRACTION,
|
||||
inference_mode=False,
|
||||
r=model_args.lora_rank,
|
||||
target_modules=target_modules,
|
||||
modules_to_save=model_args.modules_to_save,
|
||||
lora_alpha=model_args.lora_alpha,
|
||||
lora_dropout=model_args.lora_dropout
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def save_merged_model(model_args: DecoderOnlyEmbedderICLModelArguments, output_dir: str):
|
||||
"""
|
||||
Loads a model with specified configurations, merges it with PEFT layers if available.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderModelArguments): Model arguments instance.
|
||||
output_dir (str): Directory to save the model.
|
||||
"""
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = AutoModel.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
attn_implementation = "flash_attention_2" if model_args.use_flash_attn else None,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
)
|
||||
else:
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.raw_peft is not None:
|
||||
model.set_input_embeddings(torch.load(os.path.join(model_args.raw_peft, 'embedding', 'emb.pth')))
|
||||
model = PeftModel.from_pretrained(model, model_args.raw_peft)
|
||||
model = model.merge_and_unload()
|
||||
|
||||
if os.path.exists(os.path.join(output_dir, 'embedding', 'emb.pth')):
|
||||
model.set_input_embeddings(torch.load(os.path.join(output_dir, 'embedding', 'emb.pth')))
|
||||
|
||||
try:
|
||||
model = PeftModel.from_pretrained(model, output_dir)
|
||||
model = model.merge_and_unload()
|
||||
except:
|
||||
model = PeftModel.from_pretrained(model, find_largest_checkpoint(output_dir))
|
||||
model = model.merge_and_unload()
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(output_dir, trust_remote_code=model_args.trust_remote_code)
|
||||
tokenizer.save_pretrained(os.path.join(output_dir, 'merged_model'))
|
||||
|
||||
# modify the vocab size in the model configuration
|
||||
model.config.vocab_size = len(tokenizer)
|
||||
model.save_pretrained(os.path.join(output_dir, 'merged_model'))
|
||||
@@ -0,0 +1,207 @@
|
||||
import logging
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
from transformers import AutoModel, PreTrainedModel, PreTrainedTokenizer
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BiDecoderOnlyEmbedderICLModel(AbsEmbedderModel):
|
||||
"""Embedder model class for decoder only model.
|
||||
|
||||
Args:
|
||||
base_model (PreTrainedModel): The base model to train on.
|
||||
tokenizer (PreTrainedTokenizer, optional): The tokenizer to use. Defaults to ``None``.
|
||||
negatives_cross_device (bool, optional): If True, will compute cross devices negative loss. Defaults to ``False``.
|
||||
temperature (float, optional): Temperature to control the scale of scores. Defaults to ``1.0``.
|
||||
sub_batch_size (int, optional): Sub-batch size during encoding. If negative, will not split to sub-batch.
|
||||
Defaults to ``-1``.
|
||||
kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``'kl_div'``.
|
||||
use_mrl (bool, optional): Whether to use MRL for training. Defaults to ``False``.
|
||||
mrl_dims (List[int], optional): The dimensions of MRL layers. Defaults to ``[]``.
|
||||
sentence_pooling_method (str, optional): Pooling method to get sentence embedding. Defaults to ``'last_token'``.
|
||||
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to ``False``.
|
||||
"""
|
||||
TRANSFORMER_CLS = AutoModel
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_model: PreTrainedModel,
|
||||
tokenizer: PreTrainedTokenizer = None,
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1.0,
|
||||
sub_batch_size: int = -1,
|
||||
kd_loss_type: str = 'kl_div',
|
||||
use_mrl: bool = False,
|
||||
mrl_dims: List[int] = [],
|
||||
sentence_pooling_method: str = 'last_token',
|
||||
normalize_embeddings: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
base_model,
|
||||
tokenizer=tokenizer,
|
||||
negatives_cross_device=negatives_cross_device,
|
||||
temperature=temperature,
|
||||
sub_batch_size=sub_batch_size,
|
||||
kd_loss_type=kd_loss_type,
|
||||
use_mrl=use_mrl,
|
||||
mrl_dims=mrl_dims,
|
||||
)
|
||||
self.sentence_pooling_method = sentence_pooling_method
|
||||
self.normalize_embeddings = normalize_embeddings
|
||||
self.cross_entropy = torch.nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
def encode(self, features):
|
||||
"""
|
||||
Encode and get the embedding.
|
||||
|
||||
Args:
|
||||
features (Union[list, dict]): Features feed to the model.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, List[torch.Tensor]]: The embedding vectors or a list of embedding vectors if MRL is used.
|
||||
"""
|
||||
if features is None:
|
||||
return None
|
||||
if not isinstance(features, list):
|
||||
if self.sub_batch_size is not None and self.sub_batch_size > 0:
|
||||
all_p_reps = []
|
||||
for i in range(0, len(features['attention_mask']), self.sub_batch_size):
|
||||
end_inx = min(i + self.sub_batch_size, len(features['attention_mask']))
|
||||
sub_features = {}
|
||||
for k, v in features.items():
|
||||
sub_features[k] = v[i:end_inx]
|
||||
last_hidden_state = self.model(**sub_features, return_dict=True).last_hidden_state
|
||||
p_reps = self._sentence_embedding(last_hidden_state, sub_features['attention_mask'])
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
else:
|
||||
last_hidden_state = self.model(**features, return_dict=True).last_hidden_state
|
||||
all_p_reps = self._sentence_embedding(last_hidden_state, features['attention_mask'])
|
||||
else:
|
||||
all_p_reps = []
|
||||
for sub_features in features:
|
||||
last_hidden_state = self.model(**sub_features, return_dict=True).last_hidden_state
|
||||
p_reps = self._sentence_embedding(last_hidden_state, sub_features['attention_mask'])
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
|
||||
if self.use_mrl:
|
||||
p_reps_list = []
|
||||
ori_dim = all_p_reps.size(-1)
|
||||
for dim in self.mrl_dims:
|
||||
if dim > ori_dim:
|
||||
logger.warning(f"MRL dim {dim} is larger than original dimension {ori_dim}, using original dimension instead.")
|
||||
dim = min(dim, ori_dim)
|
||||
dim_p_reps = all_p_reps[:, :dim]
|
||||
if self.normalize_embeddings:
|
||||
dim_p_reps = torch.nn.functional.normalize(dim_p_reps, dim=-1)
|
||||
p_reps_list.append(dim_p_reps.contiguous())
|
||||
return p_reps_list
|
||||
else:
|
||||
if self.normalize_embeddings:
|
||||
all_p_reps = torch.nn.functional.normalize(all_p_reps, dim=-1)
|
||||
return all_p_reps.contiguous()
|
||||
|
||||
def _sentence_embedding(self, last_hidden_state, attention_mask):
|
||||
"""Use the pooling method to get the sentence embedding.
|
||||
|
||||
Args:
|
||||
last_hidden_state (torch.Tensor): The model output's last hidden state.
|
||||
attention_mask (torch.Tensor): Mask out padding tokens during pooling.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Specified pooling method not implemented.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The sentence embeddings.
|
||||
"""
|
||||
if self.sentence_pooling_method == "cls":
|
||||
return last_hidden_state[:, 0]
|
||||
elif self.sentence_pooling_method == "mean":
|
||||
s = torch.sum(
|
||||
last_hidden_state * attention_mask.unsqueeze(-1).float(), dim=1
|
||||
)
|
||||
d = attention_mask.sum(dim=1, keepdim=True).float()
|
||||
return s / d
|
||||
elif self.sentence_pooling_method == "last_token":
|
||||
left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
|
||||
if left_padding:
|
||||
return last_hidden_state[:, -1]
|
||||
else:
|
||||
sequence_lengths = attention_mask.sum(dim=1) - 1
|
||||
batch_size = last_hidden_state.shape[0]
|
||||
return last_hidden_state[
|
||||
torch.arange(batch_size, device=last_hidden_state.device),
|
||||
sequence_lengths,
|
||||
]
|
||||
else:
|
||||
raise NotImplementedError(f"pooling method {self.sentence_pooling_method} not implemented")
|
||||
|
||||
def compute_score(self, q_reps, p_reps):
|
||||
"""Computes the scores between query and passage representations.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed scores, adjusted by temperature.
|
||||
"""
|
||||
scores = self._compute_similarity(q_reps, p_reps) / self.temperature
|
||||
scores = scores.view(q_reps.size(0), -1)
|
||||
return scores
|
||||
|
||||
def _compute_similarity(self, q_reps, p_reps):
|
||||
"""Computes the similarity between query and passage representations using inner product.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed similarity matrix.
|
||||
"""
|
||||
if len(p_reps.size()) == 2:
|
||||
return torch.matmul(q_reps, p_reps.transpose(0, 1))
|
||||
return torch.matmul(q_reps, p_reps.transpose(-2, -1))
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
"""Compute the loss using cross entropy.
|
||||
|
||||
Args:
|
||||
scores (torch.Tensor): Computed score.
|
||||
target (torch.Tensor): The target value.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed cross entropy loss.
|
||||
"""
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
"""
|
||||
Activates gradient checkpointing for the current model.
|
||||
"""
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def enable_input_require_grads(self, **kwargs):
|
||||
"""
|
||||
Enables the gradients for the input embeddings.
|
||||
"""
|
||||
self.model.enable_input_require_grads(**kwargs)
|
||||
|
||||
def save(self, output_dir: str):
|
||||
"""Save the model to the directory.
|
||||
|
||||
Args:
|
||||
output_dir (str): Directory for saving the model.
|
||||
"""
|
||||
state_dict = self.model.state_dict()
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
self.model.save_pretrained(output_dir, state_dict=state_dict)
|
||||
@@ -0,0 +1,159 @@
|
||||
import logging
|
||||
from typing import Tuple
|
||||
from pathlib import Path
|
||||
from transformers import AutoConfig, AutoTokenizer, PreTrainedTokenizer
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder.AbsArguments import AbsEmbedderTrainingArguments
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderRunner, AbsEmbedderModel, EmbedderTrainerCallbackForDataRefresh
|
||||
|
||||
from .arguments import DecoderOnlyEmbedderICLModelArguments, DecoderOnlyEmbedderICLDataArguments
|
||||
from .trainer import DecoderOnlyEmbedderICLTrainer
|
||||
from .modeling import BiDecoderOnlyEmbedderICLModel
|
||||
from .dataset import DecoderOnlyEmbedderICLSameDatasetTrainDataset
|
||||
from .load_model import get_model, save_merged_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DecoderOnlyEmbedderICLRunner(AbsEmbedderRunner):
|
||||
"""Runner class for decoder only icl model.
|
||||
|
||||
Args:
|
||||
model_args (DecoderOnlyEmbedderICLModelArguments): Model arguments instance.
|
||||
data_args (DecoderOnlyEmbedderICLDataArguments): Data arguments instance.
|
||||
training_args (AbsEmbedderTrainingArguments): Trainer arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model_args: DecoderOnlyEmbedderICLModelArguments,
|
||||
data_args: DecoderOnlyEmbedderICLDataArguments,
|
||||
training_args: AbsEmbedderTrainingArguments
|
||||
):
|
||||
super().__init__(model_args, data_args, training_args)
|
||||
self.model_args: DecoderOnlyEmbedderICLModelArguments
|
||||
self.data_args: DecoderOnlyEmbedderICLDataArguments
|
||||
self.training_args: AbsEmbedderTrainingArguments
|
||||
|
||||
def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
|
||||
"""Load tokenizer and model.
|
||||
|
||||
Returns:
|
||||
Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Tokenizer and model instances.
|
||||
"""
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
self.model_args.tokenizer_name if self.model_args.tokenizer_name else self.model_args.model_name_or_path,
|
||||
token=self.model_args.token,
|
||||
cache_dir=self.model_args.cache_dir,
|
||||
use_fast=self.model_args.use_fast_tokenizer,
|
||||
add_eos_token=True,
|
||||
trust_remote_code=self.model_args.trust_remote_code,
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
if tokenizer.unk_token is not None:
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
tokenizer.pad_token_id = tokenizer.unk_token_id
|
||||
else:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
tokenizer.pad_token_id = tokenizer.eos_token_id
|
||||
tokenizer.padding_side = 'left'
|
||||
|
||||
resize = False
|
||||
if self.model_args.additional_special_tokens is not None:
|
||||
special_tokens_dict = {'additional_special_tokens': self.model_args.additional_special_tokens}
|
||||
add_num = tokenizer.add_special_tokens(special_tokens_dict)
|
||||
if add_num > 0:
|
||||
resize = True
|
||||
logger.info(f"Add {add_num} special tokens to the tokenizer. Special tokens: {self.model_args.additional_special_tokens}")
|
||||
else:
|
||||
logger.warning(f"Special tokens {self.model_args.additional_special_tokens} already exists in the tokenizer.")
|
||||
base_model = get_model(self.model_args, self.training_args.output_dir, resize, len(tokenizer))
|
||||
|
||||
num_labels = 1
|
||||
config = AutoConfig.from_pretrained(
|
||||
self.model_args.config_name if self.model_args.config_name else self.model_args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
cache_dir=self.model_args.cache_dir,
|
||||
token=self.model_args.token,
|
||||
trust_remote_code=self.model_args.trust_remote_code,
|
||||
)
|
||||
logger.info('Config: %s', config)
|
||||
|
||||
model = BiDecoderOnlyEmbedderICLModel(
|
||||
base_model,
|
||||
tokenizer=tokenizer,
|
||||
negatives_cross_device=self.training_args.negatives_cross_device,
|
||||
temperature=self.training_args.temperature,
|
||||
sub_batch_size=self.training_args.sub_batch_size,
|
||||
kd_loss_type=self.training_args.kd_loss_type,
|
||||
use_mrl=self.training_args.use_mrl,
|
||||
mrl_dims=self.training_args.mrl_dims,
|
||||
sentence_pooling_method=self.training_args.sentence_pooling_method,
|
||||
normalize_embeddings=self.training_args.normalize_embeddings
|
||||
)
|
||||
|
||||
if self.training_args.gradient_checkpointing:
|
||||
model.enable_input_require_grads()
|
||||
|
||||
if self.training_args.fix_position_embedding:
|
||||
for k, v in model.named_parameters():
|
||||
if "position_embeddings" in k:
|
||||
logging.info(f"Freeze the parameters for {k}")
|
||||
v.requires_grad = False
|
||||
return tokenizer, model
|
||||
|
||||
def load_trainer(self) -> DecoderOnlyEmbedderICLTrainer:
|
||||
"""Load the trainer.
|
||||
|
||||
Returns:
|
||||
DecoderOnlyEmbedderICLTrainer: Loaded trainer instance.
|
||||
"""
|
||||
trainer = DecoderOnlyEmbedderICLTrainer(
|
||||
model=self.model,
|
||||
args=self.training_args,
|
||||
train_dataset=self.train_dataset,
|
||||
data_collator=self.data_collator,
|
||||
processing_class=self.tokenizer
|
||||
)
|
||||
if self.data_args.same_dataset_within_batch:
|
||||
trainer.add_callback(EmbedderTrainerCallbackForDataRefresh(self.train_dataset))
|
||||
return trainer
|
||||
|
||||
def load_train_dataset(self) -> DecoderOnlyEmbedderICLSameDatasetTrainDataset:
|
||||
"""Load the dataset instance for training.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Only support `same_dataset_within_batch` for `DecoderOnlyEmbedderICLRunner`.
|
||||
|
||||
Returns:
|
||||
DecoderOnlyEmbedderICLSameDatasetTrainDataset: The dataset instance.
|
||||
"""
|
||||
if self.data_args.same_dataset_within_batch:
|
||||
train_dataset = DecoderOnlyEmbedderICLSameDatasetTrainDataset(
|
||||
args=self.data_args,
|
||||
default_batch_size=self.training_args.per_device_train_batch_size,
|
||||
seed=self.training_args.seed,
|
||||
tokenizer=self.tokenizer,
|
||||
process_index=self.training_args.process_index,
|
||||
num_processes=self.training_args.world_size
|
||||
)
|
||||
self.training_args.per_device_train_batch_size = 1
|
||||
self.training_args.dataloader_num_workers = 0 # avoid multi-processing
|
||||
else:
|
||||
raise NotImplementedError("Only support `same_dataset_within_batch` for `DecoderOnlyEmbedderICLRunner`.")
|
||||
return train_dataset
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Run the finetune.
|
||||
"""
|
||||
if not self.model_args.only_merge_lora_model:
|
||||
Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
|
||||
self.trainer.save_model()
|
||||
|
||||
# save merged model
|
||||
if self.model_args.save_merged_lora_model and self.training_args.process_index == 0:
|
||||
save_merged_model(self.model_args, self.training_args.output_dir)
|
||||
@@ -0,0 +1,45 @@
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DecoderOnlyEmbedderICLTrainer(AbsEmbedderTrainer):
|
||||
"""
|
||||
Trainer class for base encoder models.
|
||||
"""
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
"""Save the model to directory.
|
||||
|
||||
Args:
|
||||
output_dir (Optional[str], optional): Output directory to save the model. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError
|
||||
"""
|
||||
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
logger.info("Saving model checkpoint to %s", output_dir)
|
||||
# Save a trained model and configuration using `save_pretrained()`.
|
||||
# They can then be reloaded using `from_pretrained()`
|
||||
if not hasattr(self.model, 'save'):
|
||||
raise NotImplementedError(
|
||||
f'MODEL {self.model.__class__.__name__} '
|
||||
f'does not support save interface')
|
||||
else:
|
||||
self.model.save(output_dir)
|
||||
|
||||
if self.tokenizer is not None and self.is_world_process_zero():
|
||||
self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
|
||||
# save the checkpoint for sentence-transformers library
|
||||
# if self.is_world_process_zero():
|
||||
# save_ckpt_for_sentence_transformers(output_dir,
|
||||
# pooling_mode=self.args.sentence_pooling_method,
|
||||
# normlized=self.args.normlized)
|
||||
Reference in New Issue
Block a user