Forward step.
(data, model, args, timers, mems, eval_metric=None)
| 31 | |
| 32 | |
| 33 | def lm_forward_step(data, model, args, timers, mems, eval_metric=None): |
| 34 | """Forward step.""" |
| 35 | # Get the batch. |
| 36 | if timers is not None: |
| 37 | timers('batch generator').start() |
| 38 | try: |
| 39 | data = next(data) |
| 40 | except BaseException: |
| 41 | data = data |
| 42 | |
| 43 | if 'mask' in data: |
| 44 | # finetune SQuAD |
| 45 | data['attention_mask'] = data.pop('mask') |
| 46 | data['position_id'] = data.pop('position') |
| 47 | data['loss_mask'] = data.pop('logit_mask') |
| 48 | |
| 49 | tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data, args) |
| 50 | if timers is not None: |
| 51 | timers('batch generator').stop() |
| 52 | |
| 53 | if tokens.dim() == 3: |
| 54 | tokens = tokens.squeeze(1) |
| 55 | labels = labels.squeeze(1) |
| 56 | loss_mask = loss_mask.squeeze(1) |
| 57 | attention_mask = attention_mask.squeeze(1) |
| 58 | position_ids = position_ids.squeeze(1) |
| 59 | |
| 60 | def print_masked_text(batch_id): |
| 61 | block_position_ids = position_ids[:, 1] |
| 62 | position_ids_ = position_ids[:, 0] |
| 63 | output_tokens = [] |
| 64 | sep = attention_mask[batch_id].item() |
| 65 | for i, token in enumerate(tokens[batch_id, :sep].tolist()): |
| 66 | if global_tokenizer is not None: |
| 67 | token = global_tokenizer.IdToToken(token) |
| 68 | if token.startswith('[MASK'): |
| 69 | token = f"[{position_ids_[batch_id, i].item()}, {token}]" |
| 70 | if token.startswith('##') and len(output_tokens) > 0 and not output_tokens[-1].endswith(']'): |
| 71 | output_tokens[-1] += token[2:] |
| 72 | else: |
| 73 | output_tokens.append(token) |
| 74 | else: |
| 75 | output_tokens.append(str(token)) |
| 76 | print(" ".join(output_tokens)) |
| 77 | last_index = None |
| 78 | for i in range(sep, tokens.size(1)): |
| 79 | if global_tokenizer.IdToToken(tokens[batch_id, i].item()).startswith("<|startofpiece"): |
| 80 | if last_index is not None: |
| 81 | print(global_tokenizer.DecodeIds(tokens[batch_id, last_index: i].tolist()), "|", |
| 82 | global_tokenizer.DecodeIds(labels[batch_id, last_index: i].tolist())), |
| 83 | print(position_ids_[batch_id, last_index: i].tolist(), |
| 84 | block_position_ids[batch_id, last_index:i].tolist()) |
| 85 | last_index = i |
| 86 | if last_index is not None: |
| 87 | print(global_tokenizer.DecodeIds(tokens[batch_id, last_index:].tolist()), "|", |
| 88 | global_tokenizer.DecodeIds(labels[batch_id, last_index:].tolist())) |
| 89 | print(position_ids_[batch_id, last_index:].tolist(), block_position_ids[batch_id, last_index:].tolist()) |
| 90 |
no test coverage detected