Nested/utils/vocab.py

from collections import Counter
class Vocab:
    def _init_(self, counter, specials=[]) -> None:
        self.itos = list(counter.keys()) + specials
        self.stoi = {s: i for i, s in enumerate(self.itos)}
        self.word_count = counter

    def get_itos(self) -> list[str]:
        return self.itos

    def get_stoi(self) -> dict[str, int]:
        return self.stoi

    def _len_(self):
        return len(self.itos)