Predict in one epoch
(env, args, model, loader)
| 177 | |
| 178 | @dygraph.no_grad |
| 179 | def epoch_predict(env, args, model, loader): |
| 180 | """Predict in one epoch""" |
| 181 | arcs, rels, probs = [], [], [] |
| 182 | pad_index = args.pad_index |
| 183 | bos_index = args.bos_index |
| 184 | eos_index = args.eos_index |
| 185 | for batch, inputs in enumerate(loader(), start=1): |
| 186 | if args.encoding_model.startswith("ernie"): |
| 187 | words = inputs[0] |
| 188 | s_arc, s_rel, words = model(words) |
| 189 | else: |
| 190 | words, feats = inputs |
| 191 | s_arc, s_rel, words = model(words, feats) |
| 192 | mask = layers.logical_and( |
| 193 | layers.logical_and(words != pad_index, words != bos_index), |
| 194 | words != eos_index, |
| 195 | ) |
| 196 | lens = nn.reduce_sum(mask, -1) |
| 197 | arc_preds, rel_preds = decode(args, s_arc, s_rel, mask) |
| 198 | arcs.extend(layers.split(nn.masked_select(arc_preds, mask), lens.numpy().tolist())) |
| 199 | rels.extend(layers.split(nn.masked_select(rel_preds, mask), lens.numpy().tolist())) |
| 200 | if args.prob: |
| 201 | arc_probs = nn.index_sample(layers.softmax(s_arc, axis=-1), layers.unsqueeze(arc_preds, -1)) |
| 202 | probs.extend( |
| 203 | layers.split( |
| 204 | nn.masked_select(layers.squeeze(arc_probs, axes=[-1]), mask), |
| 205 | lens.numpy().tolist(), |
| 206 | )) |
| 207 | arcs = [seq.numpy().tolist() for seq in arcs] |
| 208 | rels = [env.REL.vocab[seq.numpy().tolist()] for seq in rels] |
| 209 | probs = [[round(p, 3) for p in seq.numpy().tolist()] for seq in probs] |
| 210 | |
| 211 | return arcs, rels, probs |
| 212 | |
| 213 | |
| 214 | def loss_function(s_arc, s_rel, arcs, rels, mask): |