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