# coding:utf-8
# 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 glob
import json
import math
import os
import copy
import itertools
import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from ..transformers import ErnieTokenizer, ErnieModel
from ..transformers import is_chinese_char
from ..datasets import load_dataset
from ..data import Stack, Pad, Tuple, Vocab
from .utils import download_file, add_docstrings, static_mode_guard
from .models import ErnieForCSC
from .task import Task
usage = r"""
from paddlenlp import Taskflow
text_correction = Taskflow("text_correction")
text_correction('遇到逆竟时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。')
'''
[{'source': '遇到逆竟时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。',
'target': '遇到逆境时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。',
'errors': [{'position': 3, 'correction': {'竟': '境'}}]}
]
'''
text_correction(['遇到逆竟时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。',
'人生就是如此，经过磨练才能让自己更加拙壮，才能使自己更加乐观。'])
'''
[{'source': '遇到逆竟时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。',
'target': '遇到逆境时，我们必须勇于面对，而且要愈挫愈勇，这样我们才能朝著成功之路前进。',
'errors': [{'position': 3, 'correction': {'竟': '境'}}]},
{'source': '人生就是如此，经过磨练才能让自己更加拙壮，才能使自己更加乐观。',
'target': '人生就是如此，经过磨练才能让自己更加茁壮，才能使自己更加乐观。',
'errors': [{'position': 18, 'correction': {'拙': '茁'}}]}
]
'''
"""
TASK_MODEL_MAP = {"ernie-csc": "ernie-1.0"}
[文档]class CSCTask(Task):
"""
The text generation model to predict the question or chinese poetry.
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",
"pinyin_vocab": "pinyin_vocab.txt"
}
resource_files_urls = {
"ernie-csc": {
"model_state": [
"https://bj.bcebos.com/paddlenlp/taskflow/text_correction/ernie-csc/model_state.pdparams",
"cdc53e7e3985ffc78fedcdf8e6dca6d2"
],
"pinyin_vocab": [
"https://bj.bcebos.com/paddlenlp/taskflow/text_correction/ernie-csc/pinyin_vocab.txt",
"5599a8116b6016af573d08f8e686b4b2"
],
}
}
def __init__(self, task, model, **kwargs):
super().__init__(task=task, model=model, **kwargs)
self._usage = usage
self._check_task_files()
self._construct_vocabs()
self._get_inference_model()
self._construct_tokenizer(model)
try:
import pypinyin
except:
raise ImportError(
"Please install the dependencies first, pip install pypinyin --upgrade"
)
self._pypinyin = pypinyin
self._batchify_fn = lambda samples, fn=Tuple(
Pad(axis=0, pad_val=self._tokenizer.pad_token_id, dtype='int64'), # input
Pad(axis=0, pad_val=self._tokenizer.pad_token_type_id, dtype='int64'), # segment
Pad(axis=0, pad_val=self._pinyin_vocab.token_to_idx[self._pinyin_vocab.pad_token], dtype='int64'), # pinyin
Stack(axis=0, dtype='int64'), # length
): [data for data in fn(samples)]
self._num_workers = self.kwargs[
'num_workers'] if 'num_workers' in self.kwargs else 0
self._batch_size = self.kwargs[
'batch_size'] if 'batch_size' in self.kwargs else 1
self._lazy_load = self.kwargs[
'lazy_load'] if 'lazy_load' in self.kwargs else False
self._max_seq_len = self.kwargs[
'max_seq_len'] if 'max_seq_len' in self.kwargs else 128
self._split_sentence = self.kwargs[
'split_sentence'] if 'split_sentence' in self.kwargs else False
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='pinyin_ids'),
]
def _construct_vocabs(self):
pinyin_vocab_path = os.path.join(self._task_path, "pinyin_vocab.txt")
self._pinyin_vocab = Vocab.load_vocabulary(
pinyin_vocab_path, unk_token='[UNK]', pad_token='[PAD]')
def _construct_model(self, model):
"""
Construct the inference model for the predictor.
"""
ernie = ErnieModel.from_pretrained(TASK_MODEL_MAP[model])
model_instance = ErnieForCSC(
ernie,
pinyin_vocab_size=len(self._pinyin_vocab),
pad_pinyin_id=self._pinyin_vocab[self._pinyin_vocab.pad_token])
# Load the model parameter for the predict
model_path = os.path.join(self._task_path, "model_state.pdparams")
state_dict = paddle.load(model_path)
model_instance.set_state_dict(state_dict)
self._model = model_instance
self._model.eval()
def _construct_tokenizer(self, model):
"""
Construct the tokenizer for the predictor.
"""
self._tokenizer = ErnieTokenizer.from_pretrained(TASK_MODEL_MAP[model])
def _preprocess(self, inputs, padding=True, add_special_tokens=True):
input_texts = self._check_input_text(inputs)
examples = []
texts = []
max_predict_len = self._max_seq_len - 2
short_input_texts, self.input_mapping = self._auto_splitter(
input_texts, max_predict_len, split_sentence=self._split_sentence)
for text in short_input_texts:
if not (isinstance(text, str) and len(text) > 0):
continue
example = {"source": text.strip()}
input_ids, token_type_ids, pinyin_ids, length = self._convert_example(
example)
examples.append((input_ids, token_type_ids, pinyin_ids, length))
texts.append(example["source"])
batch_examples = [
examples[idx:idx + self._batch_size]
for idx in range(0, len(examples), self._batch_size)
]
batch_texts = [
short_input_texts[idx:idx + self._batch_size]
for idx in range(0, len(examples), self._batch_size)
]
outputs = {}
outputs['batch_examples'] = batch_examples
outputs['batch_texts'] = batch_texts
return outputs
def _run_model(self, inputs):
"""
Run the task model from the outputs of the `_tokenize` function.
"""
results = []
with static_mode_guard():
for examples in inputs['batch_examples']:
token_ids, token_type_ids, pinyin_ids, lengths = self._batchify_fn(
examples)
self.input_handles[0].copy_from_cpu(token_ids)
self.input_handles[1].copy_from_cpu(pinyin_ids)
self.predictor.run()
det_preds = self.output_handle[0].copy_to_cpu()
char_preds = self.output_handle[1].copy_to_cpu()
batch_result = []
for i in range(len(lengths)):
batch_result.append(
(det_preds[i], char_preds[i], lengths[i]))
results.append(batch_result)
inputs['batch_results'] = results
return inputs
def _postprocess(self, inputs):
"""
The model output is the logits and probs, this function will convert the model output to raw text.
"""
results = []
for examples, texts, temp_results in zip(inputs['batch_examples'],
inputs['batch_texts'],
inputs['batch_results']):
for i in range(len(examples)):
result = {}
det_pred, char_preds, length = temp_results[i]
pred_result = self._parse_decode(texts[i], char_preds, det_pred,
length)
result['source'] = texts[i]
result['target'] = ''.join(pred_result)
results.append(result)
results = self._auto_joiner(results, self.input_mapping, is_dict=True)
for result in results:
errors_result = []
for i, (source_token, target_token
) in enumerate(zip(result['source'], result['target'])):
if source_token != target_token:
errors_result.append({
'position': i,
'correction': {
source_token: target_token
}
})
result['errors'] = errors_result
return results
def _convert_example(self, example):
source = example["source"]
words = list(source)
length = len(words)
words = ['[CLS]'] + words + ['[SEP]']
input_ids = self._tokenizer.convert_tokens_to_ids(words)
token_type_ids = [0] * len(input_ids)
# Use pad token in pinyin emb to map word emb [CLS], [SEP]
pinyins = self._pypinyin.lazy_pinyin(
source,
style=self._pypinyin.Style.TONE3,
neutral_tone_with_five=True)
pinyin_ids = [0]
# Align pinyin and chinese char
pinyin_offset = 0
for i, word in enumerate(words[1:-1]):
pinyin = '[UNK]' if word != '[PAD]' else '[PAD]'
if len(word) == 1 and is_chinese_char(ord(word)):
while pinyin_offset < len(pinyins):
current_pinyin = pinyins[pinyin_offset][:-1]
pinyin_offset += 1
if current_pinyin in self._pinyin_vocab:
pinyin = current_pinyin
break
pinyin_ids.append(self._pinyin_vocab[pinyin])
pinyin_ids.append(0)
assert len(input_ids) == len(
pinyin_ids
), "length of input_ids must be equal to length of pinyin_ids"
return input_ids, token_type_ids, pinyin_ids, length
def _parse_decode(self, words, corr_preds, det_preds, lengths):
UNK = self._tokenizer.unk_token
UNK_id = self._tokenizer.convert_tokens_to_ids(UNK)
corr_pred = corr_preds[1:1 + lengths].tolist()
det_pred = det_preds[1:1 + lengths].tolist()
words = list(words)
rest_words = []
max_seq_length = self._max_seq_len - 2
if len(words) > max_seq_length:
rest_words = words[max_seq_length:]
words = words[:max_seq_length]
pred_result = ""
for j, word in enumerate(words):
candidates = self._tokenizer.convert_ids_to_tokens(corr_pred[
j] if corr_pred[j] < self._tokenizer.vocab_size else UNK_id)
word_icc = is_chinese_char(ord(word))
cand_icc = is_chinese_char(ord(candidates)) if len(
candidates) == 1 else False
if not word_icc or det_pred[j] == 0\
or candidates in [UNK, '[PAD]']\
or (word_icc and not cand_icc):
pred_result += word
else:
pred_result += candidates.lstrip("##")
pred_result += ''.join(rest_words)
return pred_result