def train():
print('Preprocessing raw data')
preprocessor = Preprocessor()
preprocessor.preprocess()
dataset = Dataset(preprocessor)
print('Training MF')
mf = MF(preprocessor, dataset)
mf.train_or_load_if_exists()
print('Building I2I')
i2i = Item2Item(dataset)
print('Generating candidates')
candidate_generator = CandidateGenerator(preprocessor, dataset, mf, i2i)
X_train, y_train, q_train, q_train_reader = candidate_generator.generate_train()
X_val, y_val, q_val, q_val_reader = candidate_generator.generate_val()
import pickle
try:
with open('puke.pkl', 'wb') as f:
pickle.dump((X_train, y_train, q_train, q_train_reader,
X_val, y_val, q_val, q_val_reader), f)
except:
print("Couldn't save puke")
print('Training ranker')
ranker = Ranker()
ranker.train(X_train, y_train, q_train, X_val, y_val, q_val)
ranker.save()
print('Validating ranker')
rank_scores = ranker.rank(X_val)
print('ndcg', dataset.validate_ndcg(y_val, q_val, q_val_reader, rank_scores))
if args.which_ranker == 'ranker':
from ranker import Ranker
elif args.which_ranker == 'masker_ranker':
from masker_ranker import Ranker
model = Ranker(vocab_src, vocab_tgt, args.embed_dim, args.ff_embed_dim,
args.num_heads, args.dropout, args.num_layers)
model = model.cuda()
optimizer = torch.optim.Adam(model.parameters(), args.lr)
train_data = DataLoader(args.train_data, vocab_src, vocab_tgt,
args.train_batch_size, True)
dev_data = DataLoader(args.dev_data, vocab_src, vocab_tgt,
args.dev_batch_size, True)
model.train()
loss_accumulated = 0.
acc_accumulated = 0.
batches_processed = 0
best_dev_acc = 0
for epoch in range(args.epochs):
for src_input, tgt_input in train_data:
optimizer.zero_grad()
loss, acc = model(src_input, tgt_input)
loss_accumulated += loss.item()
acc_accumulated += acc
batches_processed += 1
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
optimizer.step()
