wehub-resource-sync
179 lines
7.4 KiB
Python
179 lines
7.4 KiB
Python
import os
|
|
import torch
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
from transformers import (
|
|
AutoModel, AutoConfig,
|
|
AutoTokenizer, PreTrainedTokenizer
|
|
)
|
|
from huggingface_hub import snapshot_download
|
|
|
|
from FlagEmbedding.abc.finetune.embedder import (
|
|
AbsEmbedderRunner, AbsEmbedderModel,
|
|
AbsEmbedderDataArguments, EmbedderTrainerCallbackForDataRefresh
|
|
)
|
|
from .modeling import EncoderOnlyEmbedderM3Model
|
|
from .trainer import EncoderOnlyEmbedderM3Trainer
|
|
from .arguments import EncoderOnlyEmbedderM3ModelArguments, EncoderOnlyEmbedderM3TrainingArguments
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class EncoderOnlyEmbedderM3Runner(AbsEmbedderRunner):
|
|
"""
|
|
M3 model runner for finetuning.
|
|
|
|
Args:
|
|
model_args (EncoderOnlyEmbedderM3ModelArguments): Model arguments
|
|
data_args (AbsEmbedderDataArguments): Data arguments.
|
|
training_args (EncoderOnlyEmbedderM3TrainingArguments): Training arguments.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
model_args: EncoderOnlyEmbedderM3ModelArguments,
|
|
data_args: AbsEmbedderDataArguments,
|
|
training_args: EncoderOnlyEmbedderM3TrainingArguments
|
|
):
|
|
super().__init__(model_args, data_args, training_args)
|
|
self.model_args: EncoderOnlyEmbedderM3ModelArguments
|
|
self.data_args: AbsEmbedderDataArguments
|
|
self.training_args: EncoderOnlyEmbedderM3TrainingArguments
|
|
|
|
@staticmethod
|
|
def get_model(
|
|
model_name_or_path: str,
|
|
trust_remote_code: bool = False,
|
|
colbert_dim: int = -1,
|
|
cache_dir: str = None,
|
|
torch_dtype: Optional[torch.dtype] = None,
|
|
):
|
|
"""Get the model.
|
|
|
|
Args:
|
|
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
|
|
load a model from HuggingFace Hub with the name.
|
|
trust_remote_code (bool, optional): trust_remote_code to use when loading models from HF. Defaults to ``False``.
|
|
colbert_dim (int, optional): Colbert dim to set. Defaults to ``-1``.
|
|
cache_dir (str, optional): HF cache dir to store the model. Defaults to ``None``.
|
|
torch_dtype (Optional[torch.dtype], optional): Torch dtype used when loading model weights. Defaults to ``None``.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the model, colbert linear and sparse linear.
|
|
"""
|
|
cache_folder = os.getenv('HF_HUB_CACHE', None) if cache_dir is None else cache_dir
|
|
if not os.path.exists(model_name_or_path):
|
|
model_name_or_path = snapshot_download(
|
|
repo_id=model_name_or_path,
|
|
cache_dir=cache_folder,
|
|
ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5']
|
|
)
|
|
|
|
model = AutoModel.from_pretrained(
|
|
model_name_or_path,
|
|
cache_dir=cache_folder,
|
|
trust_remote_code=trust_remote_code,
|
|
dtype=torch_dtype,
|
|
)
|
|
colbert_linear = torch.nn.Linear(
|
|
in_features=model.config.hidden_size,
|
|
out_features=model.config.hidden_size if colbert_dim <= 0 else colbert_dim,
|
|
dtype=torch_dtype,
|
|
)
|
|
sparse_linear = torch.nn.Linear(
|
|
in_features=model.config.hidden_size,
|
|
out_features=1,
|
|
dtype=torch_dtype,
|
|
)
|
|
|
|
colbert_model_path = os.path.join(model_name_or_path, 'colbert_linear.pt')
|
|
sparse_model_path = os.path.join(model_name_or_path, 'sparse_linear.pt')
|
|
if os.path.exists(colbert_model_path) and os.path.exists(sparse_model_path):
|
|
logger.info('loading existing colbert_linear and sparse_linear---------')
|
|
colbert_state_dict = torch.load(colbert_model_path, map_location='cpu', weights_only=True)
|
|
sparse_state_dict = torch.load(sparse_model_path, map_location='cpu', weights_only=True)
|
|
colbert_linear.load_state_dict(colbert_state_dict)
|
|
sparse_linear.load_state_dict(sparse_state_dict)
|
|
else:
|
|
logger.info('The parameters of colbert_linear and sparse linear is new initialize. Make sure the model is loaded for training, not inferencing')
|
|
|
|
return {
|
|
'model': model,
|
|
'colbert_linear': colbert_linear,
|
|
'sparse_linear': sparse_linear
|
|
}
|
|
|
|
def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
|
|
"""Load the tokenizer and model.
|
|
|
|
Returns:
|
|
Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Tokenizer and model instances.
|
|
"""
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
self.model_args.model_name_or_path,
|
|
cache_dir=self.model_args.cache_dir,
|
|
token=self.model_args.token,
|
|
use_fast=self.model_args.use_fast_tokenizer,
|
|
trust_remote_code=self.model_args.trust_remote_code
|
|
)
|
|
|
|
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 = EncoderOnlyEmbedderM3Model(
|
|
self.get_model(self.model_args.model_name_or_path, self.model_args.trust_remote_code, self.model_args.colbert_dim),
|
|
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,
|
|
unified_finetuning=self.training_args.unified_finetuning,
|
|
use_self_distill=self.training_args.use_self_distill,
|
|
self_distill_start_step=self.training_args.self_distill_start_step
|
|
)
|
|
|
|
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
|
|
|
|
if self.training_args.fix_encoder:
|
|
for k, v in model.named_parameters():
|
|
if "colbert_linear" in k or 'sparse_linear' in k:
|
|
logging.info(f"train the parameters for {k}")
|
|
else:
|
|
v.requires_grad = False
|
|
|
|
return tokenizer, model
|
|
|
|
def load_trainer(self) -> EncoderOnlyEmbedderM3Trainer:
|
|
"""Load the M3 trainer.
|
|
|
|
Returns:
|
|
EncoderOnlyEmbedderM3Trainer: M3 Trainer instance.
|
|
"""
|
|
trainer = EncoderOnlyEmbedderM3Trainer(
|
|
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
|