202 lines
8.1 KiB
Python
202 lines
8.1 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2019-10-24 15:07
|
|
import functools
|
|
from abc import ABC
|
|
from typing import Tuple, Union, List, Iterable
|
|
|
|
import tensorflow as tf
|
|
|
|
from hanlp.common.transform_tf import Transform
|
|
from hanlp.common.vocab_tf import VocabTF
|
|
from hanlp.utils.io_util import get_resource
|
|
from hanlp.utils.lang.zh.char_table import CharTable
|
|
from hanlp.utils.span_util import bmes_of, bmes_to_words
|
|
from hanlp.utils.string_util import split_long_sent
|
|
|
|
|
|
def generate_words_per_line(file_path):
|
|
with open(file_path, encoding='utf-8') as src:
|
|
for line in src:
|
|
cells = line.strip().split()
|
|
if not cells:
|
|
continue
|
|
yield cells
|
|
|
|
|
|
def words_to_bmes(words):
|
|
tags = []
|
|
for w in words:
|
|
if not w:
|
|
raise ValueError('{} contains None or zero-length word {}'.format(str(words), w))
|
|
if len(w) == 1:
|
|
tags.append('S')
|
|
else:
|
|
tags.extend(['B'] + ['M'] * (len(w) - 2) + ['E'])
|
|
return tags
|
|
|
|
|
|
def extract_ngram_features_and_tags(sentence, bigram_only=False, window_size=4, segmented=True):
|
|
"""
|
|
Feature extraction for windowed approaches
|
|
See Also https://github.com/chqiwang/convseg/
|
|
Parameters
|
|
----------
|
|
sentence
|
|
bigram_only
|
|
window_size
|
|
segmented
|
|
|
|
Returns
|
|
-------
|
|
|
|
"""
|
|
chars, tags = bmes_of(sentence, segmented)
|
|
chars = CharTable.normalize_chars(chars)
|
|
ret = []
|
|
ret.append(chars)
|
|
# TODO: optimize ngram generation using https://www.tensorflow.org/api_docs/python/tf/strings/ngrams
|
|
ret.extend(extract_ngram_features(chars, bigram_only, window_size))
|
|
ret.append(tags)
|
|
return tuple(ret[:-1]), ret[-1] # x, y
|
|
|
|
|
|
def extract_ngram_features(chars, bigram_only, window_size):
|
|
ret = []
|
|
if bigram_only:
|
|
chars = ['', ''] + chars + ['', '']
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[:-4], chars[1:])])
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[1:-3], chars[2:])])
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[2:-2], chars[3:])])
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[3:-1], chars[4:])])
|
|
elif window_size > 0:
|
|
chars = ['', '', ''] + chars + ['', '', '']
|
|
# single char
|
|
if window_size >= 1:
|
|
ret.append(chars[3:-3])
|
|
if window_size >= 2:
|
|
# bi chars
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[2:], chars[3:-3])])
|
|
ret.append([a + b if a and b else '' for a, b in zip(chars[3:-3], chars[4:])])
|
|
if window_size >= 3:
|
|
# tri chars
|
|
ret.append(
|
|
[a + b + c if a and b and c else '' for a, b, c in zip(chars[1:], chars[2:], chars[3:-3])])
|
|
ret.append(
|
|
[a + b + c if a and b and c else '' for a, b, c in zip(chars[2:], chars[3:-3], chars[4:])])
|
|
ret.append(
|
|
[a + b + c if a and b and c else '' for a, b, c in zip(chars[3:-3], chars[4:], chars[5:])])
|
|
if window_size >= 4:
|
|
# four chars
|
|
ret.append([a + b + c + d if a and b and c and d else '' for a, b, c, d in
|
|
zip(chars[0:], chars[1:], chars[2:], chars[3:-3])])
|
|
ret.append([a + b + c + d if a and b and c and d else '' for a, b, c, d in
|
|
zip(chars[1:], chars[2:], chars[3:-3], chars[4:])])
|
|
ret.append([a + b + c + d if a and b and c and d else '' for a, b, c, d in
|
|
zip(chars[2:], chars[3:-3], chars[4:], chars[5:])])
|
|
ret.append([a + b + c + d if a and b and c and d else '' for a, b, c, d in
|
|
zip(chars[3:-3], chars[4:], chars[5:], chars[6:])])
|
|
return ret
|
|
|
|
|
|
def generate_ngram_bmes(file_path, bigram_only=False, window_size=4, gold=True):
|
|
with open(file_path, encoding='utf-8') as src:
|
|
for line in src:
|
|
sentence = line.strip()
|
|
if not sentence:
|
|
continue
|
|
yield extract_ngram_features_and_tags(sentence, bigram_only, window_size, gold)
|
|
|
|
|
|
def vocab_from_txt(txt_file_path, bigram_only=False, window_size=4, **kwargs) -> Tuple[VocabTF, VocabTF, VocabTF]:
|
|
char_vocab, ngram_vocab, tag_vocab = VocabTF(), VocabTF(), VocabTF(pad_token=None, unk_token=None)
|
|
for X, Y in generate_ngram_bmes(txt_file_path, bigram_only, window_size, gold=True):
|
|
char_vocab.update(X[0])
|
|
for ngram in X[1:]:
|
|
ngram_vocab.update(filter(lambda x: x, ngram))
|
|
tag_vocab.update(Y)
|
|
return char_vocab, ngram_vocab, tag_vocab
|
|
|
|
|
|
def dataset_from_txt(txt_file_path: str, char_vocab: VocabTF, ngram_vocab: VocabTF, tag_vocab: VocabTF,
|
|
bigram_only=False,
|
|
window_size=4, segmented=True, batch_size=32, shuffle=None, repeat=None, prefetch=1):
|
|
generator = functools.partial(generate_ngram_bmes, txt_file_path, bigram_only, window_size, segmented)
|
|
return dataset_from_generator(generator, char_vocab, ngram_vocab, tag_vocab, bigram_only, window_size, batch_size,
|
|
shuffle, repeat, prefetch)
|
|
|
|
|
|
def dataset_from_generator(generator, char_vocab, ngram_vocab, tag_vocab, bigram_only=False, window_size=4,
|
|
batch_size=32, shuffle=None, repeat=None, prefetch=1):
|
|
if bigram_only:
|
|
ngram_size = 4
|
|
else:
|
|
ngram_size = window_size * (window_size + 1) // 2
|
|
vec_dim = 2 + ngram_size
|
|
shapes = tuple([[None]] * (vec_dim - 1)), [None]
|
|
types = tuple([tf.string] * (vec_dim - 1)), tf.string
|
|
defaults = tuple([char_vocab.pad_token] + [
|
|
ngram_vocab.pad_token if ngram_vocab else char_vocab.pad_token] * ngram_size), (
|
|
tag_vocab.pad_token if tag_vocab.pad_token else tag_vocab.first_token)
|
|
dataset = tf.data.Dataset.from_generator(generator, output_shapes=shapes, output_types=types)
|
|
if shuffle:
|
|
if isinstance(shuffle, bool):
|
|
shuffle = 1024
|
|
dataset = dataset.shuffle(shuffle)
|
|
if repeat:
|
|
dataset = dataset.repeat(repeat)
|
|
dataset = dataset.padded_batch(batch_size, shapes, defaults).prefetch(prefetch)
|
|
return dataset
|
|
|
|
|
|
class TxtFormat(Transform, ABC):
|
|
def file_to_inputs(self, filepath: str, gold=True):
|
|
filepath = get_resource(filepath)
|
|
with open(filepath, encoding='utf-8') as src:
|
|
for line in src:
|
|
sentence = line.strip()
|
|
if not sentence:
|
|
continue
|
|
yield sentence
|
|
|
|
|
|
class TxtBMESFormat(TxtFormat, ABC):
|
|
def file_to_inputs(self, filepath: str, gold=True):
|
|
max_seq_length = self.config.get('max_seq_length', False)
|
|
if max_seq_length:
|
|
if 'transformer' in self.config:
|
|
max_seq_length -= 2 # allow for [CLS] and [SEP]
|
|
delimiter = set()
|
|
delimiter.update('。!?:;、,,;!?、,')
|
|
for text in super().file_to_inputs(filepath, gold):
|
|
chars, tags = bmes_of(text, gold)
|
|
if max_seq_length:
|
|
start = 0
|
|
for short_chars in split_long_sent(chars, delimiter, max_seq_length):
|
|
end = start + len(short_chars)
|
|
yield short_chars, tags[start:end]
|
|
start = end
|
|
else:
|
|
yield chars, tags
|
|
|
|
def input_is_single_sample(self, input: Union[List[str], List[List[str]]]) -> bool:
|
|
return isinstance(input, str)
|
|
|
|
def inputs_to_samples(self, inputs, gold=False):
|
|
for chars, tags in (inputs if gold else zip(inputs, [None] * len(inputs))):
|
|
if not gold:
|
|
tags = [self.tag_vocab.safe_pad_token] * len(chars)
|
|
chars = CharTable.normalize_chars(chars)
|
|
yield chars, tags
|
|
|
|
def Y_to_outputs(self, Y: Union[tf.Tensor, Tuple[tf.Tensor]], gold=False, inputs=None, X=None,
|
|
batch=None) -> Iterable:
|
|
yield from self.Y_to_tokens(self.tag_vocab, Y, gold, inputs)
|
|
|
|
def Y_to_tokens(self, tag_vocab, Y, gold, inputs):
|
|
if not gold:
|
|
Y = tf.argmax(Y, axis=2)
|
|
for text, ys in zip(inputs, Y):
|
|
tags = [tag_vocab.idx_to_token[int(y)] for y in ys[:len(text)]]
|
|
yield bmes_to_words(list(text), tags)
|