Predict
(env)
| 189 | |
| 190 | |
| 191 | def predict(env): |
| 192 | """Predict""" |
| 193 | args = env.args |
| 194 | |
| 195 | logging.info("Load the dataset") |
| 196 | if args.prob: |
| 197 | env.fields = env.fields._replace(PHEAD=Field("prob")) |
| 198 | predicts = Corpus.load(args.infer_data_path, env.fields) |
| 199 | dataset = TextDataset(predicts, [env.WORD, env.FEAT], args.buckets) |
| 200 | # set the data loader |
| 201 | dataset.loader = batchify(dataset, args.batch_size) |
| 202 | logging.info("{} sentences, {} batches".format(len(dataset), len(dataset.loader))) |
| 203 | |
| 204 | logging.info("Load the model") |
| 205 | model = load(args.model_path) |
| 206 | model.args = args |
| 207 | |
| 208 | logging.info("Make predictions on the dataset") |
| 209 | start = datetime.datetime.now() |
| 210 | model.eval() |
| 211 | pred_arcs, pred_rels, pred_probs = epoch_predict(env, args, model, dataset.loader) |
| 212 | total_time = datetime.datetime.now() - start |
| 213 | # restore the order of sentences in the buckets |
| 214 | indices = np.argsort(np.array([i for bucket in dataset.buckets.values() for i in bucket])) |
| 215 | predicts.head = [pred_arcs[i] for i in indices] |
| 216 | predicts.deprel = [pred_rels[i] for i in indices] |
| 217 | if args.prob: |
| 218 | predicts.prob = [pred_probs[i] for i in indices] |
| 219 | logging.info("Save the predicted result to {}".format(args.infer_result_path)) |
| 220 | predicts.save(args.infer_result_path) |
| 221 | logging.info("{}s elapsed, {:.2f} Sents/s".format(total_time, len(dataset) / total_time.total_seconds())) |
| 222 | |
| 223 | |
| 224 | def predict_query(env): |
no test coverage detected