def make_generator_inputs(self, data): source = preproc.mr2source_inputs(data) tokens = [self.source_vocab.start_token] + source \ + [self.source_vocab.stop_token] inputs = Variable( torch.LongTensor([[self.source_vocab[t] for t in tokens]]).t(), lengths=torch.LongTensor([len(tokens)]), length_dim=0, batch_dim=1, pad_value=self.source_vocab.pad_index) if self._gpu > -1: inputs = inputs.cuda(self._gpu) return {"source_inputs": inputs}
def run(self, env, verbose=False): output_path = env["proj_dir"] / "output" / self.filename output_path.parent.mkdir(exist_ok=True, parents=True) if self.checkpoint is None: ckpt = self._get_default_checkpoint(env) else: ckpt = self.checkpoint if ckpt is None: raise RuntimeError("No checkpoints found!") ckpt_path = env["checkpoints"][ckpt]["path"] if verbose: print("Loading model from {}".format(ckpt_path)) model = plum.load(ckpt_path).eval() if env["gpu"] > -1: model.cuda(env["gpu"]) self._gpu = env["gpu"] samples = self.make_samples() with open(output_path, "w") as out_fp: for i, mr in enumerate(samples, 1): print("{}/{}".format(i, len(samples)), end="\r" if i < len(samples) else "\n", flush=True) gen_input = self.make_generator_inputs(mr) source = preproc.mr2source_inputs(mr) tokens = self._get_outputs(model, gen_input) data = json.dumps({ "source": source, "mr": mr, "text": " ".join(tokens), }) print(data, file=out_fp, flush=True)