Add text semantic matching for taskflow

docs/model_zoo/taskflow.md

@@ -35,15 +35,14 @@ PaddleNLP提供**开箱即用**的产业级NLP预置任务能力，无需训练
 | [信息抽取](#信息抽取)           | `Taskflow("information_extraction")`| ✅        | ✅        | ✅        | ✅         | ✅          | 适配多场景的开放域通用信息抽取工具                     |
 | [『解语』-知识标注](#解语知识标注) | `Taskflow("knowledge_mining")`     | ✅        | ✅        | ✅        | ✅          | ✅          | 覆盖所有中文词汇的知识标注工具                         |
 | [文本纠错](#文本纠错)              | `Taskflow("text_correction")`    | ✅        | ✅        | ✅        | ✅          | ✅          | 融合拼音特征的端到端文本纠错模型ERNIE-CSC              |
-| [文本相似度](#文本相似度)          | `Taskflow("text_similarity")`    | ✅        | ✅        | ✅        |            |            | 基于百度知道2200万对相似句组训练                       |
+| [文本相似度](#文本相似度)          | `Taskflow("text_similarity")`    | ✅        | ✅        | ✅        |            |            | 基于百万量级Dureader Retrieval数据集训练RocketQA并达到前沿文本相似效果|
 | [情感倾向分析](#情感倾向分析)      | `Taskflow("sentiment_analysis")`  | ✅        | ✅        | ✅        |            | ✅          | 基于情感知识增强预训练模型SKEP达到业界SOTA             |
 | [生成式问答](#生成式问答)          | `Taskflow("question_answering")` | ✅        | ✅        | ✅        |            |            | 使用最大中文开源CPM模型完成问答                        |
 | [智能写诗](#智能写诗)              | `Taskflow("poetry_generation")`  | ✅        | ✅        | ✅        |            |            | 使用最大中文开源CPM模型完成写诗                        |
 | [开放域对话](#开放域对话)          | `Taskflow("dialogue")`           | ✅        | ✅        | ✅        |            |            | 十亿级语料训练最强中文闲聊模型PLATO-Mini，支持多轮对话 |
 | [代码生成](#代码生成)          | `Taskflow("code_generation")`        | ✅        | ✅        | ✅        |            |            | 代码生成大模型 |
 | [文图生成](#文图生成)          | `Taskflow("text2image_generation")`        | ✅        | ✅        | ✅        |            |            | 文图生成大模型 |
 
-
 ## QuickStart
 
 **环境依赖**
@@ -1156,23 +1155,23 @@ from paddlenlp import Taskflow
 </div></details>
 
 ### 文本相似度
-<details><summary>&emsp;基于百度知道2200万对相似句组训练SimBERT达到前沿文本相似效果</summary><div>
+<details><summary>&emsp;基于百万量级Dureader Retrieval数据集训练RocketQA并达到前沿文本相似效果</summary><div>
 
 #### 单条输入
 
 ```python
 >>> from paddlenlp import Taskflow
->>> similarity = Taskflow("text_similarity")
+>>> similarity = Taskflow("text_similarity",model="rocketqa-zh-dureader-cross-encoder")
 >>> similarity([["春天适合种什么花？", "春天适合种什么菜？"]])
-[{'text1': '春天适合种什么花？', 'text2': '春天适合种什么菜？', 'similarity': 0.8340253}]
+[{'text1': '春天适合种什么花？', 'text2': '春天适合种什么菜？', 'similarity': 0.0048632388934493065}]
 ```
 
 #### 批量样本输入，平均速度更快
 
 ```python
 >>> from paddlenlp import Taskflow
->>> similarity([["光眼睛大就好看吗", "眼睛好看吗？"], ["小蝌蚪找妈妈怎么样", "小蝌蚪找妈妈是谁画的"]])
-[{'text1': '光眼睛大就好看吗', 'text2': '眼睛好看吗？', 'similarity': 0.74502707}, {'text1': '小蝌蚪找妈妈怎么样', 'text2': '小蝌蚪找妈妈是谁画的', 'similarity': 0.8192149}]
+>>> text_similarity([['春天适合种什么花？','春天适合种什么菜？'],['谁有狂三这张高清的','这张高清图，谁有']])
+[{'text1': '春天适合种什么花？', 'text2': '春天适合种什么菜？', 'similarity': 0.0048632388934493065}, {'text1': '谁有狂三这张高清的', 'text2': '这张高清图，谁有', 'similarity': 0.7050786018371582}]
 ```
 
 #### 可配置参数说明

paddlenlp/taskflow/taskflow.py

@@ -192,9 +192,18 @@
                 "task_class": TextSimilarityTask,
                 "task_flag": "text_similarity-simbert-base-chinese"
             },
+            "rocketqa-zh-dureader-cross-encoder": {
+                "task_class": TextSimilarityTask,
+                "task_flag":
+                'text_similarity-rocketqa-zh-dureader-cross-encoder',
+            },
+            "rocketqa-mini-cross-encoder": {
+                "task_class": TextSimilarityTask,
+                "task_flag": 'text_similarity-rocketqa-mini-cross-encoder',
+            },
         },
         "default": {
-            "model": "simbert-base-chinese"
+            "model": "rocketqa-zh-dureader-cross-encoder"
         }
     },
     "word_segmentation": {

paddlenlp/taskflow/text_similarity.py

@@ -14,6 +14,7 @@
 
 import paddle
 from paddlenlp.transformers import BertModel, BertTokenizer
+from ..transformers import ErnieCrossEncoder, ErnieTokenizer
 
 from ..data import Pad, Tuple
 from .utils import static_mode_guard
@@ -59,17 +60,21 @@ class TextSimilarityTask(Task):
                 "https://bj.bcebos.com/paddlenlp/taskflow/text_similarity/simbert-base-chinese/model_config.json",
                 "1254bbd7598457a9dad0afcb2e24b70c"
             ],
-        }
+        },
     }
 
-    def __init__(self, task, model, batch_size=1, max_seq_len=128, **kwargs):
+    def __init__(self, task, model, batch_size=1, max_seq_len=384, **kwargs):
         super().__init__(task=task, model=model, **kwargs)
-        self._check_task_files()
+        if ('rocketqa' not in model):
+            self._check_task_files()
+            self._get_inference_model()
+        else:
+            self._construct_model(model)
         self._construct_tokenizer(model)
-        self._get_inference_model()
         self._batch_size = batch_size
         self._max_seq_len = max_seq_len
         self._usage = usage
+        self.model_name = model
 
     def _construct_input_spec(self):
         """
@@ -88,15 +93,21 @@ def _construct_model(self, model):
         """
         Construct the inference model for the predictor.
         """
-        self._model = BertModel.from_pretrained(self._task_path,
-                                                pool_act='linear')
+        if ("rocketqa" in model):
+            self._model = ErnieCrossEncoder(model)
+        else:
+            self._model = BertModel.from_pretrained(self._task_path,
+                                                    pool_act='linear')
         self._model.eval()
 
     def _construct_tokenizer(self, model):
         """
         Construct the tokenizer for the predictor.
         """
-        self._tokenizer = BertTokenizer.from_pretrained(model)
+        if ("rocketqa" in model):
+            self._tokenizer = ErnieTokenizer.from_pretrained(model)
+        else:
+            self._tokenizer = BertTokenizer.from_pretrained(model)
 
     def _check_input_text(self, inputs):
         inputs = inputs[0]
@@ -118,40 +129,52 @@ def _preprocess(self, inputs):
             'lazy_load'] if 'lazy_load' in self.kwargs else False
 
         examples = []
-
         for data in inputs:
             text1, text2 = data[0], data[1]
-
-            text1_encoded_inputs = self._tokenizer(
-                text=text1, max_seq_len=self._max_seq_len)
-            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_seq_len=self._max_seq_len)
-            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))
+            if ("rocketqa" in self.model_name):
+                encoded_inputs = self._tokenizer(text=text1,
+                                                 text_pair=text2,
+                                                 max_seq_len=self._max_seq_len)
+                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_seq_len=self._max_seq_len)
+                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_seq_len=self._max_seq_len)
+                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)
         ]
-
-        batchify_fn = lambda samples, fn=Tuple(
-            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)]
+        if ("rocketqa" in self.model_name):
+            batchify_fn = lambda samples, fn=Tuple(
+                Pad(axis=0, pad_val=self._tokenizer.pad_token_id),  # input ids
+                Pad(axis=0, pad_val=self._tokenizer.pad_token_type_id
+                    ),  # token type ids
+            ): [data for data in fn(samples)]
+        else:
+            batchify_fn = lambda samples, fn=Tuple(
+                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
@@ -164,26 +187,35 @@ def _run_model(self, inputs):
         Run the task model from the outputs of the `_tokenize` function.
         """
         results = []
-        with static_mode_guard():
+        if ("rocketqa" in self.model_name):
             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)
+                input_ids, segment_ids = self._batchify_fn(batch)
+                input_ids = paddle.to_tensor(input_ids, dtype='int64')
+                segment_ids = paddle.to_tensor(segment_ids, dtype='int64')
+                scores = self._model.matching(input_ids=input_ids,
+                                              token_type_ids=segment_ids)
+                results.extend(scores.numpy().tolist())
+        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