|
|
# 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 os
|
|
|
import random
|
|
|
|
|
|
import numpy as np
|
|
|
import paddle
|
|
|
from paddle.io import Dataset
|
|
|
from paddlenlp.transformers import BertTokenizer
|
|
|
# from speechtask.punctuation_restoration.utils.punct_prepro import load_dataset
|
|
|
|
|
|
__all__ = ["PuncDataset", "PuncDatasetFromBertTokenizer"]
|
|
|
|
|
|
|
|
|
class PuncDataset(Dataset):
|
|
|
"""Representing a Dataset
|
|
|
superclass
|
|
|
----------
|
|
|
data.Dataset :
|
|
|
Dataset is a abstract class, representing the real data.
|
|
|
"""
|
|
|
|
|
|
def __init__(self, train_path, vocab_path, punc_path, seq_len=100):
|
|
|
# 检查文件是否存在
|
|
|
print(train_path)
|
|
|
print(vocab_path)
|
|
|
assert os.path.exists(train_path), "train文件不存在"
|
|
|
assert os.path.exists(vocab_path), "词典文件不存在"
|
|
|
assert os.path.exists(punc_path), "标点文件不存在"
|
|
|
self.seq_len = seq_len
|
|
|
|
|
|
self.word2id = self.load_vocab(
|
|
|
vocab_path, extra_word_list=['<UNK>', '<END>'])
|
|
|
self.id2word = {v: k for k, v in self.word2id.items()}
|
|
|
self.punc2id = self.load_vocab(punc_path, extra_word_list=[" "])
|
|
|
self.id2punc = {k: v for (v, k) in self.punc2id.items()}
|
|
|
|
|
|
tmp_seqs = open(train_path, encoding='utf-8').readlines()
|
|
|
self.txt_seqs = [i for seq in tmp_seqs for i in seq.split()]
|
|
|
# print(self.txt_seqs[:10])
|
|
|
# with open('./txt_seq', 'w', encoding='utf-8') as w:
|
|
|
# print(self.txt_seqs, file=w)
|
|
|
self.preprocess(self.txt_seqs)
|
|
|
print('---punc-')
|
|
|
print(self.punc2id)
|
|
|
|
|
|
def __len__(self):
|
|
|
"""return the sentence nums in .txt
|
|
|
"""
|
|
|
return self.in_len
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
"""返回指定索引的张量对 (输入文本id的序列 , 其对应的标点id序列)
|
|
|
Parameters
|
|
|
----------
|
|
|
index : int 索引
|
|
|
"""
|
|
|
return self.input_data[index], self.label[index]
|
|
|
|
|
|
def load_vocab(self, vocab_path, extra_word_list=[], encoding='utf-8'):
|
|
|
n = len(extra_word_list)
|
|
|
with open(vocab_path, encoding='utf-8') as vf:
|
|
|
vocab = {word.strip(): i + n for i, word in enumerate(vf)}
|
|
|
for i, word in enumerate(extra_word_list):
|
|
|
vocab[word] = i
|
|
|
return vocab
|
|
|
|
|
|
def preprocess(self, txt_seqs: list):
|
|
|
"""将文本转为单词和应预测标点的id pair
|
|
|
Parameters
|
|
|
----------
|
|
|
txt : 文本
|
|
|
文本每个单词跟随一个空格，符号也跟一个空格
|
|
|
"""
|
|
|
input_data = []
|
|
|
label = []
|
|
|
input_r = []
|
|
|
label_r = []
|
|
|
# txt_seqs is a list like: ['char', 'char', 'char', '*，*', 'char', ......]
|
|
|
count = 0
|
|
|
length = len(txt_seqs)
|
|
|
for token in txt_seqs:
|
|
|
count += 1
|
|
|
if count == length:
|
|
|
break
|
|
|
if token in self.punc2id:
|
|
|
continue
|
|
|
punc = txt_seqs[count]
|
|
|
if punc not in self.punc2id:
|
|
|
# print('标点{}：'.format(count), self.punc2id[" "])
|
|
|
label.append(self.punc2id[" "])
|
|
|
input_data.append(
|
|
|
self.word2id.get(token, self.word2id["<UNK>"]))
|
|
|
input_r.append(token)
|
|
|
label_r.append(' ')
|
|
|
else:
|
|
|
# print('标点{}：'.format(count), self.punc2id[punc])
|
|
|
label.append(self.punc2id[punc])
|
|
|
input_data.append(
|
|
|
self.word2id.get(token, self.word2id["<UNK>"]))
|
|
|
input_r.append(token)
|
|
|
label_r.append(punc)
|
|
|
if len(input_data) != len(label):
|
|
|
assert 'error: length input_data != label'
|
|
|
# code below is for using 100 as a hidden size
|
|
|
print(len(input_data))
|
|
|
self.in_len = len(input_data) // self.seq_len
|
|
|
len_tmp = self.in_len * self.seq_len
|
|
|
input_data = input_data[:len_tmp]
|
|
|
label = label[:len_tmp]
|
|
|
|
|
|
self.input_data = paddle.to_tensor(
|
|
|
np.array(input_data, dtype='int64').reshape(-1, self.seq_len))
|
|
|
self.label = paddle.to_tensor(
|
|
|
np.array(label, dtype='int64').reshape(-1, self.seq_len))
|
|
|
|
|
|
|
|
|
# unk_token='[UNK]'
|
|
|
# sep_token='[SEP]'
|
|
|
# pad_token='[PAD]'
|
|
|
# cls_token='[CLS]'
|
|
|
# mask_token='[MASK]'
|
|
|
|
|
|
|
|
|
class PuncDatasetFromBertTokenizer(Dataset):
|
|
|
"""Representing a Dataset
|
|
|
superclass
|
|
|
----------
|
|
|
data.Dataset :
|
|
|
Dataset is a abstract class, representing the real data.
|
|
|
"""
|
|
|
|
|
|
def __init__(self,
|
|
|
train_path,
|
|
|
is_eval,
|
|
|
pretrained_token,
|
|
|
punc_path,
|
|
|
seq_len=100):
|
|
|
# 检查文件是否存在
|
|
|
print(train_path)
|
|
|
self.tokenizer = BertTokenizer.from_pretrained(
|
|
|
pretrained_token, do_lower_case=True)
|
|
|
self.paddingID = self.tokenizer.pad_token_id
|
|
|
assert os.path.exists(train_path), "train文件不存在"
|
|
|
assert os.path.exists(punc_path), "标点文件不存在"
|
|
|
self.seq_len = seq_len
|
|
|
|
|
|
self.punc2id = self.load_vocab(punc_path, extra_word_list=[" "])
|
|
|
self.id2punc = {k: v for (v, k) in self.punc2id.items()}
|
|
|
|
|
|
tmp_seqs = open(train_path, encoding='utf-8').readlines()
|
|
|
self.txt_seqs = [i for seq in tmp_seqs for i in seq.split()]
|
|
|
# print(self.txt_seqs[:10])
|
|
|
# with open('./txt_seq', 'w', encoding='utf-8') as w:
|
|
|
# print(self.txt_seqs, file=w)
|
|
|
if (is_eval):
|
|
|
self.preprocess(self.txt_seqs)
|
|
|
else:
|
|
|
self.preprocess_shift(self.txt_seqs)
|
|
|
print("data len: %d" % (len(self.input_data)))
|
|
|
print('---punc-')
|
|
|
print(self.punc2id)
|
|
|
|
|
|
def __len__(self):
|
|
|
"""return the sentence nums in .txt
|
|
|
"""
|
|
|
return self.in_len
|
|
|
|
|
|
def __getitem__(self, index):
|
|
|
"""返回指定索引的张量对 (输入文本id的序列 , 其对应的标点id序列)
|
|
|
Parameters
|
|
|
----------
|
|
|
index : int 索引
|
|
|
"""
|
|
|
return self.input_data[index], self.label[index]
|
|
|
|
|
|
def load_vocab(self, vocab_path, extra_word_list=[], encoding='utf-8'):
|
|
|
n = len(extra_word_list)
|
|
|
with open(vocab_path, encoding='utf-8') as vf:
|
|
|
vocab = {word.strip(): i + n for i, word in enumerate(vf)}
|
|
|
for i, word in enumerate(extra_word_list):
|
|
|
vocab[word] = i
|
|
|
return vocab
|
|
|
|
|
|
def preprocess(self, txt_seqs: list):
|
|
|
"""将文本转为单词和应预测标点的id pair
|
|
|
Parameters
|
|
|
----------
|
|
|
txt : 文本
|
|
|
文本每个单词跟随一个空格，符号也跟一个空格
|
|
|
"""
|
|
|
input_data = []
|
|
|
label = []
|
|
|
# txt_seqs is a list like: ['char', 'char', 'char', '*，*', 'char', ......]
|
|
|
count = 0
|
|
|
for i in range(len(txt_seqs) - 1):
|
|
|
word = txt_seqs[i]
|
|
|
punc = txt_seqs[i + 1]
|
|
|
if word in self.punc2id:
|
|
|
continue
|
|
|
|
|
|
token = self.tokenizer(word)
|
|
|
x = token["input_ids"][1:-1]
|
|
|
input_data.extend(x)
|
|
|
|
|
|
for i in range(len(x) - 1):
|
|
|
label.append(self.punc2id[" "])
|
|
|
|
|
|
if punc not in self.punc2id:
|
|
|
# print('标点{}：'.format(count), self.punc2id[" "])
|
|
|
label.append(self.punc2id[" "])
|
|
|
else:
|
|
|
label.append(self.punc2id[punc])
|
|
|
|
|
|
if len(input_data) != len(label):
|
|
|
assert 'error: length input_data != label'
|
|
|
# code below is for using 100 as a hidden size
|
|
|
|
|
|
# print(len(input_data[0]))
|
|
|
# print(len(label))
|
|
|
self.in_len = len(input_data) // self.seq_len
|
|
|
len_tmp = self.in_len * self.seq_len
|
|
|
input_data = input_data[:len_tmp]
|
|
|
label = label[:len_tmp]
|
|
|
# # print(input_data)
|
|
|
# print(type(input_data))
|
|
|
# tmp=np.array(input_data)
|
|
|
# print('--~~~~~~~~~~~~~')
|
|
|
# print(type(tmp))
|
|
|
# print(tmp.shape)
|
|
|
self.input_data = paddle.to_tensor(
|
|
|
np.array(input_data, dtype='int64').reshape(
|
|
|
-1, self.seq_len)) #, dtype='int64'
|
|
|
self.label = paddle.to_tensor(
|
|
|
np.array(label, dtype='int64').reshape(
|
|
|
-1, self.seq_len)) #, dtype='int64'
|
|
|
|
|
|
def preprocess_shift(self, txt_seqs: list):
|
|
|
"""将文本转为单词和应预测标点的id pair
|
|
|
Parameters
|
|
|
----------
|
|
|
txt : 文本
|
|
|
文本每个单词跟随一个空格，符号也跟一个空格
|
|
|
"""
|
|
|
input_data = []
|
|
|
label = []
|
|
|
# txt_seqs is a list like: ['char', 'char', 'char', '*，*', 'char', ......]
|
|
|
count = 0
|
|
|
for i in range(len(txt_seqs) - 1):
|
|
|
word = txt_seqs[i]
|
|
|
punc = txt_seqs[i + 1]
|
|
|
if word in self.punc2id:
|
|
|
continue
|
|
|
|
|
|
token = self.tokenizer(word)
|
|
|
x = token["input_ids"][1:-1]
|
|
|
input_data.extend(x)
|
|
|
|
|
|
for i in range(len(x) - 1):
|
|
|
label.append(self.punc2id[" "])
|
|
|
|
|
|
if punc not in self.punc2id:
|
|
|
# print('标点{}：'.format(count), self.punc2id[" "])
|
|
|
label.append(self.punc2id[" "])
|
|
|
else:
|
|
|
label.append(self.punc2id[punc])
|
|
|
|
|
|
if len(input_data) != len(label):
|
|
|
assert 'error: length input_data != label'
|
|
|
|
|
|
# print(len(input_data[0]))
|
|
|
# print(len(label))
|
|
|
start = 0
|
|
|
processed_data = []
|
|
|
processed_label = []
|
|
|
while (start < len(input_data) - self.seq_len):
|
|
|
# end=start+self.seq_len
|
|
|
end = random.randint(start + self.seq_len // 2,
|
|
|
start + self.seq_len)
|
|
|
processed_data.append(input_data[start:end])
|
|
|
processed_label.append(label[start:end])
|
|
|
|
|
|
start = start + random.randint(1, self.seq_len // 2)
|
|
|
|
|
|
self.in_len = len(processed_data)
|
|
|
# # print(input_data)
|
|
|
# print(type(input_data))
|
|
|
# tmp=np.array(input_data)
|
|
|
# print('--~~~~~~~~~~~~~')
|
|
|
# print(type(tmp))
|
|
|
# print(tmp.shape)
|
|
|
self.input_data = processed_data
|
|
|
#paddle.to_tensor(np.array(processed_data, dtype='int64')) #, dtype='int64'
|
|
|
self.label = processed_label
|
|
|
#paddle.to_tensor(np.array(processed_label, dtype='int64')) #, dtype='int64'
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
dataset = PuncDataset()
|
