wehub-resource-sync
49 lines
2.0 KiB
Python
49 lines
2.0 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2020-06-03 18:50
|
|
from typing import Union, List, Callable
|
|
|
|
from hanlp.common.dataset import TransformableDataset
|
|
from hanlp.utils.io_util import get_resource
|
|
from hanlp.utils.span_util import bmes_of
|
|
from hanlp.utils.string_util import ispunct
|
|
|
|
|
|
class ChunkingDataset(TransformableDataset):
|
|
|
|
def __init__(self, data: Union[str, List], transform: Union[Callable, List] = None, cache=None,
|
|
generate_idx=None, max_seq_len=None, sent_delimiter=None) -> None:
|
|
if not sent_delimiter:
|
|
sent_delimiter = lambda x: ispunct(x)
|
|
elif isinstance(sent_delimiter, str):
|
|
sent_delimiter = set(list(sent_delimiter))
|
|
sent_delimiter = lambda x: x in sent_delimiter
|
|
self.sent_delimiter = sent_delimiter
|
|
self.max_seq_len = max_seq_len
|
|
super().__init__(data, transform, cache, generate_idx)
|
|
|
|
def load_file(self, filepath):
|
|
max_seq_len = self.max_seq_len
|
|
delimiter = self.sent_delimiter
|
|
for chars, tags in self._generate_chars_tags(filepath, delimiter, max_seq_len):
|
|
yield {'char': chars, 'tag': tags}
|
|
|
|
@staticmethod
|
|
def _generate_chars_tags(filepath, delimiter, max_seq_len):
|
|
filepath = get_resource(filepath)
|
|
with open(filepath, encoding='utf8') as src:
|
|
for text in src:
|
|
chars, tags = bmes_of(text, True)
|
|
if max_seq_len and delimiter and len(chars) > max_seq_len:
|
|
short_chars, short_tags = [], []
|
|
for idx, (char, tag) in enumerate(zip(chars, tags)):
|
|
short_chars.append(char)
|
|
short_tags.append(tag)
|
|
if len(short_chars) >= max_seq_len and delimiter(char):
|
|
yield short_chars, short_tags
|
|
short_chars, short_tags = [], []
|
|
if short_chars:
|
|
yield short_chars, short_tags
|
|
else:
|
|
yield chars, tags
|