def __init__(self, model_path, device="cpu"):
super().__init__(model_path, device)
self.vocab = Vocabulary()
self.races = Races()
self.genders = Genders()
self.to_tensor = ToTensor()
self.name_transform = Compose([self.vocab, OneHot(self.vocab.size), ToTensor()])
self.race_transform = Compose([self.races, OneHot(self.races.size), ToTensor()])
self.gender_transform = Compose([self.genders, OneHot(self.genders.size), ToTensor()])
def init_dataset(self):
return DnDCharacterNameDataset(
root_dir=self.root_dir,
name_transform=Compose(
[self.vocab, OneHot(self.vocab.size),
ToTensor()]),
race_transform=Compose(
[self.races, OneHot(self.races.size),
ToTensor()]),
gender_transform=Compose(
[self.genders,
OneHot(self.genders.size),
ToTensor()]),
target_transform=Compose([self.vocab, ToTensor()]))