paddlenlp.taskflow.text_similarity 源代码

# Copyright (c) 2021 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.

import paddle

from paddlenlp.transformers import AutoModel, AutoTokenizer

from ..data import Pad, Tuple
from ..transformers import ErnieCrossEncoder, ErnieTokenizer
from ..utils.log import logger
from .task import Task
from .utils import static_mode_guard

usage = r"""
         from paddlenlp import Taskflow

         similarity = Taskflow("text_similarity")
         similarity([["世界上什么东西最小", "世界上什么东西最小?"]])
         '''
         [{'text1': '世界上什么东西最小', 'text2': '世界上什么东西最小?', 'similarity': 0.992725}]
         '''

         similarity = Taskflow("text_similarity", batch_size=2)
         similarity([["光眼睛大就好看吗", "眼睛好看吗?"], ["小蝌蚪找妈妈怎么样", "小蝌蚪找妈妈是谁画的"]])
         '''
         [{'text1': '光眼睛大就好看吗', 'text2': '眼睛好看吗?', 'similarity': 0.74502707}, {'text1': '小蝌蚪找妈妈怎么样', 'text2': '小蝌蚪找妈妈是谁画的', 'similarity': 0.8192149}]
         '''
         """
MATCH_TYPE = {
    "rocketqa-zh-dureader-cross-encoder": "matching",
    "rocketqa-base-cross-encoder": "matching",
    "rocketqa-medium-cross-encoder": "matching",
    "rocketqa-mini-cross-encoder": "matching",
    "rocketqa-micro-cross-encoder": "matching",
    "rocketqa-nano-cross-encoder": "matching",
    "rocketqav2-en-marco-cross-encoder": "matching_v2",
    "ernie-search-large-cross-encoder-marco-en": "matching_v3",
}


