wehub-resource-sync
586 lines
26 KiB
Python
586 lines
26 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 typing import Optional
|
|
|
|
import numpy as np
|
|
import paddle
|
|
|
|
from paddlenlp.data import DataCollatorWithPadding
|
|
from paddlenlp.transformers import AutoModel, AutoTokenizer, ErnieDualEncoder
|
|
|
|
from ..utils.log import logger
|
|
from .task import Task
|
|
from .utils import dygraph_mode_guard, static_mode_guard
|
|
|
|
ENCODER_TYPE = {
|
|
"rocketqa-zh-dureader-query-encoder": "query",
|
|
"rocketqa-zh-dureader-para-encoder": "paragraph",
|
|
"rocketqa-zh-base-query-encoder": "query",
|
|
"rocketqa-zh-base-para-encoder": "paragraph",
|
|
"rocketqa-zh-medium-query-encoder": "query",
|
|
"rocketqa-zh-medium-para-encoder": "paragraph",
|
|
"rocketqa-zh-mini-query-encoder": "query",
|
|
"rocketqa-zh-mini-para-encoder": "paragraph",
|
|
"rocketqa-zh-micro-query-encoder": "query",
|
|
"rocketqa-zh-micro-para-encoder": "paragraph",
|
|
"rocketqa-zh-nano-query-encoder": "query",
|
|
"rocketqa-zh-nano-para-encoder": "paragraph",
|
|
"rocketqav2-en-marco-query-encoder": "query",
|
|
"rocketqav2-en-marco-para-encoder": "paragraph",
|
|
"ernie-search-base-dual-encoder-marco-en": "query_paragraph",
|
|
}
|
|
|
|
|
|
usage = r"""
|
|
from paddlenlp import Taskflow
|
|
import paddle.nn.functional as F
|
|
# Text feature_extraction with rocketqa-zh-base-query-encoder
|
|
text_encoder = Taskflow("feature_extraction", model='rocketqa-zh-base-query-encoder')
|
|
text_embeds = text_encoder(['春天适合种什么花？','谁有狂三这张高清的?'])
|
|
text_features1 = text_embeds["features"]
|
|
print(text_features1)
|
|
'''
|
|
Tensor(shape=[2, 768], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[[ 0.27640465, -0.13405125, 0.00612330, ..., -0.15600294,
|
|
-0.18932408, -0.03029604],
|
|
[-0.12041329, -0.07424965, 0.07895312, ..., -0.17068857,
|
|
0.04485796, -0.18887770]])
|
|
'''
|
|
text_embeds = text_encoder('春天适合种什么菜？')
|
|
text_features2 = text_embeds["features"]
|
|
print(text_features2)
|
|
'''
|
|
Tensor(shape=[1, 768], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[[ 0.32578075, -0.02398480, -0.18929179, -0.18639392, -0.04062131,
|
|
0.06708499, -0.04631376, -0.41177100, -0.23074438, -0.23627219,
|
|
......
|
|
'''
|
|
probs = F.cosine_similarity(text_features1, text_features2)
|
|
print(probs)
|
|
'''
|
|
Tensor(shape=[2], dtype=float32, place=Place(gpu:0), stop_gradient=True,
|
|
[0.86455142, 0.41222256])
|
|
'''
|
|
"""
|
|
|
|
|
|
class TextFeatureExtractionTask(Task):
|
|
|
|
resource_files_names = {
|
|
"model_state": "model_state.pdparams",
|
|
"config": "config.json",
|
|
"vocab_file": "vocab.txt",
|
|
"special_tokens_map": "special_tokens_map.json",
|
|
"tokenizer_config": "tokenizer_config.json",
|
|
}
|
|
|
|
resource_files_urls = {
|
|
"rocketqa-zh-dureader-query-encoder": {
|
|
"model_state": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-dureader-query-encoder/model_state.pdparams",
|
|
"6125930530fd55ed715b0595e65789aa",
|
|
],
|
|
"config": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-dureader-query-encoder/config.json",
|
|
"efc1280069bb22b5bd06dc44b780bc6a",
|
|
],
|
|
"vocab_file": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-dureader-query-encoder/vocab.txt",
|
|
"062f696cad47bb62da86d8ae187b0ef4",
|
|
],
|
|
"special_tokens_map": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-dureader-query-encoder/special_tokens_map.json",
|
|
"8b3fb1023167bb4ab9d70708eb05f6ec",
|
|
],
|
|
"tokenizer_config": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-dureader-query-encoder/tokenizer_config.json",
|
|
"3a50349b8514e744fed72e59baca51b5",
|
|
],
|
|
},
|
|
"rocketqa-zh-base-query-encoder": {
|
|
"model_state": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-base-query-encoder/model_state.pdparams",
|
|
"3bb1a7870792146c6dd2fa47a45e15cc",
|
|
],
|
|
"config": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-base-query-encoder/config.json",
|
|
"be88115dd8a00e9de6b44f8c9a055e1a",
|
|
],
|
|
"vocab_file": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-base-query-encoder/vocab.txt",
|
|
"1c1c1f4fd93c5bed3b4eebec4de976a8",
|
|
],
|
|
"special_tokens_map": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-base-query-encoder/special_tokens_map.json",
|
|
"8b3fb1023167bb4ab9d70708eb05f6ec",
|
|
],
|
|
"tokenizer_config": [
|
|
"https://paddlenlp.bj.bcebos.com/taskflow/feature_extraction/rocketqa-zh-base-query-encoder/tokenizer_config.json",
|
|
"be86466f6769fde498690269d099ea7c",
|
|
],
|
|
},
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
task: str = None,
|
|
model: str = None,
|
|
batch_size: int = 1,
|
|
max_seq_len: int = 128,
|
|
_static_mode: bool = True,
|
|
return_tensors: str = "pd",
|
|
reinitialize: bool = False,
|
|
share_parameters: bool = False,
|
|
is_paragraph: bool = False,
|
|
output_emb_size: Optional[int] = None,
|
|
**kwargs
|
|
):
|
|
super().__init__(task=task, model=model, **kwargs)
|
|
self._seed = None
|
|
self.export_type = "text"
|
|
self._batch_size = batch_size
|
|
self.max_seq_len = max_seq_len
|
|
self.model = model
|
|
self._static_mode = _static_mode
|
|
self.return_tensors = return_tensors
|
|
|
|
self.reinitialize = reinitialize
|
|
self.share_parameters = share_parameters
|
|
self.output_emb_size = output_emb_size
|
|
self.is_paragraph = is_paragraph
|
|
self._check_para_encoder()
|
|
# self._check_task_files()
|
|
self._check_predictor_type()
|
|
self._construct_tokenizer()
|
|
# self._get_inference_model()
|
|
if self._static_mode:
|
|
self._get_inference_model()
|
|
else:
|
|
self._construct_model(model)
|
|
|
|
def _check_para_encoder(self):
|
|
if self.model in ENCODER_TYPE:
|
|
if ENCODER_TYPE[self.model] == "paragraph":
|
|
self.is_paragraph = True
|
|
else:
|
|
self.is_paragraph = False
|
|
else:
|
|
self.is_paragraph = False
|
|
|
|
def _construct_model(self, model):
|
|
"""
|
|
Construct the inference model for the predictor.
|
|
"""
|
|
# self._model = ErnieDualEncoder(self._task_path)
|
|
self._model = ErnieDualEncoder(
|
|
query_model_name_or_path=self.model,
|
|
output_emb_size=self.output_emb_size,
|
|
reinitialize=self.reinitialize,
|
|
share_parameters=self.share_parameters,
|
|
)
|
|
self._model.eval()
|
|
|
|
def _construct_tokenizer(self):
|
|
"""
|
|
Construct the tokenizer for the predictor.
|
|
"""
|
|
self._tokenizer = AutoTokenizer.from_pretrained(self.model)
|
|
if self._static_mode:
|
|
self._collator = DataCollatorWithPadding(self._tokenizer, return_tensors="np")
|
|
else:
|
|
self._collator = DataCollatorWithPadding(self._tokenizer, return_tensors="pd")
|
|
|
|
def _construct_input_spec(self):
|
|
"""
|
|
Construct the input spec for the convert dygraph model to static model.
|
|
"""
|
|
self._input_spec = [
|
|
paddle.static.InputSpec(shape=[None, None], dtype="int64", name="input_ids"),
|
|
paddle.static.InputSpec(shape=[None, None], dtype="int64", name="token_type_ids"),
|
|
]
|
|
|
|
def _batchify(self, data, batch_size):
|
|
"""
|
|
Generate input batches.
|
|
"""
|
|
|
|
def _parse_batch(batch_examples):
|
|
if self.is_paragraph:
|
|
# The input of the passage encoder is [CLS][SEP]...[SEP].
|
|
tokenized_inputs = self._tokenizer(
|
|
text=[""] * len(batch_examples),
|
|
text_pair=batch_examples,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_seq_len=self.max_seq_len,
|
|
)
|
|
else:
|
|
tokenized_inputs = self._tokenizer(
|
|
text=batch_examples,
|
|
padding="max_length",
|
|
truncation=True,
|
|
max_seq_len=self.max_seq_len,
|
|
)
|
|
return tokenized_inputs
|
|
|
|
# Separates data into some batches.
|
|
one_batch = []
|
|
for example in data:
|
|
one_batch.append(example)
|
|
if len(one_batch) == batch_size:
|
|
yield _parse_batch(one_batch)
|
|
one_batch = []
|
|
if one_batch:
|
|
yield _parse_batch(one_batch)
|
|
|
|
def _preprocess(self, inputs):
|
|
"""
|
|
Transform the raw inputs to the model inputs, two steps involved:
|
|
1) Transform the raw text/image to token ids/pixel_values.
|
|
2) Generate the other model inputs from the raw text/image and token ids/pixel_values.
|
|
"""
|
|
inputs = self._check_input_text(inputs)
|
|
batches = self._batchify(inputs, self._batch_size)
|
|
outputs = {"batches": batches, "inputs": inputs}
|
|
return outputs
|
|
|
|
def _run_model(self, inputs, **kwargs):
|
|
"""
|
|
Run the task model from the outputs of the `_preprocess` function.
|
|
"""
|
|
all_feats = []
|
|
if self._static_mode:
|
|
with static_mode_guard():
|
|
for batch_inputs in inputs["batches"]:
|
|
batch_inputs = self._collator(batch_inputs)
|
|
if self._predictor_type == "paddle-inference":
|
|
if "input_ids" in batch_inputs:
|
|
self.input_handles[0].copy_from_cpu(batch_inputs["input_ids"])
|
|
self.input_handles[1].copy_from_cpu(batch_inputs["token_type_ids"])
|
|
self.predictor.run()
|
|
text_features = self.output_handle[0].copy_to_cpu()
|
|
all_feats.append(text_features)
|
|
else:
|
|
# onnx mode
|
|
if "input_ids" in batch_inputs:
|
|
input_dict = {}
|
|
input_dict["input_ids"] = batch_inputs["input_ids"]
|
|
input_dict["token_type_ids"] = batch_inputs["token_type_ids"]
|
|
text_features = self.predictor.run(None, input_dict)[0].tolist()
|
|
all_feats.append(text_features)
|
|
|
|
else:
|
|
with dygraph_mode_guard():
|
|
for batch_inputs in inputs["batches"]:
|
|
batch_inputs = self._collator(batch_inputs)
|
|
text_features = self._model.get_pooled_embedding(
|
|
input_ids=batch_inputs["input_ids"], token_type_ids=batch_inputs["token_type_ids"]
|
|
)
|
|
all_feats.append(text_features.detach().numpy())
|
|
inputs.update({"features": all_feats})
|
|
return inputs
|
|
|
|
def _postprocess(self, inputs):
|
|
inputs["features"] = np.concatenate(inputs["features"], axis=0)
|
|
if self.return_tensors == "pd":
|
|
inputs["features"] = paddle.to_tensor(inputs["features"])
|
|
return inputs
|
|
|
|
def _convert_dygraph_to_static(self):
|
|
"""
|
|
Convert the dygraph model to static model.
|
|
"""
|
|
assert (
|
|
self._model is not None
|
|
), "The dygraph model must be created before converting the dygraph model to static model."
|
|
assert (
|
|
self._input_spec is not None
|
|
), "The input spec must be created before converting the dygraph model to static model."
|
|
logger.info("Converting to the inference model cost a little time.")
|
|
|
|
static_model = paddle.jit.to_static(self._model.get_pooled_embedding, input_spec=self._input_spec)
|
|
paddle.jit.save(static_model, self.inference_model_path)
|
|
logger.info("The inference model save in the path:{}".format(self.inference_model_path))
|
|
|
|
|
|
def text_length(text):
|
|
# {key: value} case
|
|
if isinstance(text, dict):
|
|
return len(next(iter(text.values())))
|
|
# Object has no len() method
|
|
elif not hasattr(text, "__len__"):
|
|
return 1
|
|
# Empty string or list of ints
|
|
elif len(text) == 0 or isinstance(text[0], int):
|
|
return len(text)
|
|
# Sum of length of individual strings
|
|
else:
|
|
return sum([len(t) for t in text])
|
|
|
|
|
|
class SentenceFeatureExtractionTask(Task):
|
|
|
|
resource_files_names = {
|
|
"model_state": "model_state.pdparams",
|
|
"config": "config.json",
|
|
"vocab_file": "vocab.txt",
|
|
"special_tokens_map": "special_tokens_map.json",
|
|
"tokenizer_config": "tokenizer_config.json",
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
task: str = None,
|
|
model: str = None,
|
|
batch_size: int = 1,
|
|
max_seq_len: int = 512,
|
|
_static_mode: bool = True,
|
|
return_tensors: str = "pd",
|
|
pooling_mode: str = "cls_token",
|
|
**kwargs
|
|
):
|
|
super().__init__(
|
|
task=task,
|
|
model=model,
|
|
pooling_mode=pooling_mode,
|
|
**kwargs,
|
|
)
|
|
self._seed = None
|
|
self.export_type = "text"
|
|
self._batch_size = batch_size
|
|
self.max_seq_len = max_seq_len
|
|
self.model = model
|
|
self._static_mode = _static_mode
|
|
self.return_tensors = return_tensors
|
|
self.pooling_mode = pooling_mode
|
|
self._check_predictor_type()
|
|
self._construct_tokenizer()
|
|
if self._static_mode:
|
|
self._get_inference_model()
|
|
else:
|
|
self._construct_model(model)
|
|
|
|
def _construct_model(self, model):
|
|
"""
|
|
Construct the inference model for the predictor.
|
|
"""
|
|
self._model = AutoModel.from_pretrained(self.model)
|
|
self._model.eval()
|
|
|
|
def _construct_tokenizer(self):
|
|
"""
|
|
Construct the tokenizer for the predictor.
|
|
"""
|
|
self._tokenizer = AutoTokenizer.from_pretrained(self.model)
|
|
self.pad_token_id = self._tokenizer.convert_tokens_to_ids(self._tokenizer.pad_token)
|
|
if self._static_mode:
|
|
self._collator = DataCollatorWithPadding(self._tokenizer, return_tensors="np")
|
|
else:
|
|
self._collator = DataCollatorWithPadding(self._tokenizer, return_tensors="pd")
|
|
|
|
def _construct_input_spec(self):
|
|
"""
|
|
Construct the input spec for the convert dygraph model to static model.
|
|
"""
|
|
self._input_spec = [
|
|
paddle.static.InputSpec(shape=[None, None], dtype="int64", name="input_ids"),
|
|
paddle.static.InputSpec(shape=[None, None], dtype="int64", name="token_type_ids"),
|
|
]
|
|
|
|
def _batchify(self, data, batch_size):
|
|
"""
|
|
Generate input batches.
|
|
"""
|
|
|
|
def _parse_batch(batch_examples, max_seq_len=None):
|
|
if isinstance(batch_examples[0], str):
|
|
to_tokenize = [batch_examples]
|
|
else:
|
|
batch1, batch2 = [], []
|
|
for text_tuple in batch_examples:
|
|
batch1.append(text_tuple[0])
|
|
batch2.append(text_tuple[1])
|
|
to_tokenize = [batch1, batch2]
|
|
to_tokenize = [[str(s).strip() for s in col] for col in to_tokenize]
|
|
if max_seq_len is None:
|
|
max_seq_len = self.max_seq_len
|
|
tokenized_inputs = self._tokenizer(
|
|
to_tokenize[0],
|
|
padding=True,
|
|
truncation="longest_first",
|
|
max_seq_len=max_seq_len,
|
|
)
|
|
return tokenized_inputs
|
|
|
|
# Separates data into some batches.
|
|
one_batch = []
|
|
self.length_sorted_idx = np.argsort([-text_length(sen) for sen in data])
|
|
sentences_sorted = [data[idx] for idx in self.length_sorted_idx]
|
|
|
|
for example in range(len(sentences_sorted)):
|
|
one_batch.append(sentences_sorted[example])
|
|
if len(one_batch) == batch_size:
|
|
yield _parse_batch(one_batch)
|
|
one_batch = []
|
|
if one_batch:
|
|
yield _parse_batch(one_batch)
|
|
|
|
def _preprocess(self, inputs):
|
|
"""
|
|
Transform the raw inputs to the model inputs, two steps involved:
|
|
1) Transform the raw text/image to token ids/pixel_values.
|
|
2) Generate the other model inputs from the raw text/image and token ids/pixel_values.
|
|
"""
|
|
inputs = self._check_input_text(inputs)
|
|
batches = self._batchify(inputs, self._batch_size)
|
|
outputs = {"batches": batches, "inputs": inputs}
|
|
return outputs
|
|
|
|
def _run_model(self, inputs, **kwargs):
|
|
"""
|
|
Run the task model from the outputs of the `_preprocess` function.
|
|
"""
|
|
pooling_mode = kwargs.get("pooling_mode", None)
|
|
if pooling_mode is None:
|
|
pooling_mode = self.pooling_mode
|
|
all_feats = []
|
|
if self._static_mode:
|
|
with static_mode_guard():
|
|
for batch_inputs in inputs["batches"]:
|
|
batch_inputs = self._collator(batch_inputs)
|
|
if self._predictor_type == "paddle-inference":
|
|
if "input_ids" in batch_inputs:
|
|
self.input_handles[0].copy_from_cpu(batch_inputs["input_ids"])
|
|
self.input_handles[1].copy_from_cpu(batch_inputs["token_type_ids"])
|
|
self.predictor.run()
|
|
token_embeddings = self.output_handle[0].copy_to_cpu()
|
|
if pooling_mode == "max_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
token_embeddings.dtype
|
|
)
|
|
input_mask_expanded = np.expand_dims(attention_mask, -1).repeat(
|
|
token_embeddings.shape[-1], axis=-1
|
|
)
|
|
token_embeddings[input_mask_expanded == 0] = -1e9
|
|
max_over_time = np.max(token_embeddings, 1)
|
|
all_feats.append(max_over_time)
|
|
elif pooling_mode == "mean_tokens" or pooling_mode == "mean_sqrt_len_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
token_embeddings.dtype
|
|
)
|
|
input_mask_expanded = np.expand_dims(attention_mask, -1).repeat(
|
|
token_embeddings.shape[-1], axis=-1
|
|
)
|
|
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, 1)
|
|
sum_mask = input_mask_expanded.sum(1)
|
|
sum_mask = np.clip(sum_mask, a_min=1e-9, a_max=np.max(sum_mask))
|
|
if pooling_mode == "mean_tokens":
|
|
all_feats.append(sum_embeddings / sum_mask)
|
|
elif pooling_mode == "mean_sqrt_len_tokens":
|
|
all_feats.append(sum_embeddings / np.sqrt(sum_mask))
|
|
else:
|
|
cls_token = token_embeddings[:, 0]
|
|
all_feats.append(cls_token)
|
|
else:
|
|
# onnx mode
|
|
if "input_ids" in batch_inputs:
|
|
input_dict = {}
|
|
input_dict["input_ids"] = batch_inputs["input_ids"]
|
|
input_dict["token_type_ids"] = batch_inputs["token_type_ids"]
|
|
token_embeddings = self.predictor.run(None, input_dict)[0]
|
|
if pooling_mode == "max_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
token_embeddings.dtype
|
|
)
|
|
input_mask_expanded = np.expand_dims(attention_mask, -1).repeat(
|
|
token_embeddings.shape[-1], axis=-1
|
|
)
|
|
token_embeddings[input_mask_expanded == 0] = -1e9
|
|
max_over_time = np.max(token_embeddings, 1)
|
|
all_feats.append(max_over_time)
|
|
elif pooling_mode == "mean_tokens" or pooling_mode == "mean_sqrt_len_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
token_embeddings.dtype
|
|
)
|
|
input_mask_expanded = np.expand_dims(attention_mask, -1).repeat(
|
|
token_embeddings.shape[-1], axis=-1
|
|
)
|
|
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, 1)
|
|
sum_mask = input_mask_expanded.sum(1)
|
|
sum_mask = np.clip(sum_mask, a_min=1e-9, a_max=np.max(sum_mask))
|
|
if pooling_mode == "mean_tokens":
|
|
all_feats.append(sum_embeddings / sum_mask)
|
|
elif pooling_mode == "mean_sqrt_len_tokens":
|
|
all_feats.append(sum_embeddings / np.sqrt(sum_mask))
|
|
else:
|
|
cls_token = token_embeddings[:, 0]
|
|
all_feats.append(cls_token)
|
|
else:
|
|
with dygraph_mode_guard():
|
|
for batch_inputs in inputs["batches"]:
|
|
batch_inputs = self._collator(batch_inputs)
|
|
token_embeddings = self._model(input_ids=batch_inputs["input_ids"])[0]
|
|
if pooling_mode == "max_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
self._model.pooler.dense.weight.dtype
|
|
)
|
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.shape)
|
|
token_embeddings[input_mask_expanded == 0] = -1e9
|
|
max_over_time = paddle.max(token_embeddings, 1).detach().numpy()
|
|
all_feats.append(max_over_time)
|
|
|
|
elif pooling_mode == "mean_tokens" or pooling_mode == "mean_sqrt_len_tokens":
|
|
attention_mask = (batch_inputs["input_ids"] != self.pad_token_id).astype(
|
|
self._model.pooler.dense.weight.dtype
|
|
)
|
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.shape)
|
|
sum_embeddings = paddle.sum(token_embeddings * input_mask_expanded, 1)
|
|
sum_mask = input_mask_expanded.sum(1)
|
|
sum_mask = paddle.clip(sum_mask, min=1e-9)
|
|
if pooling_mode == "mean_tokens":
|
|
text_features = sum_embeddings / sum_mask
|
|
all_feats.append(text_features.detach().numpy())
|
|
elif pooling_mode == "mean_sqrt_len_tokens":
|
|
text_features = sum_embeddings / paddle.sqrt(sum_mask)
|
|
all_feats.append(text_features.detach().numpy())
|
|
else:
|
|
cls_token = token_embeddings[:, 0].detach().numpy()
|
|
all_feats.append(cls_token)
|
|
inputs.update({"features": all_feats})
|
|
return inputs
|
|
|
|
def _postprocess(self, inputs):
|
|
inputs["features"] = np.concatenate(inputs["features"], axis=0)
|
|
inputs["features"] = [inputs["features"][idx] for idx in np.argsort(self.length_sorted_idx)]
|
|
if self.return_tensors == "pd":
|
|
inputs["features"] = paddle.to_tensor(inputs["features"])
|
|
return inputs
|
|
|
|
def _convert_dygraph_to_static(self):
|
|
"""
|
|
Convert the dygraph model to static model.
|
|
"""
|
|
assert (
|
|
self._model is not None
|
|
), "The dygraph model must be created before converting the dygraph model to static model."
|
|
assert (
|
|
self._input_spec is not None
|
|
), "The input spec must be created before converting the dygraph model to static model."
|
|
logger.info("Converting to the inference model cost a little time.")
|
|
|
|
static_model = paddle.jit.to_static(self._model, input_spec=self._input_spec)
|
|
paddle.jit.save(static_model, self.inference_model_path)
|
|
logger.info("The inference model save in the path:{}".format(self.inference_model_path))
|