| from bidict import bidict |
| import pickle |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class Vocabulary(): |
| """This class maps strings to integers, which also allow many namespaces |
| """ |
|
|
| DEFAULT_PAD_TOKEN = '*@PAD@*' |
| DEFAULT_UNK_TOKEN = '*@UNK@*' |
|
|
| def __init__(self, |
| counters=dict(), |
| min_count=dict(), |
| pretrained_vocab=dict(), |
| intersection_namespace=dict(), |
| no_pad_namespace=list(), |
| no_unk_namespace=list(), |
| contain_pad_namespace=dict(), |
| contain_unk_namespace=dict()): |
| """initialize vocabulary |
| |
| Keyword Arguments: |
| counters {dict} -- multiple counter (default: {dict()}) |
| min_count {dict} -- min count dict (default: {dict()}) |
| pretrained_vocab {dict} -- pretrained vocabulary (default: {dict()}) |
| intersection_namespace {dict} -- intersection namespace correspond to pretrained vocabulary in case of too large pretrained vocabulary (default: {dict()}) |
| no_pad_namespace {list} -- no paddding namespace (default: {list()}) |
| no_unk_namespace {list} -- no unknown namespace (default: {list()}) |
| contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()}) |
| contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()}) |
| """ |
|
|
| self.min_count = dict(min_count) |
| self.intersection_namespace = dict(intersection_namespace) |
| self.no_pad_namespace = set(no_pad_namespace) |
| self.no_unk_namespace = set(no_unk_namespace) |
| self.contain_pad_namespace = dict(contain_pad_namespace) |
| self.contain_unk_namespace = dict(contain_unk_namespace) |
| self.vocab = dict() |
|
|
| self.extend_from_counter(counters, self.min_count, self.no_pad_namespace, |
| self.no_unk_namespace) |
|
|
| self.extend_from_pretrained_vocab(pretrained_vocab, self.intersection_namespace, |
| self.no_pad_namespace, self.no_unk_namespace) |
|
|
| logger.info("Initialize vocabulary successfully.") |
|
|
| def extend_from_pretrained_vocab(self, |
| pretrained_vocab, |
| intersection_namespace=dict(), |
| no_pad_namespace=list(), |
| no_unk_namespace=list(), |
| contain_pad_namespace=dict(), |
| contain_unk_namespace=dict()): |
| """extend vocabulary from pretrained vocab |
| |
| Arguments: |
| pretrained_vocab {dict} -- pretrained vocabulary |
| |
| Keyword Arguments: |
| intersection_namespace {dict} -- intersection namespace correspond to pretrained vocabulary in case of too large pretrained vocabulary (default: {dict()}) |
| no_pad_namespace {list} -- no paddding namespace (default: {list()}) |
| no_unk_namespace {list} -- no unknown namespace (default: {list()}) |
| contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()}) |
| contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()}) |
| """ |
|
|
| self.intersection_namespace.update(dict(intersection_namespace)) |
| self.no_pad_namespace.update(set(no_pad_namespace)) |
| self.no_unk_namespace.update(set(no_unk_namespace)) |
| self.contain_pad_namespace.update(dict(contain_pad_namespace)) |
| self.contain_unk_namespace.update(dict(contain_unk_namespace)) |
|
|
| for namespace, vocab in pretrained_vocab.items(): |
| self.__namespace_init(namespace) |
| is_intersection = namespace in self.intersection_namespace |
| intersection_vocab = self.vocab[ |
| self.intersection_namespace[namespace]] if is_intersection else [] |
| for key, value in vocab.items(): |
| if not is_intersection or key in intersection_vocab: |
| self.vocab[namespace][key] = value |
|
|
| logger.info( |
| "Vocabulay {} (size: {}) was constructed successfully from pretrained_vocab.". |
| format(namespace, len(self.vocab[namespace]))) |
|
|
| def extend_from_counter(self, |
| counters, |
| min_count=dict(), |
| no_pad_namespace=list(), |
| no_unk_namespace=list(), |
| contain_pad_namespace=dict(), |
| contain_unk_namespace=dict()): |
| """extend vocabulary from counter |
| |
| Arguments: |
| counters {dict} -- multiply counter |
| |
| Keyword Arguments: |
| min_count {dict} -- min count dict (default: {dict()}) |
| no_pad_namespace {list} -- no paddding namespace (default: {list()}) |
| no_unk_namespace {list} -- no unknown namespace (default: {list()}) |
| contain_pad_namespace {dict} -- contain padding token namespace (default: {dict()}) |
| contain_unk_namespace {dict} -- contain unknown token namespace (default: {dict()}) |
| """ |
|
|
| self.no_pad_namespace.update(set(no_pad_namespace)) |
| self.no_unk_namespace.update(set(no_unk_namespace)) |
| self.contain_pad_namespace.update(dict(contain_pad_namespace)) |
| self.contain_unk_namespace.update(dict(contain_unk_namespace)) |
| self.min_count.update(dict(min_count)) |
|
|
| for namespace, counter in counters.items(): |
| self.__namespace_init(namespace) |
| for key in counter: |
| minc = min_count[namespace] \ |
| if min_count and namespace in min_count else 1 |
| if counter[key] >= minc: |
| self.vocab[namespace][key] = len(self.vocab[namespace]) |
|
|
| logger.info("Vocabulay {} (size: {}) was constructed successfully from counter.".format( |
| namespace, len(self.vocab[namespace]))) |
|
|
| def add_tokens_to_namespace(self, tokens, namespace): |
| """This function adds tokens to one namespace for extending vocabulary |
| |
| Arguments: |
| tokens {list} -- token list |
| namespace {str} -- namespace name |
| """ |
|
|
| if namespace not in self.vocab: |
| self.__namespace_init(namespace) |
| logger.error('Add Namespace {} into vocabulary.'.format(namespace)) |
|
|
| for token in tokens: |
| if token not in self.vocab[namespace]: |
| self.vocab[namespace][token] = len(self.vocab[namespace]) |
|
|
| def get_token_index(self, token, namespace): |
| """This function gets token index in one namespace of vocabulary |
| |
| Arguments: |
| token {str} -- token |
| namespace {str} -- namespace name |
| |
| Raises: |
| RuntimeError: namespace not exists |
| |
| Returns: |
| int -- token index |
| """ |
|
|
| if token in self.vocab[namespace]: |
| return self.vocab[namespace][token] |
|
|
| if namespace not in self.no_unk_namespace: |
| return self.get_unknown_index(namespace) |
|
|
| logger.error("Can not find the index of {} from a no unknown token namespace {}.".format( |
| token, namespace)) |
| raise RuntimeError( |
| "Can not find the index of {} from a no unknown token namespace {}.".format( |
| token, namespace)) |
|
|
| def get_token_from_index(self, index, namespace): |
| """This function gets token using index in vocabulary |
| |
| Arguments: |
| index {int} -- index |
| namespace {str} -- namespace name |
| |
| Raises: |
| RuntimeError: index out of range |
| |
| Returns: |
| str -- token |
| """ |
|
|
| if index < len(self.vocab[namespace]): |
| return self.vocab[namespace].inv[index] |
|
|
| logger.error("The index {} is out of vocabulary {} range.".format(index, namespace)) |
| raise RuntimeError("The index {} is out of vocabulary {} range.".format(index, namespace)) |
|
|
| def get_vocab_size(self, namespace): |
| """This function gets the size of one namespace in vocabulary |
| |
| Arguments: |
| namespace {str} -- namespace name |
| |
| Returns: |
| int -- vocabulary size |
| """ |
|
|
| return len(self.vocab[namespace]) |
|
|
| def get_all_namespaces(self): |
| """This function gets all namespaces |
| |
| Returns: |
| list -- all namespaces vocabulary contained |
| """ |
|
|
| return set(self.vocab) |
|
|
| def get_padding_index(self, namespace): |
| """This function gets padding token index in one namespace of vocabulary |
| |
| Arguments: |
| namespace {str} -- namespace name |
| |
| Raises: |
| RuntimeError: no padding |
| |
| Returns: |
| int -- padding index |
| """ |
|
|
| if namespace not in self.vocab: |
| raise RuntimeError("Namespace {} doesn't exist.".format(namespace)) |
|
|
| if namespace not in self.no_pad_namespace: |
| if namespace not in self.contain_pad_namespace: |
| return self.vocab[namespace][Vocabulary.DEFAULT_PAD_TOKEN] |
| return self.vocab[namespace][self.contain_pad_namespace[namespace]] |
|
|
| logger.error("Namespace {} doesn't has paddding token.".format(namespace)) |
| raise RuntimeError("Namespace {} doesn't has paddding token.".format(namespace)) |
|
|
| def get_unknown_index(self, namespace): |
| """This function gets unknown token index in one namespace of vocabulary |
| |
| Arguments: |
| namespace {str} -- namespace name |
| |
| Raises: |
| RuntimeError: no unknown |
| |
| Returns: |
| int -- unknown index |
| """ |
|
|
| if namespace not in self.vocab: |
| raise RuntimeError("Namespace {} doesn't exist.".format(namespace)) |
|
|
| if namespace not in self.no_unk_namespace: |
| if namespace not in self.contain_unk_namespace: |
| return self.vocab[namespace][Vocabulary.DEFAULT_UNK_TOKEN] |
| return self.vocab[namespace][self.contain_unk_namespace[namespace]] |
|
|
| logger.error("Namespace {} doesn't has unknown token.".format(namespace)) |
| raise RuntimeError("Namespace {} doesn't has unknown token.".format(namespace)) |
|
|
| def get_namespace_tokens(self, namesapce): |
| """This function returns all tokens in one namespace |
| |
| Arguments: |
| namesapce {str} -- namespce name |
| |
| Returns: |
| dict_keys -- all tokens |
| """ |
|
|
| return self.vocab[namesapce] |
|
|
| def save(self, file_path): |
| """This function saves vocabulary into file |
| |
| Arguments: |
| file_path {str} -- file path |
| """ |
|
|
| pickle.dump(self, open(file_path, 'wb')) |
|
|
| @classmethod |
| def load(cls, file_path): |
| """This function loads vocabulary from file |
| |
| Arguments: |
| file_path {str} -- file path |
| |
| Returns: |
| Vocabulary -- vocabulary |
| """ |
|
|
| return pickle.load(open(file_path, 'rb'), encoding='utf-8') |
|
|
| def __namespace_init(self, namespace): |
| """This function initializes a namespace, |
| adds pad and unk token to one namespace of vacabulary |
| |
| Arguments: |
| namespace {str} -- namespace |
| """ |
|
|
| self.vocab[namespace] = bidict() |
|
|
| if namespace not in self.no_pad_namespace and namespace not in self.contain_pad_namespace: |
| self.vocab[namespace][Vocabulary.DEFAULT_PAD_TOKEN] = len(self.vocab[namespace]) |
|
|
| if namespace not in self.no_unk_namespace and namespace not in self.contain_unk_namespace: |
| self.vocab[namespace][Vocabulary.DEFAULT_UNK_TOKEN] = len(self.vocab[namespace]) |
|
|