[文档] class TextSimilarityTask(Task): """ Text similarity task using SimBERT to predict the similarity of sentence pair. Args: task(string): The name of task. model(string): The model name in the task. kwargs (dict, optional): Additional keyword arguments passed along to the specific task. """ resource_files_names = { "model_state": "model_state.pdparams", "model_config": "model_config.json", } resource_files_urls = { "simbert-base-chinese": { "model_state": [ "https://bj.bcebos.com/paddlenlp/taskflow/text_similarity/simbert-base-chinese/model_state.pdparams", "27d9ef240c2e8e736bdfefea52af2542", ], "model_config": [ "https://bj.bcebos.com/paddlenlp/taskflow/text_similarity/simbert-base-chinese/model_config.json", "1254bbd7598457a9dad0afcb2e24b70c", ], }, "rocketqa-zh-dureader-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-zh-dureader-cross-encoder/model_state.pdparams", "88bc3e1a64992a1bdfe4044ecba13bc7", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-zh-dureader-cross-encoder/model_config.json", "b69083c2895e8f68e1a10467b384daab", ], }, "rocketqa-base-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-base-cross-encoder/model_state.pdparams", "6d845a492a2695e62f2be79f8017be92", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-base-cross-encoder/model_config.json", "18ce260ede18bc3cb28dcb2e7df23b1a", ], }, "rocketqa-medium-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-medium-cross-encoder/model_state.pdparams", "4b929f4fc11a1df8f59fdf2784e23fa7", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-medium-cross-encoder/model_config.json", "10997db96bc86e29cd113e1bf58989d7", ], }, "rocketqa-mini-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-mini-cross-encoder/model_state.pdparams", "c411111df990132fb88c070d8b8cf3f7", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-mini-cross-encoder/model_config.json", "271e6d779acbe8e8acdd596b1c835546", ], }, "rocketqa-micro-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-micro-cross-encoder/model_state.pdparams", "3d643ff7d6029c8ceab5653680167dc0", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-micro-cross-encoder/model_config.json", "b32d1a932d8c367fab2a6216459dd0a7", ], }, "rocketqa-nano-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-nano-cross-encoder/model_state.pdparams", "4c1d36e5e94f5af09f665fc7ad0be140", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqa-nano-cross-encoder/model_config.json", "dcff14cd671e1064be2c5d63734098bb", ], }, "rocketqav2-en-marco-cross-encoder": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqav2-en-marco-cross-encoder/model_state.pdparams", "a5afc77b6a63fc32a1beca3010f40f32", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/rocketqav2-en-marco-cross-encoder/config.json", "8f5d5c71c8a891b68d0402a13e38b6f9", ], }, "ernie-search-large-cross-encoder-marco-en": { "model_state": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/ernie-search-large-cross-encoder-marco-en/model_state.pdparams", "fdf29f7de0f7fe570740d343c96165e5", ], "model_config": [ "https://paddlenlp.bj.bcebos.com/taskflow/text_similarity/ernie-search-large-cross-encoder-marco-en/config.json", "28bad2c7b36fa148fa75a8dc5b690485", ], }, "__internal_testing__/tiny-random-bert": { "model_state": [ "https://bj.bcebos.com/paddlenlp/models/community/__internal_testing__/tiny-random-bert/model_state.pdparams", "8d8814d589c21bf083fdb35de6c11a57", ], "model_config": [ "https://bj.bcebos.com/paddlenlp/models/community/__internal_testing__/tiny-random-bert/config.json", "37e28e2359f330f64fc82beff1967a1e", ], }, } def __init__(self, task, model, batch_size=1, max_length=384, **kwargs): super().__init__(task=task, model=model, **kwargs) self._static_mode = True self._check_predictor_type() if not self.from_hf_hub: self._check_task_files() if self._static_mode: self._get_inference_model() else: self._construct_model(model) self._construct_tokenizer(model) self._batch_size = batch_size self._max_length = max_length self._usage = usage self.model_name = model 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 _construct_model(self, model): """ Construct the inference model for the predictor. """ if "rocketqav2-en" in model or "ernie-search" in model: self._model = ErnieCrossEncoder(self._task_path, num_classes=1, reinitialize=True) elif "rocketqa" in model: self._model = ErnieCrossEncoder(self._task_path, num_classes=2) else: self._model = AutoModel.from_pretrained(self._task_path, pool_act="linear") self._model.eval() def _construct_tokenizer(self, model): """ Construct the tokenizer for the predictor. """ if "rocketqa" in model or "ernie-search" in model: self._tokenizer = ErnieTokenizer.from_pretrained(model) else: self._tokenizer = AutoTokenizer.from_pretrained(model) def _check_input_text(self, inputs): inputs = inputs[0] if not all([isinstance(i, list) and i and all(i) and len(i) == 2 for i in inputs]): raise TypeError("Invalid input format.") return inputs def _preprocess(self, inputs): """ Transform the raw text to the model inputs, two steps involved: 1) Transform the raw text to token ids. 2) Generate the other model inputs from the raw text and token ids. """ inputs = self._check_input_text(inputs) examples = [] for data in inputs: text1, text2 = data[0], data[1] if "rocketqa" in self.model_name or "ernie-search" in self.model_name: # Todo: wugaosheng, Add erine-search encoding support encoded_inputs = self._tokenizer(text=text1, text_pair=text2, max_length=self._max_length) ids = encoded_inputs["input_ids"] segment_ids = encoded_inputs["token_type_ids"] examples.append((ids, segment_ids)) else: text1_encoded_inputs = self._tokenizer(text=text1, max_length=self._max_length) text1_input_ids = text1_encoded_inputs["input_ids"] text1_token_type_ids = text1_encoded_inputs["token_type_ids"] text2_encoded_inputs = self._tokenizer(text=text2, max_length=self._max_length) text2_input_ids = text2_encoded_inputs["input_ids"] text2_token_type_ids = text2_encoded_inputs["token_type_ids"] examples.append((text1_input_ids, text1_token_type_ids, text2_input_ids, text2_token_type_ids)) batches = [examples[idx : idx + self._batch_size] for idx in range(0, len(examples), self._batch_size)] if "rocketqa" in self.model_name or "ernie-search" in self.model_name: batchify_fn = lambda samples, fn=Tuple( # noqa: E731 Pad(axis=0, pad_val=self._tokenizer.pad_token_id, dtype="int64"), # input ids Pad(axis=0, pad_val=self._tokenizer.pad_token_type_id, dtype="int64"), # token type ids ): [data for data in fn(samples)] else: batchify_fn = lambda samples, fn=Tuple( # noqa: E731 Pad(axis=0, pad_val=self._tokenizer.pad_token_id, dtype="int64"), # text1_input_ids Pad(axis=0, pad_val=self._tokenizer.pad_token_type_id, dtype="int64"), # text1_token_type_ids Pad(axis=0, pad_val=self._tokenizer.pad_token_id, dtype="int64"), # text2_input_ids Pad(axis=0, pad_val=self._tokenizer.pad_token_type_id, dtype="int64"), # text2_token_type_ids ): [data for data in fn(samples)] outputs = {} outputs["data_loader"] = batches outputs["text"] = inputs self._batchify_fn = batchify_fn return outputs def _run_model(self, inputs): """ Run the task model from the outputs of the `_tokenize` function. """ results = [] if "rocketqa" in self.model_name or "ernie-search" in self.model_name: with static_mode_guard(): for batch in inputs["data_loader"]: if self._predictor_type == "paddle-inference": input_ids, segment_ids = self._batchify_fn(batch) self.input_handles[0].copy_from_cpu(input_ids) self.input_handles[1].copy_from_cpu(segment_ids) self.predictor.run() scores = self.output_handle[0].copy_to_cpu().tolist() results.extend(scores) else: # onnx mode input_dict = {} input_ids, segment_ids = self._batchify_fn(batch) input_dict["input_ids"] = input_ids input_dict["token_type_ids"] = segment_ids scores = self.predictor.run(None, input_dict)[0].tolist() results.extend(scores) else: with static_mode_guard(): for batch in inputs["data_loader"]: text1_ids, text1_segment_ids, text2_ids, text2_segment_ids = self._batchify_fn(batch) self.input_handles[0].copy_from_cpu(text1_ids) self.input_handles[1].copy_from_cpu(text1_segment_ids) self.predictor.run() vecs_text1 = self.output_handle[1].copy_to_cpu() self.input_handles[0].copy_from_cpu(text2_ids) self.input_handles[1].copy_from_cpu(text2_segment_ids) self.predictor.run() vecs_text2 = self.output_handle[1].copy_to_cpu() vecs_text1 = vecs_text1 / (vecs_text1**2).sum(axis=1, keepdims=True) ** 0.5 vecs_text2 = vecs_text2 / (vecs_text2**2).sum(axis=1, keepdims=True) ** 0.5 similarity = (vecs_text1 * vecs_text2).sum(axis=1) results.extend(similarity) inputs["result"] = results return inputs def _postprocess(self, inputs): """ The model output is tag ids, this function will convert the model output to raw text. """ final_results = [] for text, similarity in zip(inputs["text"], inputs["result"]): result = {} result["text1"] = text[0] result["text2"] = text[1] # The numpy.float32 can not be converted to the json format if isinstance(similarity, list): result["similarity"] = float(similarity[0]) else: result["similarity"] = float(similarity) final_results.append(result) return final_results 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.") if self.model in MATCH_TYPE: if MATCH_TYPE[self.model] == "matching": static_model = paddle.jit.to_static(self._model.matching, input_spec=self._input_spec) elif MATCH_TYPE[self.model] == "matching_v2": static_model = paddle.jit.to_static(self._model.matching_v2, input_spec=self._input_spec) elif MATCH_TYPE[self.model] == "matching_v3": static_model = paddle.jit.to_static(self._model.matching_v3, input_spec=self._input_spec) else: 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))