relation-api / Nested /utils /vocab.py
aaljabari's picture
Create vocab.py
86d978f verified
Raw
History Blame Contribute Delete
425 Bytes
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)