from typing import Counter


def get_vocab(training_set, vocab_size_threshold=5):
  """Get the vocabulary from the training set"""
  vocab = []
  for st in training_set:
    for s in st:
      vocab.extend(s)

  vocab = Counter(vocab)
  vocab_truncate = [w for w in vocab if vocab[w] >= vocab_size_threshold]
  
  word2id = {"_GOO": 0, "_EOS": 1, "_PAD": 2, "_UNK": 3}
  id2word = {0: "_GOO", 1: "_EOS", 2: "_PAD", 3: "_UNK"}

  i = len(word2id)
  for w in vocab_truncate:
    word2id[w] = i
    id2word[i] = w
    i += 1
  
  assert(len(word2id) == len(id2word))
  print("vocabulary size: %d" % len(word2id))
  return word2id, id2word