hankcs--hanlp
wehub-resource-sync
507 行
16 KiB
Python
507 行
16 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2019-06-13 22:42
|
|
from collections import Counter
|
|
from typing import List, Dict, Union, Iterable
|
|
|
|
from hanlp_common.constant import UNK, PAD
|
|
from hanlp_common.structure import Serializable
|
|
from hanlp_common.reflection import classpath_of
|
|
|
|
|
|
class Vocab(Serializable):
|
|
def __init__(self, idx_to_token: List[str] = None, token_to_idx: Dict = None, mutable=True, pad_token=PAD,
|
|
unk_token=UNK) -> None:
|
|
"""Vocabulary base class which converts tokens to indices and vice versa.
|
|
|
|
Args:
|
|
idx_to_token: id to token mapping.
|
|
token_to_idx: token to id mapping.
|
|
mutable: ``True`` to allow adding new tokens, ``False`` to map OOV to ``unk``.
|
|
pad_token: The token representing padding.
|
|
unk_token: The token representing OOV.
|
|
"""
|
|
super().__init__()
|
|
if idx_to_token:
|
|
t2i = dict((token, idx) for idx, token in enumerate(idx_to_token))
|
|
if token_to_idx:
|
|
t2i.update(token_to_idx)
|
|
token_to_idx = t2i
|
|
if token_to_idx is None:
|
|
token_to_idx = {}
|
|
if pad_token is not None:
|
|
token_to_idx[pad_token] = len(token_to_idx)
|
|
if unk_token is not None:
|
|
token_to_idx[unk_token] = token_to_idx.get(unk_token, len(token_to_idx))
|
|
self.token_to_idx = token_to_idx
|
|
self.idx_to_token: List[str] = None
|
|
self.mutable = mutable
|
|
self.pad_token = pad_token
|
|
self.unk_token = unk_token
|
|
|
|
def __setitem__(self, token: str, idx: int):
|
|
assert self.mutable, 'Update an immutable Vocab object is not allowed'
|
|
self.token_to_idx[token] = idx
|
|
|
|
def __getitem__(self, key: Union[str, int, List]) -> Union[int, str, List]:
|
|
""" Get the index/indices associated with a token or a list of tokens or vice versa.
|
|
|
|
Args:
|
|
key: ``str`` for token(s) and ``int`` for index/indices.
|
|
|
|
Returns: Associated indices or tokens.
|
|
|
|
"""
|
|
if isinstance(key, str):
|
|
return self.get_idx(key)
|
|
elif isinstance(key, int):
|
|
return self.get_token(key)
|
|
elif isinstance(key, list):
|
|
if len(key) == 0:
|
|
return []
|
|
elif isinstance(key[0], str):
|
|
return [self.get_idx(x) for x in key]
|
|
elif isinstance(key[0], int):
|
|
return [self.get_token(x) for x in key]
|
|
|
|
def __contains__(self, key: Union[str, int]):
|
|
if isinstance(key, str):
|
|
return key in self.token_to_idx
|
|
elif isinstance(key, int):
|
|
return 0 <= key < len(self.idx_to_token)
|
|
else:
|
|
return False
|
|
|
|
def add(self, token: str) -> int:
|
|
""" Tries to add a token into a vocab and returns its id. If it has already been there, its id will be returned
|
|
and the vocab won't be updated. If the vocab is locked, an assertion failure will occur.
|
|
|
|
Args:
|
|
token: A new or existing token.
|
|
|
|
Returns:
|
|
Its associated id.
|
|
|
|
"""
|
|
assert self.mutable, 'It is not allowed to call add on an immutable Vocab'
|
|
assert isinstance(token, str), f'Token type must be str but got {type(token)} from {token}'
|
|
assert token is not None, 'Token must not be None'
|
|
idx = self.token_to_idx.get(token, None)
|
|
if idx is None:
|
|
idx = len(self.token_to_idx)
|
|
self.token_to_idx[token] = idx
|
|
return idx
|
|
|
|
def update(self, tokens: Iterable[str]) -> None:
|
|
"""Update the vocab with these tokens by adding them to vocab one by one.
|
|
|
|
Args:
|
|
tokens (Iterable[str]): A list of tokens.
|
|
"""
|
|
assert self.mutable, 'It is not allowed to update an immutable Vocab'
|
|
for token in tokens:
|
|
self.add(token)
|
|
|
|
def get_idx(self, token: str) -> int:
|
|
"""Get the idx of a token. If it's not there, it will be added to the vocab when the vocab is locked otherwise
|
|
the id of UNK will be returned.
|
|
|
|
Args:
|
|
token: A token.
|
|
|
|
Returns:
|
|
The id of that token.
|
|
|
|
"""
|
|
assert isinstance(token, str), 'token has to be `str`'
|
|
idx = self.token_to_idx.get(token, None)
|
|
if idx is None:
|
|
if self.mutable:
|
|
idx = len(self.token_to_idx)
|
|
self.token_to_idx[token] = idx
|
|
else:
|
|
idx = self.token_to_idx.get(self.unk_token, None)
|
|
return idx
|
|
|
|
def get_idx_without_add(self, token: str) -> int:
|
|
idx = self.token_to_idx.get(token, None)
|
|
if idx is None:
|
|
idx = self.token_to_idx.get(self.safe_unk_token, None)
|
|
return idx
|
|
|
|
def get_token(self, idx: int) -> str:
|
|
"""Get the token using its index.
|
|
|
|
Args:
|
|
idx: The index to a token.
|
|
|
|
Returns:
|
|
|
|
"""
|
|
if self.idx_to_token:
|
|
return self.idx_to_token[idx]
|
|
|
|
if self.mutable:
|
|
for token in self.token_to_idx:
|
|
if self.token_to_idx[token] == idx:
|
|
return token
|
|
|
|
def has_key(self, token):
|
|
return token in self.token_to_idx
|
|
|
|
def __len__(self):
|
|
return len(self.token_to_idx)
|
|
|
|
def lock(self):
|
|
"""Lock this vocab up so that it won't accept new tokens.
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
if self.locked:
|
|
return self
|
|
self.mutable = False
|
|
self.build_idx_to_token()
|
|
return self
|
|
|
|
def build_idx_to_token(self):
|
|
max_idx = max(self.token_to_idx.values())
|
|
self.idx_to_token = [None] * (max_idx + 1)
|
|
for token, idx in self.token_to_idx.items():
|
|
self.idx_to_token[idx] = token
|
|
|
|
def unlock(self):
|
|
"""Unlock this vocab so that new tokens can be added in.
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
if not self.locked:
|
|
return
|
|
self.mutable = True
|
|
self.idx_to_token = None
|
|
return self
|
|
|
|
@property
|
|
def locked(self):
|
|
"""
|
|
``True`` indicates this vocab is locked.
|
|
"""
|
|
return not self.mutable
|
|
|
|
@property
|
|
def unk_idx(self):
|
|
"""
|
|
The index of ``UNK`` token.
|
|
"""
|
|
if self.unk_token is None:
|
|
return None
|
|
else:
|
|
return self.token_to_idx.get(self.unk_token, None)
|
|
|
|
@property
|
|
def pad_idx(self):
|
|
"""
|
|
The index of ``PAD`` token.
|
|
"""
|
|
if self.pad_token is None:
|
|
return None
|
|
else:
|
|
return self.token_to_idx.get(self.pad_token, None)
|
|
|
|
@property
|
|
def tokens(self):
|
|
"""
|
|
A set of all tokens in this vocab.
|
|
"""
|
|
return self.token_to_idx.keys()
|
|
|
|
def __str__(self) -> str:
|
|
return self.token_to_idx.__str__()
|
|
|
|
def summary(self, verbose=True) -> str:
|
|
"""Get or print a summary of this vocab.
|
|
|
|
Args:
|
|
verbose: ``True`` to print the summary to stdout.
|
|
|
|
Returns:
|
|
Summary in text form.
|
|
|
|
"""
|
|
# report = 'Length: {}\n'.format(len(self))
|
|
# report += 'Samples: {}\n'.format(str(list(self.token_to_idx.keys())[:min(50, len(self))]))
|
|
# report += 'Mutable: {}'.format(self.mutable)
|
|
# report = report.strip()
|
|
report = '[{}] = '.format(len(self))
|
|
report += str(list(self.token_to_idx.keys())[:min(50, len(self))])
|
|
if verbose:
|
|
print(report)
|
|
return report
|
|
|
|
def __call__(self, some_token: Union[str, Iterable[str]]) -> Union[int, List[int]]:
|
|
if isinstance(some_token, (list, tuple, set)):
|
|
indices = []
|
|
if len(some_token) and isinstance(some_token[0], (list, tuple, set)):
|
|
for sent in some_token:
|
|
inside = []
|
|
for token in sent:
|
|
inside.append(self.get_idx(token))
|
|
indices.append(inside)
|
|
return indices
|
|
for token in some_token:
|
|
indices.append(self.get_idx(token))
|
|
return indices
|
|
else:
|
|
return self.get_idx(some_token)
|
|
|
|
def to_dict(self) -> dict:
|
|
"""Convert this vocab to a dict so that it can be json serialized.
|
|
|
|
Returns:
|
|
A dict.
|
|
|
|
"""
|
|
idx_to_token = self.idx_to_token
|
|
pad_token = self.pad_token
|
|
unk_token = self.unk_token
|
|
mutable = self.mutable
|
|
items = locals().copy()
|
|
items.pop('self')
|
|
return items
|
|
|
|
def copy_from(self, item: dict):
|
|
"""Copy properties from a dict so that it can json de-serialized.
|
|
|
|
Args:
|
|
item: A dict holding ``token_to_idx``
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
for key, value in item.items():
|
|
setattr(self, key, value)
|
|
self.token_to_idx = {k: v for v, k in enumerate(self.idx_to_token)}
|
|
return self
|
|
|
|
def lower(self):
|
|
"""Convert all tokens to lower case.
|
|
|
|
Returns:
|
|
Itself.
|
|
|
|
"""
|
|
self.unlock()
|
|
token_to_idx = self.token_to_idx
|
|
self.token_to_idx = {}
|
|
for token in token_to_idx.keys():
|
|
self.add(token.lower())
|
|
return self
|
|
|
|
@property
|
|
def first_token(self):
|
|
"""The first token in this vocab.
|
|
"""
|
|
if self.idx_to_token:
|
|
return self.idx_to_token[0]
|
|
if self.token_to_idx:
|
|
return next(iter(self.token_to_idx))
|
|
return None
|
|
|
|
def merge(self, other):
|
|
"""Merge this with another vocab inplace.
|
|
|
|
Args:
|
|
other (Vocab): Another vocab.
|
|
"""
|
|
for word, idx in other.token_to_idx.items():
|
|
self.get_idx(word)
|
|
|
|
@property
|
|
def safe_pad_token(self) -> str:
|
|
"""Get the pad token safely. It always returns a pad token, which is the pad token or the first token
|
|
if pad does not present in the vocab.
|
|
"""
|
|
if self.pad_token:
|
|
return self.pad_token
|
|
if self.first_token:
|
|
return self.first_token
|
|
return PAD
|
|
|
|
@property
|
|
def safe_pad_token_idx(self) -> int:
|
|
"""Get the idx to the pad token safely. It always returns an index, which corresponds to the pad token or the
|
|
first token if pad does not present in the vocab.
|
|
"""
|
|
return self.token_to_idx.get(self.safe_pad_token, 0)
|
|
|
|
@property
|
|
def safe_unk_token(self) -> str:
|
|
"""Get the unk token safely. It always returns a unk token, which is the unk token or the first token if unk
|
|
does not presented in the vocab.
|
|
"""
|
|
if self.unk_token:
|
|
return self.unk_token
|
|
if self.first_token:
|
|
return self.first_token
|
|
return UNK
|
|
|
|
def __repr__(self) -> str:
|
|
if self.idx_to_token is not None:
|
|
return self.idx_to_token.__repr__()
|
|
return self.token_to_idx.__repr__()
|
|
|
|
def extend(self, tokens: Iterable[str]):
|
|
self.unlock()
|
|
self(tokens)
|
|
|
|
def reload_idx_to_token(self, idx_to_token: List[str], pad_idx=0, unk_idx=1):
|
|
self.idx_to_token = idx_to_token
|
|
self.token_to_idx = dict((s, i) for i, s in enumerate(idx_to_token))
|
|
if pad_idx is not None:
|
|
self.pad_token = idx_to_token[pad_idx]
|
|
if unk_idx is not None:
|
|
self.unk_token = idx_to_token[unk_idx]
|
|
|
|
def set_unk_as_safe_unk(self):
|
|
"""Set ``self.unk_token = self.safe_unk_token``. It's useful when the dev/test set contains OOV labels.
|
|
"""
|
|
self.unk_token = self.safe_unk_token
|
|
|
|
def clear(self):
|
|
self.unlock()
|
|
self.token_to_idx.clear()
|
|
|
|
|
|
class CustomVocab(Vocab):
|
|
def to_dict(self) -> dict:
|
|
d = super().to_dict()
|
|
d['type'] = classpath_of(self)
|
|
return d
|
|
|
|
|
|
class LowercaseVocab(CustomVocab):
|
|
def get_idx(self, token: str) -> int:
|
|
idx = self.token_to_idx.get(token, None)
|
|
if idx is None:
|
|
idx = self.token_to_idx.get(token.lower(), None)
|
|
if idx is None:
|
|
if self.mutable:
|
|
idx = len(self.token_to_idx)
|
|
self.token_to_idx[token] = idx
|
|
else:
|
|
idx = self.token_to_idx.get(self.unk_token, None)
|
|
return idx
|
|
|
|
|
|
class VocabWithNone(CustomVocab):
|
|
def get_idx(self, token: str) -> int:
|
|
if token is None:
|
|
return -1
|
|
return super().get_idx(token)
|
|
|
|
|
|
class VocabWithFrequency(CustomVocab):
|
|
|
|
def __init__(self, counter: Counter = None, min_occur_cnt=0, pad_token=PAD, unk_token=UNK, specials=None) -> None:
|
|
super().__init__(None, None, True, pad_token, unk_token)
|
|
if specials:
|
|
for each in specials:
|
|
counter.pop(each, None)
|
|
self.add(each)
|
|
self.frequencies = [1] * len(self)
|
|
if counter:
|
|
for token, freq in counter.most_common():
|
|
if freq >= min_occur_cnt:
|
|
self.add(token)
|
|
self.frequencies.append(freq)
|
|
self.lock()
|
|
|
|
def to_dict(self) -> dict:
|
|
d = super().to_dict()
|
|
d['frequencies'] = self.frequencies
|
|
return d
|
|
|
|
def copy_from(self, item: dict):
|
|
super().copy_from(item)
|
|
self.frequencies = item['frequencies']
|
|
|
|
def get_frequency(self, token):
|
|
idx = self.get_idx(token)
|
|
if idx is not None:
|
|
return self.frequencies[idx]
|
|
return 0
|
|
|
|
|
|
class VocabCounter(CustomVocab):
|
|
|
|
def __init__(self, idx_to_token: List[str] = None, token_to_idx: Dict = None, mutable=True, pad_token=PAD,
|
|
unk_token=UNK) -> None:
|
|
super().__init__(idx_to_token, token_to_idx, mutable, pad_token, unk_token)
|
|
self.counter = Counter()
|
|
|
|
def get_idx(self, token: str) -> int:
|
|
if self.mutable:
|
|
self.counter[token] += 1
|
|
return super().get_idx(token)
|
|
|
|
def trim(self, min_frequency):
|
|
assert self.mutable
|
|
specials = {self.unk_token, self.pad_token}
|
|
survivors = list((token, freq) for token, freq in self.counter.most_common()
|
|
if freq >= min_frequency and token not in specials)
|
|
survivors = [(x, -1) for x in specials if x] + survivors
|
|
self.counter = Counter(dict(survivors))
|
|
self.token_to_idx = dict()
|
|
self.idx_to_token = None
|
|
for token, freq in survivors:
|
|
idx = len(self.token_to_idx)
|
|
self.token_to_idx[token] = idx
|
|
|
|
def copy_from(self, item: dict):
|
|
super().copy_from(item)
|
|
self.counter = Counter(item['counter'].items()) if 'counter' in item else Counter()
|
|
|
|
def to_dict(self) -> dict:
|
|
d = super().to_dict()
|
|
d['counter'] = dict(self.counter.items())
|
|
return d
|
|
|
|
|
|
class Vocab3D(CustomVocab):
|
|
def __call__(self, some_token: Union[str, Iterable[str], Iterable[Iterable[str]]]) \
|
|
-> Union[int, List[int], List[List[int]]]:
|
|
"""It supports 3D arrays of tokens.
|
|
|
|
Args:
|
|
some_token: Tokens of 1D to 3D
|
|
|
|
Returns:
|
|
A list of indices.
|
|
|
|
"""
|
|
if isinstance(some_token, (list, tuple, set)):
|
|
indices = []
|
|
if len(some_token) and isinstance(some_token[0], (list, tuple, set)):
|
|
for sent in some_token:
|
|
inside = []
|
|
for token in sent:
|
|
inside.append(self.get_idx(token))
|
|
indices.append(inside)
|
|
return indices
|
|
for token in some_token:
|
|
if isinstance(token, str):
|
|
indices.append(self.get_idx(token))
|
|
else:
|
|
indices.append([self.get_idx(x) for x in token])
|
|
return indices
|
|
else:
|
|
return self.get_idx(some_token)
|
|
|
|
|
|
def create_label_vocab() -> Vocab:
|
|
return Vocab(pad_token=None, unk_token=None)
|