def main(config, local):
n_gpu = int(GPU_NUM)
n_gpu = 1 if n_gpu == 0 else n_gpu
np.random.seed(config.random_seed)
if n_gpu > 0:
torch.cuda.manual_seed_all(config.random_seed)
# Create data instances
vocab = Vocabulary(config.vocab_path)
if config.mode == 'train':
# Prepare train data loader
train_dataset, val_dataset = Dataset(vocab), Dataset(vocab)
train_path = os.path.join(config.data_dir, 'train_data/train_data')
val_path = os.path.join(config.data_dir, 'train_data/val_data')
train_dataset.create_instances(train_path,
config.max_seq_length,
type='train')
val_dataset.create_instances(val_path,
config.max_seq_length,
type='val')
train_loader = DataLoader(train_dataset,
batch_size=config.batch_size * n_gpu,
shuffle=True)
val_loader = DataLoader(val_dataset,
batch_size=config.batch_size * n_gpu)
else:
train_loader, val_loader = None, None
trainer = Trainer(config, n_gpu, vocab, train_loader, val_loader)
if nsml.IS_ON_NSML:
bind_model(trainer.model, vocab, config)
if config.pause:
nsml.paused(scope=local)
if config.mode == 'train':
trainer.train()
def inference(path, model, vocab, config, **kwargs):
model.eval()
test_dataset = Dataset(vocab)
test_path = os.path.join(path, 'test_data')
test_dataset.create_instances(test_path,
config.max_seq_length,
type='test')
test_loader = DataLoader(test_dataset, batch_size=1)
pred_results = []
for step, batch in enumerate(test_loader):
batch = tuple(t.to(device) for t in batch)
batch = sort_batch(batch)
input_ids, input_lengths, labels = batch
outputs = model(input_ids)
top_1_result = outputs['predicted_intents'][0].item()
pred_results.append([step, top_1_result])
return pred_results
