wehub-resource-sync
430 lines
18 KiB
Python
430 lines
18 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Dict, List, Optional, Union
|
|
|
|
import numpy as np
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
import paddle.nn as nn
|
|
from tqdm import tqdm
|
|
|
|
from ...utils.log import logger
|
|
from .. import AutoConfig, AutoModel, AutoTokenizer, PretrainedModel
|
|
from ..model_outputs import ModelOutput
|
|
|
|
|
|
@dataclass
|
|
class EncoderOutput(ModelOutput):
|
|
q_reps: Optional[paddle.Tensor] = None
|
|
p_reps: Optional[paddle.Tensor] = None
|
|
loss: Optional[paddle.Tensor] = None
|
|
scores: Optional[paddle.Tensor] = None
|
|
|
|
|
|
__all__ = ["BiEncoderModel"]
|
|
|
|
|
|
class BiEncoderModel(PretrainedModel):
|
|
def __init__(
|
|
self,
|
|
model_name_or_path: str = None,
|
|
corpus_model_name_or_path: str = None,
|
|
query_model_name_or_path: str = None,
|
|
dtype: str = "float16",
|
|
normalized: bool = False,
|
|
sentence_pooling_method: str = "cls",
|
|
negatives_cross_device: bool = False,
|
|
temperature: float = 1.0,
|
|
use_inbatch_neg: bool = True,
|
|
margin: float = 0.3,
|
|
matryoshka_dims: Optional[List[int]] = None,
|
|
matryoshka_loss_weights: Optional[List[float]] = None,
|
|
query_instruction: Optional[str] = None,
|
|
document_instruction: Optional[str] = None,
|
|
eval_batch_size: int = 8,
|
|
tokenizer=None,
|
|
max_seq_length: int = 4096,
|
|
model_flag: str = None,
|
|
):
|
|
super().__init__()
|
|
|
|
# Load Model
|
|
self.model = None
|
|
self.model_config = None
|
|
self.corpus_model = None
|
|
self.query_model = None
|
|
if model_name_or_path is not None:
|
|
self.model = AutoModel.from_pretrained(model_name_or_path, dtype=dtype, convert_from_torch=True)
|
|
self.model_config = AutoConfig.from_pretrained(model_name_or_path)
|
|
if corpus_model_name_or_path is not None:
|
|
self.corpus_model = AutoModel.from_pretrained(
|
|
corpus_model_name_or_path, dtype=dtype, convert_from_torch=True
|
|
)
|
|
if query_model_name_or_path is not None:
|
|
self.query_model = AutoModel.from_pretrained(
|
|
query_model_name_or_path, dtype=dtype, convert_from_torch=True
|
|
)
|
|
if self.corpus_model is None:
|
|
self.corpus_model = self.model
|
|
if self.query_model is None:
|
|
self.query_model = self.model
|
|
assert self.corpus_model is not None and self.query_model is not None
|
|
|
|
self.cross_entropy = nn.CrossEntropyLoss(reduction="mean")
|
|
|
|
self.normalized = normalized
|
|
self.sentence_pooling_method = sentence_pooling_method
|
|
self.temperature = temperature
|
|
self.use_inbatch_neg = use_inbatch_neg
|
|
self.config = self.model_config
|
|
self.margin = margin
|
|
self.matryoshka_dims = matryoshka_dims
|
|
|
|
self.query_instruction = query_instruction
|
|
self.document_instruction = document_instruction
|
|
self.eval_batch_size = eval_batch_size
|
|
self.tokenizer = tokenizer
|
|
self.max_seq_length = max_seq_length
|
|
|
|
if self.matryoshka_dims:
|
|
self.matryoshka_loss_weights = (
|
|
matryoshka_loss_weights if matryoshka_loss_weights else [1] * len(self.matryoshka_dims)
|
|
)
|
|
else:
|
|
self.matryoshka_loss_weights = None
|
|
|
|
if not normalized:
|
|
self.temperature = 1.0
|
|
logger.info("reset temperature = 1.0 due to using inner product to compute similarity")
|
|
|
|
self.negatives_cross_device = negatives_cross_device
|
|
if self.negatives_cross_device:
|
|
if not dist.is_initialized():
|
|
raise ValueError("Distributed training has not been initialized for representation all gather.")
|
|
self.process_rank = dist.get_rank()
|
|
self.world_size = dist.get_world_size()
|
|
|
|
self.model_flag = model_flag
|
|
|
|
def sentence_embedding(self, hidden_state, mask):
|
|
if self.sentence_pooling_method == "mean":
|
|
s = paddle.sum(hidden_state * mask.unsqueeze(-1).float(), axis=1)
|
|
d = mask.sum(axis=1, keepdim=True).float()
|
|
return s / d
|
|
elif self.sentence_pooling_method == "cls":
|
|
return hidden_state[:, 0]
|
|
elif self.sentence_pooling_method == "last":
|
|
# return hidden_state[:, -1] # this is for padding side is left
|
|
sequence_lengths = mask.sum(axis=1)
|
|
last_token_indices = sequence_lengths - 1
|
|
embeddings = hidden_state[paddle.arange(hidden_state.shape[0]), last_token_indices]
|
|
return embeddings
|
|
elif self.sentence_pooling_method == "last_8":
|
|
last_8_embeddings = hidden_state[paddle.arange(hidden_state.shape[0]), -8:]
|
|
embeddings = paddle.mean(last_8_embeddings, axis=1)
|
|
return embeddings
|
|
else:
|
|
raise ValueError(f"Invalid sentence pooling method: {self.sentence_pooling_method}")
|
|
|
|
def get_model_config(
|
|
self,
|
|
):
|
|
return self.model_config.to_dict()
|
|
|
|
def encode(self, features, model: AutoModel):
|
|
psg_out = model(**features, return_dict=True, output_hidden_states=True)
|
|
p_reps = self.sentence_embedding(psg_out.hidden_states[-1], features["attention_mask"])
|
|
return p_reps
|
|
|
|
def compute_similarity(self, q_reps, p_reps):
|
|
# q_reps [batch_size, embedding_dim]
|
|
# p_reps [batch_size, embedding_dim]
|
|
return paddle.matmul(q_reps, p_reps.transpose([1, 0]))
|
|
|
|
def hard_negative_loss(self, q_reps, p_reps):
|
|
scores = self.compute_similarity(q_reps, p_reps)
|
|
scores = scores / self.temperature
|
|
scores = scores.reshape([q_reps.shape[0], -1])
|
|
|
|
target = paddle.arange(scores.shape[0], dtype="int64")
|
|
target = target * (p_reps.shape[0] // q_reps.shape[0])
|
|
loss = self.compute_loss(scores, target)
|
|
return scores, loss
|
|
|
|
def in_batch_negative_loss(self, q_reps, p_reps):
|
|
# In batch negatives
|
|
scores = self.compute_similarity(q_reps, p_reps)
|
|
# Subtract margin from all positive samples cosine_sim()
|
|
margin_diag = paddle.full(shape=[q_reps.shape[0]], fill_value=self.margin, dtype=q_reps.dtype)
|
|
scores = scores - paddle.diag(margin_diag)
|
|
# Scale cosine to ease training converge
|
|
scores = scores / self.temperature
|
|
target = paddle.arange(0, q_reps.shape[0], dtype="int64")
|
|
loss = self.compute_loss(scores, target)
|
|
return scores, loss
|
|
|
|
def forward(
|
|
self,
|
|
query: Dict[str, paddle.Tensor] = None,
|
|
passage: Dict[str, paddle.Tensor] = None,
|
|
teacher_score: paddle.Tensor = None,
|
|
):
|
|
q_reps = self.encode(query, self.query_model)
|
|
p_reps = self.encode(passage, self.corpus_model)
|
|
|
|
# For non-matryoshka loss, we normalize the representations
|
|
if not self.matryoshka_dims:
|
|
if self.normalized:
|
|
q_reps = paddle.nn.functional.normalize(q_reps, axis=-1)
|
|
p_reps = paddle.nn.functional.normalize(p_reps, axis=-1)
|
|
|
|
if self.training:
|
|
# Cross device negatives
|
|
if self.negatives_cross_device:
|
|
q_reps = self._dist_gather_tensor(q_reps)
|
|
p_reps = self._dist_gather_tensor(p_reps)
|
|
|
|
if self.matryoshka_dims:
|
|
loss = 0.0
|
|
scores = 0.0
|
|
for loss_weight, dim in zip(self.matryoshka_loss_weights, self.matryoshka_dims):
|
|
reduced_q = q_reps[:, :dim]
|
|
reduced_d = p_reps[:, :dim]
|
|
if self.normalized:
|
|
reduced_q = paddle.nn.functional.normalize(reduced_q, axis=-1)
|
|
reduced_d = paddle.nn.functional.normalize(reduced_d, axis=-1)
|
|
|
|
if self.use_inbatch_neg:
|
|
dim_score, dim_loss = self.in_batch_negative_loss(reduced_q, reduced_d)
|
|
else:
|
|
dim_score, dim_loss = self.hard_negative_loss(reduced_q, reduced_d)
|
|
scores += dim_score
|
|
loss += loss_weight * dim_loss
|
|
|
|
elif self.use_inbatch_neg:
|
|
scores, loss = self.in_batch_negative_loss(q_reps, p_reps)
|
|
else:
|
|
scores, loss = self.hard_negative_loss(q_reps, p_reps)
|
|
|
|
else:
|
|
scores = self.compute_similarity(q_reps, p_reps)
|
|
loss = None
|
|
return EncoderOutput(
|
|
loss=loss,
|
|
scores=scores,
|
|
q_reps=q_reps,
|
|
p_reps=p_reps,
|
|
)
|
|
|
|
def compute_loss(self, scores, target):
|
|
return self.cross_entropy(scores, target)
|
|
|
|
def _dist_gather_tensor(self, t: Optional[paddle.Tensor]):
|
|
if t is None:
|
|
return None
|
|
|
|
all_tensors = [paddle.empty_like(t) for _ in range(self.world_size)]
|
|
dist.all_gather(all_tensors, t)
|
|
|
|
all_tensors[self.process_rank] = t
|
|
all_tensors = paddle.concat(all_tensors, axis=0)
|
|
|
|
return all_tensors
|
|
|
|
def save_pretrained(self, output_dir: str, **kwargs):
|
|
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)
|
|
|
|
@paddle.no_grad()
|
|
def encode_sentences(
|
|
self, sentences: List[str], model: AutoModel, tokenizer: AutoTokenizer, titles: List[str] = None, **kwargs
|
|
) -> np.ndarray:
|
|
model.eval()
|
|
all_embeddings = []
|
|
for start_index in tqdm(range(0, len(sentences), self.eval_batch_size), desc="Batches"):
|
|
sentences_batch = sentences[start_index : start_index + self.eval_batch_size]
|
|
if titles:
|
|
titles_batch = titles[start_index : start_index + self.eval_batch_size]
|
|
assert len(sentences_batch) == len(titles_batch)
|
|
inputs = tokenizer(
|
|
titles_batch,
|
|
sentences_batch,
|
|
padding=True,
|
|
truncation=True,
|
|
return_tensors="pd",
|
|
max_length=self.max_seq_length,
|
|
return_attention_mask=True,
|
|
)
|
|
else:
|
|
inputs = tokenizer(
|
|
sentences_batch,
|
|
padding=True,
|
|
truncation=True,
|
|
return_tensors="pd",
|
|
max_length=self.max_seq_length,
|
|
return_attention_mask=True,
|
|
)
|
|
outputs = model(
|
|
**inputs, # 注意 bert 类型有 token_type_ids
|
|
return_dict=True,
|
|
output_hidden_states=True,
|
|
)
|
|
last_hidden_state = outputs.hidden_states[-1]
|
|
|
|
if self.sentence_pooling_method == "last":
|
|
if tokenizer.padding_side == "right":
|
|
sequence_lengths = inputs.attention_mask.sum(axis=1)
|
|
last_token_indices = sequence_lengths - 1
|
|
embeddings = last_hidden_state[paddle.arange(last_hidden_state.shape[0]), last_token_indices]
|
|
elif tokenizer.padding_side == "left":
|
|
embeddings = last_hidden_state[:, -1]
|
|
else:
|
|
raise NotImplementedError(f"Padding side {tokenizer.padding_side} not supported.")
|
|
elif self.sentence_pooling_method == "cls":
|
|
embeddings = last_hidden_state[:, 0]
|
|
elif self.sentence_pooling_method == "mean":
|
|
inputs.attention_mask = paddle.cast(
|
|
inputs.attention_mask, dtype="float32"
|
|
) # float cannot * int64, maybe paddle's bug
|
|
s = paddle.sum(last_hidden_state * inputs.attention_mask.unsqueeze(-1), axis=1)
|
|
d = inputs.attention_mask.sum(axis=1, keepdim=True)
|
|
embeddings = s / d
|
|
elif self.sentence_pooling_method == "last_8":
|
|
last_8_embeddings = last_hidden_state[:, -8:, :]
|
|
embeddings = paddle.mean(last_8_embeddings, axis=1)
|
|
else:
|
|
raise NotImplementedError(f"Pooling method {self.pooling_method} not supported.")
|
|
|
|
if self.normalized:
|
|
embeddings = paddle.nn.functional.normalize(embeddings, p=2, axis=-1)
|
|
|
|
all_embeddings.append(embeddings.cpu().numpy().astype("float32"))
|
|
|
|
return np.concatenate(all_embeddings, axis=0)
|
|
|
|
def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
|
|
"""
|
|
This function will be used to encode queries for retrieval task
|
|
if there is a instruction for queries, we will add it to the query text
|
|
"""
|
|
if self.query_instruction is not None:
|
|
input_texts = [f"{self.query_instruction}{query}" for query in queries]
|
|
else:
|
|
input_texts = queries
|
|
|
|
if self.model_flag == "llara":
|
|
input_texts = self.preprocess_sentences_for_llara(input_texts, query_or_doc="query")
|
|
if self.model_flag == "bge-en-icl":
|
|
input_texts = self.preprocess_sentences_for_bge_en_icl(input_texts, query_or_doc="query")
|
|
if self.model_flag == "qwen3":
|
|
input_texts = self.preprocess_sentences_for_qwen3(input_texts, query_or_doc="query")
|
|
|
|
encode_results = self.encode_sentences(sentences=input_texts, model=self.query_model, tokenizer=self.tokenizer)
|
|
return encode_results
|
|
|
|
def encode_corpus(self, corpus: List[Union[Dict[str, str], str]], **kwargs) -> np.ndarray:
|
|
"""
|
|
This function will be used to encode corpus for retrieval task
|
|
if there is a instruction for docs, we will add it to the doc text
|
|
"""
|
|
if isinstance(corpus[0], dict):
|
|
if self.document_instruction is not None:
|
|
input_texts = [
|
|
"{}{} {}".format(self.document_instruction, doc.get("title", ""), doc["text"]).strip()
|
|
for doc in corpus
|
|
]
|
|
else:
|
|
input_texts = ["{} {}".format(doc.get("title", ""), doc["text"]).strip() for doc in corpus]
|
|
else:
|
|
if self.document_instruction is not None:
|
|
input_texts = [f"{self.document_instruction}{doc}" for doc in corpus]
|
|
else:
|
|
input_texts = corpus
|
|
input_titles = None
|
|
|
|
if self.model_flag == "llara":
|
|
input_texts = self.preprocess_sentences_for_llara(input_texts, query_or_doc="doc")
|
|
if "RocketQA" in self.model_flag:
|
|
if isinstance(corpus[0], dict):
|
|
input_texts = [doc["text"] for doc in corpus]
|
|
input_titles = [doc.get("title", "") for doc in corpus]
|
|
if self.model_flag == "qwen3":
|
|
input_texts = self.preprocess_sentences_for_qwen3(input_texts, query_or_doc="doc")
|
|
|
|
encode_results = self.encode_sentences(
|
|
sentences=input_texts, titles=input_titles, model=self.corpus_model, tokenizer=self.tokenizer
|
|
)
|
|
return encode_results
|
|
|
|
def preprocess_sentences_for_bge_en_icl(self, sentences: List[str], query_or_doc: str, **kwargs) -> List[str]:
|
|
if query_or_doc == "query":
|
|
query_suffix = "\n<response> "
|
|
else:
|
|
raise ValueError(f"Invalid query_or_doc: {query_or_doc}")
|
|
|
|
input_texts = []
|
|
for query in sentences:
|
|
new_query = f"{query}{query_suffix}"
|
|
input_length = len(self.tokenizer(new_query)["input_ids"])
|
|
if input_length > self.max_seq_length:
|
|
cur_len = 0
|
|
add_len = 1
|
|
while add_len < len(query):
|
|
add_len *= 2
|
|
while add_len > 1:
|
|
add_len //= 2
|
|
assert isinstance(cur_len, int) and isinstance(
|
|
add_len, int
|
|
), f"cur_len={cur_len} add_len={add_len}"
|
|
new_query = f"{query[:cur_len+add_len]}{query_suffix}"
|
|
input_length = len(self.tokenizer(new_query)["input_ids"])
|
|
if input_length <= self.max_seq_length:
|
|
cur_len += add_len
|
|
new_query = f"{query[:cur_len]}{query_suffix}"
|
|
input_texts.append(new_query)
|
|
|
|
return input_texts
|
|
|
|
def preprocess_sentences_for_llara(self, sentences: List[str], query_or_doc: str, **kwargs) -> List[str]:
|
|
|
|
prefix = '"'
|
|
if query_or_doc == "query":
|
|
suffix = '", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>'
|
|
elif query_or_doc == "doc":
|
|
suffix = '", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>'
|
|
else:
|
|
raise ValueError(f"Invalid query_or_doc: {query_or_doc}")
|
|
|
|
sentences_after_process = []
|
|
for sentence in sentences:
|
|
inputs = self.tokenizer(
|
|
sentence,
|
|
return_tensors=None,
|
|
max_length=self.max_seq_length - 20,
|
|
truncation=True,
|
|
add_special_tokens=False,
|
|
)
|
|
sentences_after_process.append(self.tokenizer.decode(inputs["input_ids"], skip_special_tokens=True))
|
|
|
|
sentences_after_process = [prefix + " " + sentence + " " + suffix for sentence in sentences_after_process]
|
|
|
|
return sentences_after_process
|
|
|
|
def preprocess_sentences_for_qwen3(self, sentences: List[str], query_or_doc: str, **kwargs) -> List[str]:
|
|
sentences = [f"{sentence}<|endoftext|>" for sentence in sentences]
|
|
return sentences
|