# Path: sockeye/vocab.py
# def build_vocab(data: Iterable[str], num_words: Optional[int] = None, min_count: int = 1,
#                 pad_to_multiple_of: Optional[int] = None) -> Vocab:
#     """
#     Creates a vocabulary mapping from words to ids. Increasing integer ids are assigned by word frequency,
#     using lexical sorting as a tie breaker. The only exception to this are special symbols such as the padding symbol
#     (PAD).
# 
#     :param data: Sequence of sentences containing whitespace-delimited tokens.
#     :param num_words: Optional maximum number of words in the vocabulary.
#     :param min_count: Minimum occurrences of words to be included in the vocabulary.
#     :param pad_to_multiple_of: If not None, pads the vocabulary to a size that is the next multiple of this int.
#     :return: Word-to-id mapping.
#     """
#     raw_vocab = count_tokens(data)
#     return build_pruned_vocab(raw_vocab=raw_vocab,
#                               num_words=num_words,
#                               min_count=min_count,
#                               pad_to_multiple_of=pad_to_multiple_of)
# 
# def get_ordered_tokens_from_vocab(vocab: Vocab) -> List[str]:
#     """
#     Returns the list of tokens in a vocabulary, ordered by increasing vocabulary id.
# 
#     :param vocab: Input vocabulary.
#     :return: List of tokens.
#     """
#     return [token for token, token_id in sorted(vocab.items(), key=lambda i: i[1])]
# 
# def is_valid_vocab(vocab: Vocab) -> bool:
#     """
#     Checks if a vocabulary is valid. We define valid as:
#     1. All indices from 0 to num_words - 1 are present without duplicates.
#     2. PAD_SYMBOL has word id 0, UNK_SYMBOL has word id 1, BOS_SYMBOL has word id 2, EOS_SYMBOL has word id 3.
#     """
#     if vocab[C.PAD_SYMBOL] != C.PAD_ID:
#         logger.warning("PAD_SYMBOL does not have word id 0 in vocabulary.")
#         return False
#     if vocab[C.UNK_SYMBOL] != C.UNK_ID:
#         logger.warning("UNK_SYMBOL does not have word id 1 in vocabulary.")
#         return False
#     if vocab[C.BOS_SYMBOL] != C.BOS_ID:
#         logger.warning("BOS_SYMBOL does not have word id 2 in vocabulary.")
#         return False
#     if vocab[C.EOS_SYMBOL] != C.EOS_ID:
#         logger.warning("EOS_SYMBOL does not have word id 3 in vocabulary.")
#         return False
#     word_ids = []
#     for word, word_id in vocab.items():
#         word_ids.append(word_id)
#     word_ids_set = set(word_ids)
#     if len(word_ids_set) != len(word_ids):
#         logger.warning("Duplicate word_ids in vocabulary.")
#         return False
# 
#     expected_word_ids = set(range(0, len(vocab)))
#     if expected_word_ids != word_ids_set:
#         logger.warning("Not all word_ids from 0 to len(vocabulary) present in vocabulary.")
#         return False
# 
#     return True
# 
# def _get_sorted_source_vocab_fnames(folder) -> List[str]:
#     _key = lambda x: int(x.split('.', 3)[-2])
#     return sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)], key=_key)
# 
# def count_tokens(data: Iterable[str]) -> Counter:
#     """
#     Count whitespace delimited tokens.
# 
#     :param data: Sequence of sentences containing whitespace-delimited tokens.
#     :return: Token counter.
#     """
#     return Counter(token for line in data for token in utils.get_tokens(line))


import pytest
import sockeye.constants as C
from unittest import mock
from collections import Counter
from sockeye.vocab import (build_vocab, get_ordered_tokens_from_vocab, is_valid_vocab, \
    _get_sorted_source_vocab_fnames, count_tokens)
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5, "one": 6}),
        (["one two three", "one two three"], 2, 2,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5}),
        # Example 2
        (["one one two three ", "one two three"], 3, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5, "three": 6}),
        (["one one two three ", "one two three"], 3, 2,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5, "three": 6}),
        (["one one two three ", "one two three"], 3, 3,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4}),
        (["one one two three ", "one two three"], 2, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5}),
         # Example 3 (including special symbols)
        (["one two three <s> <s>", "one two three <s> <s>"], None, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5, "one": 6}),
        ]


@pytest.mark.parametrize("data,size,min_count,expected", test_vocab)
def test_build_vocab(data, size, min_count, expected):
    vocab = build_vocab(data=data, num_words=size, min_count=min_count)
    assert vocab == expected


@pytest.mark.parametrize("num_types,pad_to_multiple_of,expected_vocab_size",
                         [(4, None, 8), (2, 8, 8), (4, 8, 8), (8, 8, 16), (10, 16, 16), (13, 16, 32)])
def test_padded_build_vocab(num_types, pad_to_multiple_of, expected_vocab_size):
    data = [" ".join('word%d' % i for i in range(num_types))]
    size = None
    min_count = 1

