# coding: utf-8
"""
Vocabulary module
Source: https://github.com/joeynmt/joeynmt/blob/main/joeynmt/vocabulary.py
"""
# Base Dependencies
# -----------------
import sys
import logging
import numpy as np
from collections import Counter
from pathlib import Path
from typing import Dict, List, Tuple, Optional
# Local Dependencies
# ------------------
from constants import (
BOS_ID,
BOS_TOKEN,
EOS_ID,
EOS_TOKEN,
PAD_ID,
PAD_TOKEN,
UNK_ID,
UNK_TOKEN,
)
from models.relation_collection import RelationCollection
from utils import read_list_from_file, write_list_to_file
# Constants
# ---------
from constants import DATASETS_PATHS, N2C2_VOCAB_PATH, DDI_VOCAB_PATH
VOC_MIN_FREQ = 10
logger = logging.getLogger(__name__)
class Vocabulary:
"""Vocabulary represents mapping between tokens and indices."""
def __init__(self, tokens: List[str]) -> None:
"""
Create vocabulary from list of tokens.
Special tokens are added if not already in list.
Args:
tokens (List[str]): list of tokens
"""
# warning: stoi grows with unknown tokens, don't use for saving or size
# special symbols
self.specials = [PAD_TOKEN, BOS_TOKEN, EOS_TOKEN, UNK_TOKEN]
# don't allow to access _stoi and _itos outside of this class
self._stoi: Dict[str, int] = {} # string to index
self._itos: List[str] = [] # index to string
# construct
self.add_tokens(tokens=self.specials + tokens)
assert len(self._stoi) == len(self._itos)
# assign after stoi is built
self.pad_index = self.lookup(PAD_TOKEN)
self.bos_index = self.lookup(BOS_TOKEN)
self.eos_index = self.lookup(EOS_TOKEN)
self.unk_index = self.lookup(UNK_TOKEN)
assert self.pad_index == PAD_ID
assert self.bos_index == BOS_ID
assert self.eos_index == EOS_ID
assert self.unk_index == UNK_ID
assert self._itos[UNK_ID] == UNK_TOKEN
def add_tokens(self, tokens: List[str]) -> None:
"""
Add list of tokens to vocabulary
Args:
tokens (List[str]): list of tokens to add to the vocabulary
"""
for t in tokens:
new_index = len(self._itos)
# add to vocab if not already there
if t not in self._itos:
self._itos.append(t)
self._stoi[t] = new_index
def to_file(self, file: Path) -> None:
"""
Save the vocabulary to a file, by writing token with index i in line i.
Args:
file (Path): path to file where the vocabulary is written
"""
write_list_to_file(file, self._itos)
def is_unk(self, token: str) -> bool:
"""
Check whether a token is covered by the vocabulary
Args:
token (str):
Returns:
bool: True if covered, False otherwise
"""
return self.lookup(token) == UNK_ID
def lookup(self, token: str) -> int:
"""
look up the encoding dictionary. (needed for multiprocessing)
Args:
token (str): surface str
Returns:
int: token id
"""
return self._stoi.get(token, UNK_ID)
def __len__(self) -> int:
return len(self._itos)
def __eq__(self, other) -> bool:
if isinstance(other, Vocabulary):
return self._itos == other._itos
return False
def array_to_sentence(
self, array: np.ndarray, cut_at_eos: bool = True, skip_pad: bool = True
) -> List[str]:
"""
Converts an array of IDs to a sentence, optionally cutting the result off at the
end-of-sequence token.
Args:
array (numpy.ndarray): 1D array containing indices
cut_at_eos (bool): cut the decoded sentences at the first <eos>
skip_pad (bool): skip generated <pad> tokens
Returns:
List[str]: list of strings (tokens)
"""
sentence = []
for i in array:
s = self._itos[i]
if skip_pad and s == PAD_TOKEN:
continue
sentence.append(s)
# break at the position AFTER eos
if cut_at_eos and s == EOS_TOKEN:
break
return sentence
def arrays_to_sentences(
self, arrays: np.ndarray, cut_at_eos: bool = True, skip_pad: bool = True
) -> List[List[str]]:
"""
Convert multiple arrays containing sequences of token IDs to their sentences,
optionally cutting them off at the end-of-sequence token.
Args:
arrays (numpy.ndarray): 2D array containing indices
cut_at_eos (bool): cut the decoded sentences at the first <eos>
skip_pad (bool): skip generated <pad> tokens
Returns:
List[List[str]]: list of list of strings (tokens)
"""
return [
self.array_to_sentence(
array=array, cut_at_eos=cut_at_eos, skip_pad=skip_pad
)
for array in arrays
]
def sentences_to_ids(
self,
sentences: List[List[str]],
padded: bool = False,
bos: bool = False,
eos: bool = False,
) -> Tuple[List[List[int]], List[int]]:
"""
Encode sentences to indices and pad sequences to the maximum length of the
sentences given if necessary
Args:
sentences List[List[str]]: list of tokenized sentences
Returns:
- padded ids
- original lengths before padding
"""
max_len = max([len(sent) for sent in sentences])
if bos:
max_len += 1
if eos:
max_len += 1
sentences_enc, lengths = [], []
for sent in sentences:
encoded = [self.lookup(s) for s in sent]
if bos:
encoded = [self.bos_index] + encoded
if eos:
encoded = encoded + [self.eos_index]
if padded:
offset = max(0, max_len - len(encoded))
sentences_enc.append(encoded + [self.pad_index] * offset)
else:
sentences_enc.append(encoded)
lengths.append(len(encoded))
return sentences_enc, lengths
def log_vocab(self, k: int) -> str:
"""first k vocab entities"""
return " ".join(f"({i}) {t}" for i, t in enumerate(self._itos[:k]))
def __repr__(self) -> str:
return (
f"{self.__class__.__name__}(len={self.__len__()}, "
f"specials={self.specials})"
)
@staticmethod
def sort_and_cut(
counter: Counter, max_size: int = sys.maxsize, min_freq: int = -1
) -> List[str]:
"""
Cut counter to most frequent, sorted numerically and alphabetically
Args:
counter (Counter): flattened token list in Counter object
max_size (int): maximum size of vocabulary
min_freq (int): minimum frequency for an item to be included
Returns:
List[str]: valid tokens
"""
# filter counter by min frequency
if min_freq > -1:
counter = Counter({t: c for t, c in counter.items() if c >= min_freq})
# sort by frequency, then alphabetically
tokens_and_frequencies = sorted(counter.items(), key=lambda tup: tup[0])
tokens_and_frequencies.sort(key=lambda tup: tup[1], reverse=True)
# cut off
vocab_tokens = [i[0] for i in tokens_and_frequencies[:max_size]]
assert len(vocab_tokens) <= max_size, (len(vocab_tokens), max_size)
return vocab_tokens
@staticmethod
def build_vocab(
cfg: Dict, collection: Optional[RelationCollection] = None
) -> "Vocabulary":
"""
Builds vocabulary either from file or sentences.
Args:
cfg (Dict): data cfg
Returns:
Vocabulary: created from either `tokens` or `vocab_file`
"""
vocab_file = cfg.get("voc_file", None)
min_freq = cfg.get("voc_min_freq", 1) # min freq for an item to be included
max_size = int(cfg.get("voc_limit", sys.maxsize)) # max size of vocabulary
assert max_size > 0
if vocab_file is not None:
# load it from file (not to apply `sort_and_cut()`)
unique_tokens = read_list_from_file(Path(vocab_file))
elif collection is not None:
# tokenize sentences
tokens = []
for doc in collection.tokens:
for t in doc:
tokens.append(t.text.lower())
# newly create unique token list (language-wise)
counter = Counter(tokens)
unique_tokens = Vocabulary.sort_and_cut(counter, max_size, min_freq)
else:
raise Exception("Please provide a vocab file path or a relation collection.")
vocab = Vocabulary(unique_tokens)
assert len(vocab) <= max_size + len(vocab.specials), (len(vocab), max_size)
# check for all except for UNK token whether they are OOVs
for s in vocab.specials:
assert s == UNK_TOKEN or not vocab.is_unk(s)
return vocab
@staticmethod
def create_vocabulary(dataset: str, train_collection: RelationCollection, save_to_disk: bool = True) -> "Vocabulary":
"""Creates the vocabulary of a dataset
Args:
dataset (str): dataset's name
train_collection (RelationCollection): train split of the dataset
Returns:
Vocabulary: _description_
"""
# configuration
cfg = {
"voc_min_freq": VOC_MIN_FREQ,
}
# create vocabulary
vocabulary = Vocabulary.build_vocab(cfg=cfg, collection=train_collection)
print(
"Vocabulary created for {} dataset: {} tokens".format(dataset, len(vocabulary))
)
# save vocab to file
if save_to_disk:
vocab_file = DATASETS_PATHS[dataset]
vocabulary.to_file(vocab_file)
return vocabulary
def load_vocab(dataset: str) -> "Vocabulary":
"""Loads the vocabulary of a dataset
Args:
dataset (str): dataset's name
Returns:
Vocabulary: vocabulary of the dataset
"""
path = {"n2c2": N2C2_VOCAB_PATH, "DDI": DDI_VOCAB_PATH}[dataset]
return Vocabulary(read_list_from_file(path))
