def test_itf():
tokenizer = Tokenizer.from_pretrained(Config.model_name)
if Config.use_pickle:
with open(f'{Config.pickle_path}', 'rb') as f:
train_data = pickle.load(f)
else:
train_data = make_train_data_from_txt(Config, tokenizer)
counter, itf = make_itf(train_data, Config.vocab_size, tokenizer)
# itf = (itf - itf.min()) / (itf.max() - itf.min())
# for i in range(itf.size(0)):
# print(i, itf[i])
# itf[itf == 0] += 1e-6
for k, v in counter.most_common(len(counter)):
print(tokenizer.decode([k]), v)
if not os.path.isdir(Config.data_dir):
os.mkdir(Config.data_dir)
seed_everything(Config.seed)
device = torch.device(Config.device)
start_epoch = 0
tokenizer = Tokenizer.from_pretrained(Config.model_name)
logging.info('Preparing training data')
if Config.use_pickle:
with open(f'{Config.pickle_path}', 'rb') as f:
train_data = pickle.load(f)
else:
train_data = make_train_data_from_txt(Config, tokenizer)
itf = make_itf(train_data, Config.vocab_size)
dataset = DialogDataset(train_data, tokenizer)
logging.info('Define Models')
model = build_model(Config).to(device)
model.unfreeze()
logging.info('Define Loss and Optimizer')
criterion = nn.CrossEntropyLoss(reduction='none')
optimizer = optim.AdamW(model.parameters(), lr=Config.lr, betas=Config.betas, eps=1e-9)
if Config.load:
state_dict = torch.load(f'{Config.data_dir}/{Config.fn}.pth')
start_epoch = 10
print(f'Start Epoch: {start_epoch}')
model.load_state_dict(state_dict['model'])