# Path: sockeye/vocab.py
# def build_vocab(data: Iterable[str], num_words: Optional[int] = None, min_count: int = 1,
#                 pad_to_multiple_of: Optional[int] = None) -> Vocab:
#     """
#     Creates a vocabulary mapping from words to ids. Increasing integer ids are assigned by word frequency,
#     using lexical sorting as a tie breaker. The only exception to this are special symbols such as the padding symbol
#     (PAD).
# 
#     :param data: Sequence of sentences containing whitespace-delimited tokens.
#     :param num_words: Optional maximum number of words in the vocabulary.
#     :param min_count: Minimum occurrences of words to be included in the vocabulary.
#     :param pad_to_multiple_of: If not None, pads the vocabulary to a size that is the next multiple of this int.
#     :return: Word-to-id mapping.
#     """
#     raw_vocab = count_tokens(data)
#     return build_pruned_vocab(raw_vocab=raw_vocab,
#                               num_words=num_words,
#                               min_count=min_count,
#                               pad_to_multiple_of=pad_to_multiple_of)
# 
# def get_ordered_tokens_from_vocab(vocab: Vocab) -> List[str]:
#     """
#     Returns the list of tokens in a vocabulary, ordered by increasing vocabulary id.
# 
#     :param vocab: Input vocabulary.
#     :return: List of tokens.
#     """
#     return [token for token, token_id in sorted(vocab.items(), key=lambda i: i[1])]
# 
# def is_valid_vocab(vocab: Vocab) -> bool:
#     """
#     Checks if a vocabulary is valid. We define valid as:
#     1. All indices from 0 to num_words - 1 are present without duplicates.
#     2. PAD_SYMBOL has word id 0, UNK_SYMBOL has word id 1, BOS_SYMBOL has word id 2, EOS_SYMBOL has word id 3.
#     """
#     if vocab[C.PAD_SYMBOL] != C.PAD_ID:
#         logger.warning("PAD_SYMBOL does not have word id 0 in vocabulary.")
#         return False
#     if vocab[C.UNK_SYMBOL] != C.UNK_ID:
#         logger.warning("UNK_SYMBOL does not have word id 1 in vocabulary.")
#         return False
#     if vocab[C.BOS_SYMBOL] != C.BOS_ID:
#         logger.warning("BOS_SYMBOL does not have word id 2 in vocabulary.")
#         return False
#     if vocab[C.EOS_SYMBOL] != C.EOS_ID:
#         logger.warning("EOS_SYMBOL does not have word id 3 in vocabulary.")
#         return False
#     word_ids = []
#     for word, word_id in vocab.items():
#         word_ids.append(word_id)
#     word_ids_set = set(word_ids)
#     if len(word_ids_set) != len(word_ids):
#         logger.warning("Duplicate word_ids in vocabulary.")
#         return False
# 
#     expected_word_ids = set(range(0, len(vocab)))
#     if expected_word_ids != word_ids_set:
#         logger.warning("Not all word_ids from 0 to len(vocabulary) present in vocabulary.")
#         return False
# 
#     return True
# 
# def _get_sorted_source_vocab_fnames(folder) -> List[str]:
#     _key = lambda x: int(x.split('.', 3)[-2])
#     return sorted([f for f in os.listdir(folder) if f.startswith(C.VOCAB_SRC_PREFIX)], key=_key)
# 
# def count_tokens(data: Iterable[str]) -> Counter:
#     """
#     Count whitespace delimited tokens.
# 
#     :param data: Sequence of sentences containing whitespace-delimited tokens.
#     :return: Token counter.
#     """
#     return Counter(token for line in data for token in utils.get_tokens(line))


# Path: test/unit/test_vocab.py
import pytest
import sockeye.constants as C
from unittest import mock
from collections import Counter
from sockeye.vocab import (build_vocab, get_ordered_tokens_from_vocab, is_valid_vocab, \
    _get_sorted_source_vocab_fnames, count_tokens)

         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5, "one": 6}),
        (["one two three", "one two three"], 2, 2,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5}),
        # Example 2
        (["one one two three ", "one two three"], 3, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5, "three": 6}),
        (["one one two three ", "one two three"], 3, 2,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5, "three": 6}),
        (["one one two three ", "one two three"], 3, 3,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4}),
        (["one one two three ", "one two three"], 2, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "one": 4, "two": 5}),
         # Example 3 (including special symbols)
        (["one two three <s> <s>", "one two three <s> <s>"], None, 1,
         {"<pad>": 0, "<unk>": 1, "<s>": 2, "</s>": 3, "two": 4, "three": 5, "one": 6}),
        ]


@pytest.mark.parametrize("data,size,min_count,expected", test_vocab)
def test_build_vocab(data, size, min_count, expected):
    vocab = build_vocab(data=data, num_words=size, min_count=min_count)
    assert vocab == expected


@pytest.mark.parametrize("num_types,pad_to_multiple_of,expected_vocab_size",
                         [(4, None, 8), (2, 8, 8), (4, 8, 8), (8, 8, 16), (10, 16, 16), (13, 16, 32)])
def test_padded_build_vocab(num_types, pad_to_multiple_of, expected_vocab_size):
    data = [" ".join('word%d' % i for i in range(num_types))]
    size = None
    min_count = 1

    vocab = build_vocab(data, size, min_count, pad_to_multiple_of=pad_to_multiple_of)