chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,18 @@
|
||||
from FlagEmbedding.abc.finetune.embedder import (
|
||||
AbsEmbedderModelArguments as EncoderOnlyEmbedderModelArguments,
|
||||
AbsEmbedderDataArguments as EncoderOnlyEmbedderDataArguments,
|
||||
AbsEmbedderTrainingArguments as EncoderOnlyEmbedderTrainingArguments,
|
||||
)
|
||||
|
||||
from .modeling import BiEncoderOnlyEmbedderModel
|
||||
from .trainer import EncoderOnlyEmbedderTrainer
|
||||
from .runner import EncoderOnlyEmbedderRunner
|
||||
|
||||
__all__ = [
|
||||
'EncoderOnlyEmbedderModelArguments',
|
||||
'EncoderOnlyEmbedderDataArguments',
|
||||
'EncoderOnlyEmbedderTrainingArguments',
|
||||
'BiEncoderOnlyEmbedderModel',
|
||||
'EncoderOnlyEmbedderTrainer',
|
||||
'EncoderOnlyEmbedderRunner',
|
||||
]
|
||||
@@ -0,0 +1,31 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.finetune.embedder.encoder_only.base import (
|
||||
EncoderOnlyEmbedderDataArguments,
|
||||
EncoderOnlyEmbedderTrainingArguments,
|
||||
EncoderOnlyEmbedderModelArguments,
|
||||
EncoderOnlyEmbedderRunner,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((
|
||||
EncoderOnlyEmbedderModelArguments,
|
||||
EncoderOnlyEmbedderDataArguments,
|
||||
EncoderOnlyEmbedderTrainingArguments
|
||||
))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: EncoderOnlyEmbedderModelArguments
|
||||
data_args: EncoderOnlyEmbedderDataArguments
|
||||
training_args: EncoderOnlyEmbedderTrainingArguments
|
||||
|
||||
runner = EncoderOnlyEmbedderRunner(
|
||||
model_args=model_args,
|
||||
data_args=data_args,
|
||||
training_args=training_args
|
||||
)
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,206 @@
|
||||
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 BiEncoderOnlyEmbedderModel(AbsEmbedderModel):
|
||||
"""Embedder class for encoder 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 ``'cls'``.
|
||||
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 = 'cls',
|
||||
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,87 @@
|
||||
import logging
|
||||
from typing import Tuple
|
||||
from transformers import (
|
||||
AutoModel, AutoConfig,
|
||||
AutoTokenizer, PreTrainedTokenizer
|
||||
)
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderRunner, AbsEmbedderModel, EmbedderTrainerCallbackForDataRefresh
|
||||
from .modeling import BiEncoderOnlyEmbedderModel
|
||||
from .trainer import EncoderOnlyEmbedderTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EncoderOnlyEmbedderRunner(AbsEmbedderRunner):
|
||||
"""
|
||||
Finetune Runner for base embedding models.
|
||||
"""
|
||||
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.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
|
||||
)
|
||||
base_model = AutoModel.from_pretrained(
|
||||
self.model_args.model_name_or_path,
|
||||
cache_dir=self.model_args.cache_dir,
|
||||
token=self.model_args.token,
|
||||
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 = BiEncoderOnlyEmbedderModel(
|
||||
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) -> EncoderOnlyEmbedderTrainer:
|
||||
"""Load the trainer.
|
||||
|
||||
Returns:
|
||||
EncoderOnlyEmbedderTrainer: Loaded trainer instance.
|
||||
"""
|
||||
trainer = EncoderOnlyEmbedderTrainer(
|
||||
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
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EncoderOnlyEmbedderTrainer(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,17 @@
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderDataArguments as EncoderOnlyEmbedderM3DataArguments
|
||||
|
||||
from .arguments import EncoderOnlyEmbedderM3ModelArguments, EncoderOnlyEmbedderM3TrainingArguments
|
||||
from .modeling import EncoderOnlyEmbedderM3Model, EncoderOnlyEmbedderM3ModelForInference
|
||||
from .trainer import EncoderOnlyEmbedderM3Trainer
|
||||
from .runner import EncoderOnlyEmbedderM3Runner
|
||||
|
||||
|
||||
__all__ = [
|
||||
'EncoderOnlyEmbedderM3ModelArguments',
|
||||
'EncoderOnlyEmbedderM3DataArguments',
|
||||
'EncoderOnlyEmbedderM3TrainingArguments',
|
||||
'EncoderOnlyEmbedderM3Model',
|
||||
'EncoderOnlyEmbedderM3ModelForInference',
|
||||
'EncoderOnlyEmbedderM3Trainer',
|
||||
'EncoderOnlyEmbedderM3Runner',
|
||||
]
|
||||
@@ -0,0 +1,27 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.finetune.embedder.encoder_only.m3 import (
|
||||
EncoderOnlyEmbedderM3DataArguments,
|
||||
EncoderOnlyEmbedderM3TrainingArguments,
|
||||
EncoderOnlyEmbedderM3ModelArguments,
|
||||
EncoderOnlyEmbedderM3Runner,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((EncoderOnlyEmbedderM3ModelArguments, EncoderOnlyEmbedderM3DataArguments, EncoderOnlyEmbedderM3TrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: EncoderOnlyEmbedderM3ModelArguments
|
||||
data_args: EncoderOnlyEmbedderM3DataArguments
|
||||
training_args: EncoderOnlyEmbedderM3TrainingArguments
|
||||
|
||||
runner = EncoderOnlyEmbedderM3Runner(
|
||||
model_args=model_args,
|
||||
data_args=data_args,
|
||||
training_args=training_args
|
||||
)
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,25 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import (
|
||||
AbsEmbedderTrainingArguments,
|
||||
AbsEmbedderModelArguments
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderOnlyEmbedderM3ModelArguments(AbsEmbedderModelArguments):
|
||||
"""
|
||||
Model argument class for M3.
|
||||
"""
|
||||
colbert_dim: int = field(default=-1, metadata={"help": "Dim of colbert linear"})
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderOnlyEmbedderM3TrainingArguments(AbsEmbedderTrainingArguments):
|
||||
"""
|
||||
Training argument class for M3.
|
||||
"""
|
||||
unified_finetuning: bool = field(default=False, metadata={"help": "use unify fine-tuning"})
|
||||
use_self_distill: bool = field(default=False, metadata={"help": "use self-distill when using unify fine-tuning"})
|
||||
fix_encoder: bool = field(default=False, metadata={"help": "Freeze the parameters of encoder"})
|
||||
self_distill_start_step: int = field(default=-1, metadata={"help": "Num of step when using self-distill"})
|
||||
@@ -0,0 +1,583 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import Dict, List, Union, Any, Optional
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
import torch.nn.functional as F
|
||||
from transformers import PreTrainedTokenizer
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderModel, EmbedderOutput
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EncoderOnlyEmbedderM3Model(AbsEmbedderModel):
|
||||
"""Embedder class for M3 model.
|
||||
|
||||
Args:
|
||||
base_model (dict[str, Any]): 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 ``'m3_kd_loss'``.
|
||||
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 ``'cls'``.
|
||||
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to ``False``.
|
||||
unified_finetuning (bool, optional): If True, will finetune colbert vector and sparce embedding. Defaults to ``True``.
|
||||
use_self_distill (bool, optional): If True, will do self distillation. Defaults to ``False``.
|
||||
self_distill_start_step (int, optional): Step num to start self distillation. Defaults to ``-1``.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
base_model: Dict[str, Any],
|
||||
tokenizer: PreTrainedTokenizer = None,
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1,
|
||||
sub_batch_size: int = -1,
|
||||
kd_loss_type: str = 'm3_kd_loss',
|
||||
use_mrl: bool = False,
|
||||
mrl_dims: List[int] = [],
|
||||
sentence_pooling_method: str = 'cls',
|
||||
normalize_embeddings: bool = False,
|
||||
unified_finetuning: bool = True,
|
||||
use_self_distill: bool = False,
|
||||
self_distill_start_step: int = -1
|
||||
):
|
||||
if use_mrl is True:
|
||||
raise NotImplementedError(f"use_mrl method not implemented for M3 model")
|
||||
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,
|
||||
)
|
||||
self.sentence_pooling_method = sentence_pooling_method
|
||||
self.normalize_embeddings = normalize_embeddings
|
||||
self.cross_entropy = torch.nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
self.unified_finetuning = unified_finetuning
|
||||
if not self.unified_finetuning:
|
||||
self.model = base_model['model']
|
||||
self.colbert_linear = None
|
||||
self.sparse_linear = None
|
||||
else:
|
||||
self.model = base_model['model']
|
||||
self.colbert_linear = base_model['colbert_linear']
|
||||
self.sparse_linear = base_model['sparse_linear']
|
||||
|
||||
self.config = self.model.config
|
||||
|
||||
self.vocab_size = self.model.config.vocab_size
|
||||
self.use_self_distill = use_self_distill
|
||||
self.self_distill_start_step = self_distill_start_step
|
||||
self.step = 0
|
||||
|
||||
def _dense_embedding(self, last_hidden_state, attention_mask):
|
||||
"""Use the pooling method to get the dense 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 dense 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 _sparse_embedding(self, hidden_state, input_ids, return_embedding: bool = True):
|
||||
"""Compute and return the sparse embedding.
|
||||
|
||||
Args:
|
||||
hidden_state (torch.Tensor): The model output's last hidden state.
|
||||
input_ids (_type_): Ids from input features.
|
||||
return_embedding (bool, optional): If True, return the computed embedding, otherwise just return the token weights.
|
||||
Defaults to ``True``.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The sparse embedding or just the token weights.
|
||||
"""
|
||||
token_weights = torch.relu(self.sparse_linear(hidden_state))
|
||||
if not return_embedding: return token_weights
|
||||
|
||||
if self.training:
|
||||
sparse_embedding = torch.zeros(
|
||||
input_ids.size(0), input_ids.size(1), self.vocab_size,
|
||||
dtype=token_weights.dtype,
|
||||
device=token_weights.device
|
||||
)
|
||||
sparse_embedding = torch.scatter(sparse_embedding, dim=-1, index=input_ids.unsqueeze(-1), src=token_weights)
|
||||
sparse_embedding = torch.max(sparse_embedding, dim=1).values
|
||||
else:
|
||||
# Optimize suggestion from issue #1364: https://github.com/FlagOpen/FlagEmbedding/issues/1364
|
||||
# Disable when self.training = True, otherwise will cause:
|
||||
# RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation
|
||||
sparse_embedding = torch.zeros(
|
||||
input_ids.size(0), self.vocab_size,
|
||||
dtype=token_weights.dtype,
|
||||
device=token_weights.device
|
||||
)
|
||||
sparse_embedding = sparse_embedding.scatter_reduce(
|
||||
dim=-1, index=input_ids, src=token_weights.squeeze(-1), reduce="amax"
|
||||
)
|
||||
|
||||
unused_tokens = [
|
||||
self.tokenizer.cls_token_id, self.tokenizer.eos_token_id,
|
||||
self.tokenizer.pad_token_id, self.tokenizer.unk_token_id
|
||||
]
|
||||
sparse_embedding[:, unused_tokens] *= 0.
|
||||
return sparse_embedding
|
||||
|
||||
def _colbert_embedding(self, last_hidden_state, mask):
|
||||
"""Get the colbert vectors.
|
||||
|
||||
Args:
|
||||
last_hidden_state (torch.Tensor): The model output's last hidden state.
|
||||
attention_mask (torch.Tensor): Mask out padding tokens during pooling.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The colbert vectors.
|
||||
"""
|
||||
colbert_vecs = self.colbert_linear(last_hidden_state[:, 1:])
|
||||
colbert_vecs = colbert_vecs * mask[:, 1:][:, :, None].float()
|
||||
return colbert_vecs
|
||||
|
||||
def compute_score(
|
||||
self, q_reps, p_reps, q_mask: torch.Tensor,
|
||||
dense_weight: float = 1.0, sparse_weight: float = 0.3, colbert_weight: float = 1.0
|
||||
):
|
||||
"""_summary_
|
||||
|
||||
Args:
|
||||
q_reps (_type_): Query representations.
|
||||
p_reps (_type_): Passage representations.
|
||||
q_mask (torch.Tensor): _description_
|
||||
dense_weight (float, optional): _description_. Defaults to 1.0.
|
||||
sparse_weight (float, optional): _description_. Defaults to 0.3.
|
||||
colbert_weight (float, optional): _description_. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
dense_score = self.compute_dense_score(q_reps, p_reps)
|
||||
sparse_score = self.compute_sparse_score(q_reps, p_reps)
|
||||
colbert_score = self.compute_colbert_score(q_reps, p_reps, q_mask=q_mask)
|
||||
return dense_score * dense_weight + sparse_score * sparse_weight + colbert_score * colbert_weight
|
||||
|
||||
def compute_dense_score(self, q_reps, p_reps):
|
||||
"""Compute the dense score.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed dense 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_sparse_score(self, q_reps, p_reps):
|
||||
"""Compute the sparse score.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed sparse 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_colbert_score(self, q_reps, p_reps, q_mask: torch.Tensor=None):
|
||||
"""Compute the colbert score.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed colber scores, adjusted by temperature.
|
||||
"""
|
||||
token_scores = torch.einsum('qin,pjn->qipj', q_reps, p_reps)
|
||||
scores, _ = token_scores.max(-1)
|
||||
scores = scores.sum(1) / q_mask[:, 1:].sum(-1, keepdim=True)
|
||||
scores = scores / self.temperature
|
||||
return scores
|
||||
|
||||
def ensemble_score(self, q_reps, p_reps, dense_scores=None, sparse_scores=None, colbert_scores=None):
|
||||
"""Compute the ensemble score of the three methods.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Query representations.
|
||||
p_reps (torch.Tensor): Passage representations.
|
||||
dense_scores (torch.Tensor, optional): The dense scores. Defaults to ``None``.
|
||||
sparse_scores (torch.Tensor, optional): The sparse scores. Defaults to ``None``.
|
||||
colbert_scores (torch.Tensor, optional): The colbert scores. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
ValueError: dense_scores, sparse_scores, colbert_scores must be provided
|
||||
|
||||
Returns:
|
||||
_type_: The ensemble score of the three methods.
|
||||
"""
|
||||
if dense_scores is None or sparse_scores is None or colbert_scores is None:
|
||||
raise ValueError("dense_scores, sparse_scores, colbert_scores must be provided!")
|
||||
return dense_scores + 0.3 * sparse_scores + colbert_scores
|
||||
|
||||
def _encode(self, features):
|
||||
"""Helper function to encode using input features.
|
||||
|
||||
Args:
|
||||
features (Union[list, dict]): Features feed to the model.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Dense embedding.
|
||||
torch.Tensor: Sparce embedding.
|
||||
torch.Tensor: Colbert vector.
|
||||
"""
|
||||
dense_vecs, sparse_vecs, colbert_vecs = None, None, None
|
||||
last_hidden_state = self.model(**features, return_dict=True).last_hidden_state
|
||||
dense_vecs = self._dense_embedding(last_hidden_state, features['attention_mask'])
|
||||
if self.unified_finetuning:
|
||||
sparse_vecs = self._sparse_embedding(last_hidden_state, features['input_ids'])
|
||||
colbert_vecs = self._colbert_embedding(last_hidden_state, features['attention_mask'])
|
||||
if self.normalize_embeddings:
|
||||
dense_vecs = F.normalize(dense_vecs, dim=-1)
|
||||
if self.unified_finetuning:
|
||||
colbert_vecs = F.normalize(colbert_vecs, dim=-1)
|
||||
return dense_vecs, sparse_vecs, colbert_vecs
|
||||
|
||||
def encode(self, features):
|
||||
"""Encode and get the embedding.
|
||||
|
||||
Args:
|
||||
features (Union[list, dict]): Features feed to the model.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Dense embeddings.
|
||||
torch.Tensor: Sparce embeddings.
|
||||
torch.Tensor: Colbert vectors.
|
||||
"""
|
||||
if features is None:
|
||||
return None
|
||||
|
||||
if not isinstance(features, list):
|
||||
if self.sub_batch_size is not None and self.sub_batch_size != -1:
|
||||
all_dense_vecs, all_sparse_vecs, all_colbert_vecs = [], [], []
|
||||
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]
|
||||
|
||||
dense_vecs, sparse_vecs, colbert_vecs = self._encode(sub_features)
|
||||
all_dense_vecs.append(dense_vecs)
|
||||
all_sparse_vecs.append(sparse_vecs)
|
||||
all_colbert_vecs.append(colbert_vecs)
|
||||
|
||||
dense_vecs = torch.cat(all_dense_vecs, 0)
|
||||
if self.unified_finetuning:
|
||||
sparse_vecs = torch.cat(all_sparse_vecs, 0)
|
||||
colbert_vecs = torch.cat(all_colbert_vecs, 0)
|
||||
else:
|
||||
dense_vecs, sparse_vecs, colbert_vecs = self._encode(features)
|
||||
else:
|
||||
all_dense_vecs, all_sparse_vecs, all_colbert_vecs = [], [], []
|
||||
for sub_features in features:
|
||||
dense_vecs, sparse_vecs, colbert_vecs = self._encode(sub_features)
|
||||
all_dense_vecs.append(dense_vecs)
|
||||
all_sparse_vecs.append(sparse_vecs)
|
||||
all_colbert_vecs.append(colbert_vecs)
|
||||
|
||||
dense_vecs = torch.cat(all_dense_vecs, 0)
|
||||
if self.unified_finetuning:
|
||||
sparse_vecs = torch.cat(all_sparse_vecs, 0)
|
||||
colbert_vecs = torch.cat(all_colbert_vecs, 0)
|
||||
|
||||
if self.unified_finetuning:
|
||||
return dense_vecs.contiguous(), sparse_vecs.contiguous(), colbert_vecs.contiguous()
|
||||
else:
|
||||
return dense_vecs.contiguous(), None, None
|
||||
|
||||
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 _get_queries_attention_mask(self, queries: Union[Dict[str, Tensor], List[Dict[str, Tensor]]]):
|
||||
"""padding attention mask for colbert
|
||||
|
||||
Args:
|
||||
queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]]): Input queries.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The query attention mask.
|
||||
"""
|
||||
if not isinstance(queries, list):
|
||||
q_mask = queries['attention_mask']
|
||||
else:
|
||||
q_mask_list = [sub_features['attention_mask'] for sub_features in queries]
|
||||
_length = max([mask.shape[1] for mask in q_mask_list])
|
||||
if self.tokenizer.padding_side == 'right':
|
||||
q_mask = torch.cat([
|
||||
F.pad(mask, (0, _length - mask.shape[1]), value=0)
|
||||
for mask in q_mask_list
|
||||
], dim=0)
|
||||
else:
|
||||
q_mask = torch.cat([
|
||||
F.pad(mask, (_length - mask.shape[1], 0), value=0)
|
||||
for mask in q_mask_list
|
||||
], dim=0)
|
||||
return q_mask
|
||||
|
||||
def forward(
|
||||
self,
|
||||
queries: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None,
|
||||
passages: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None,
|
||||
teacher_scores: Union[None, List[float]] = None,
|
||||
no_in_batch_neg_flag: bool = False,
|
||||
):
|
||||
"""The computation performed at every call.
|
||||
|
||||
Args:
|
||||
queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input queries. Defaults to ``None``.
|
||||
passages (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input passages. Defaults to ``None``.
|
||||
teacher_scores (Union[None, List[float]], optional): Teacher scores for distillation. Defaults to ``None``.
|
||||
no_in_batch_neg_flag (bool, optional): If True, use no in-batch negatives and no cross-device negatives. Defaults to ``False``.
|
||||
|
||||
Returns:
|
||||
EmbedderOutput: Output of the forward call of model.
|
||||
"""
|
||||
q_dense_vecs, q_sparse_vecs, q_colbert_vecs = self.encode(queries) # (batch_size, dim)
|
||||
p_dense_vecs, p_sparse_vecs, p_colbert_vecs = self.encode(passages) # (batch_size * group_size, dim)
|
||||
|
||||
if self.training:
|
||||
if teacher_scores is not None:
|
||||
teacher_scores = torch.tensor(teacher_scores, device=q_dense_vecs.device)
|
||||
teacher_scores = teacher_scores.view(q_dense_vecs.size(0), -1).detach() # (batch_size, group_size)
|
||||
teacher_targets = F.softmax(teacher_scores, dim=-1) # (batch_size, group_size)
|
||||
else:
|
||||
teacher_targets = None
|
||||
|
||||
if no_in_batch_neg_flag:
|
||||
compute_loss_func = self._compute_no_in_batch_neg_loss
|
||||
else:
|
||||
if self.negatives_cross_device:
|
||||
compute_loss_func = self._compute_cross_device_neg_loss
|
||||
else:
|
||||
compute_loss_func = self._compute_in_batch_neg_loss
|
||||
|
||||
# dense loss
|
||||
dense_scores, loss = compute_loss_func(
|
||||
q_dense_vecs, p_dense_vecs, teacher_targets=teacher_targets,
|
||||
compute_score_func=self.compute_dense_score
|
||||
)
|
||||
|
||||
if self.unified_finetuning:
|
||||
# disable cross device negatives for unified finetuning
|
||||
if no_in_batch_neg_flag:
|
||||
compute_loss_func = self._compute_no_in_batch_neg_loss
|
||||
else:
|
||||
compute_loss_func = self._compute_in_batch_neg_loss
|
||||
|
||||
# sparse loss
|
||||
sparse_scores, sparse_loss = compute_loss_func(
|
||||
q_sparse_vecs, p_sparse_vecs, teacher_targets=teacher_targets,
|
||||
compute_score_func=self.compute_sparse_score
|
||||
)
|
||||
|
||||
# colbert loss
|
||||
colbert_scores, colbert_loss = compute_loss_func(
|
||||
q_colbert_vecs, p_colbert_vecs, teacher_targets=teacher_targets,
|
||||
compute_score_func=self.compute_colbert_score,
|
||||
q_mask=self._get_queries_attention_mask(queries)
|
||||
)
|
||||
|
||||
# get dense scores of current process
|
||||
if not no_in_batch_neg_flag and self.negatives_cross_device:
|
||||
dense_scores = dense_scores[
|
||||
q_dense_vecs.size(0)*self.process_rank : q_dense_vecs.size(0)*(self.process_rank+1),
|
||||
p_dense_vecs.size(0)*self.process_rank : p_dense_vecs.size(0)*(self.process_rank+1)
|
||||
] # (batch_size, batch_size * group_size)
|
||||
elif no_in_batch_neg_flag:
|
||||
# get local p_dense_vecs: fix a bug described in
|
||||
# https://github.com/FlagOpen/FlagEmbedding/issues/1410
|
||||
group_size = p_dense_vecs.size(0) // q_dense_vecs.size(0)
|
||||
indices = torch.arange(0, q_dense_vecs.size(0), device=q_dense_vecs.device) * group_size
|
||||
p_dense_vecs = p_dense_vecs[indices, :]
|
||||
|
||||
# ensemble loss
|
||||
ensemble_scores, ensemble_loss = compute_loss_func(
|
||||
q_dense_vecs, p_dense_vecs, teacher_targets=teacher_targets,
|
||||
compute_score_func=self.ensemble_score,
|
||||
dense_scores=dense_scores,
|
||||
sparse_scores=sparse_scores,
|
||||
colbert_scores=colbert_scores
|
||||
)
|
||||
|
||||
loss = (loss + ensemble_loss + 0.1 * sparse_loss + colbert_loss) / 4
|
||||
|
||||
if self.use_self_distill and self.step > self.self_distill_start_step:
|
||||
self_teacher_targets = torch.softmax(ensemble_scores.detach(), dim=-1)
|
||||
|
||||
dense_self_distill_loss = self.distill_loss("kl_div", self_teacher_targets, dense_scores)
|
||||
sparse_self_distill_loss = self.distill_loss("kl_div", self_teacher_targets, sparse_scores)
|
||||
colbert_self_distill_loss = self.distill_loss("kl_div", self_teacher_targets, colbert_scores)
|
||||
|
||||
loss += (dense_self_distill_loss + 0.1 * sparse_self_distill_loss + colbert_self_distill_loss) / 3
|
||||
loss = loss / 2
|
||||
self.step += 1
|
||||
else:
|
||||
loss = None
|
||||
|
||||
return EmbedderOutput(
|
||||
loss=loss,
|
||||
)
|
||||
|
||||
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.
|
||||
"""
|
||||
def _trans_state_dict(state_dict):
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
return state_dict
|
||||
|
||||
self.model.save_pretrained(output_dir, state_dict=_trans_state_dict(self.model.state_dict()))
|
||||
|
||||
if self.unified_finetuning:
|
||||
torch.save(_trans_state_dict(self.colbert_linear.state_dict()),
|
||||
os.path.join(output_dir, 'colbert_linear.pt'))
|
||||
torch.save(_trans_state_dict(self.sparse_linear.state_dict()),
|
||||
os.path.join(output_dir, 'sparse_linear.pt'))
|
||||
|
||||
|
||||
class EncoderOnlyEmbedderM3ModelForInference(EncoderOnlyEmbedderM3Model):
|
||||
"""
|
||||
Inference class of M3 model.
|
||||
"""
|
||||
def forward(self,
|
||||
text_input: Dict[str, Tensor] = None,
|
||||
return_dense: bool = True,
|
||||
return_sparse: bool = False,
|
||||
return_colbert_vecs: bool = False,
|
||||
return_sparse_embedding: bool = False,
|
||||
truncate_dim: Optional[int] = None
|
||||
):
|
||||
"""Encode the text input using the selected way.
|
||||
|
||||
Args:
|
||||
text_input (Dict[str, Tensor], optional): Text inputs. Defaults to ``None``.
|
||||
return_dense (bool, optional): If True, return the dense embedding. Defaults to ``True``.
|
||||
return_sparse (bool, optional): If True, return the sparse embedding. Defaults to ``False``.
|
||||
return_colbert_vecs (bool, optional): If True, return the colbert vectors. Defaults to ``False``.
|
||||
return_sparse_embedding (bool, optional): Parameter for :meth:`_sparse_embedding()`. If True, will return sparse embedding.
|
||||
Otherwise, return the token weights. Defaults to ``False``.
|
||||
truncate_dim (Optional[int], optional): The dimension to truncate the output embeddings to. Useful for Matryoshka
|
||||
Representation Learning models. If None, no truncation is performed. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the three types of embeddings.
|
||||
"""
|
||||
assert return_dense or return_sparse or return_colbert_vecs, 'Must choose one or more from `return_colbert_vecs`, `return_sparse`, `return_dense` to set `True`!'
|
||||
|
||||
# this is for sparse embedding computation: using optimization suggestion from
|
||||
# issue #1364: https://github.com/FlagOpen/FlagEmbedding/issues/1364
|
||||
self.training = False
|
||||
|
||||
last_hidden_state = self.model(**text_input, return_dict=True).last_hidden_state
|
||||
|
||||
output = {}
|
||||
if return_dense:
|
||||
dense_vecs = self._dense_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
if truncate_dim is not None:
|
||||
dense_vecs = dense_vecs[..., :truncate_dim]
|
||||
output['dense_vecs'] = dense_vecs
|
||||
if return_sparse:
|
||||
sparse_vecs = self._sparse_embedding(
|
||||
last_hidden_state, text_input['input_ids'],
|
||||
return_embedding=return_sparse_embedding
|
||||
)
|
||||
output['sparse_vecs'] = sparse_vecs
|
||||
if return_colbert_vecs:
|
||||
colbert_vecs = self._colbert_embedding(last_hidden_state, text_input['attention_mask'])
|
||||
if truncate_dim is not None:
|
||||
colbert_vecs = colbert_vecs[..., :truncate_dim]
|
||||
output['colbert_vecs'] = colbert_vecs
|
||||
|
||||
if self.normalize_embeddings:
|
||||
if 'dense_vecs' in output:
|
||||
output['dense_vecs'] = F.normalize(output['dense_vecs'], dim=-1)
|
||||
if 'colbert_vecs' in output:
|
||||
output['colbert_vecs'] = F.normalize(output['colbert_vecs'], dim=-1)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,178 @@
|
||||
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
|
||||
@@ -0,0 +1,44 @@
|
||||
import os
|
||||
import torch
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from FlagEmbedding.abc.finetune.embedder import AbsEmbedderTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EncoderOnlyEmbedderM3Trainer(AbsEmbedderTrainer):
|
||||
"""
|
||||
Trainer class for M3.
|
||||
"""
|
||||
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