(input_ids, cls_id=CLS_ID, eos_id=EOS_ID)
| 73 | |
| 74 | |
| 75 | def get_segment_index(input_ids, cls_id=CLS_ID, eos_id=EOS_ID): |
| 76 | mask = (input_ids == cls_id).to( |
| 77 | dtype=torch.long) + (input_ids == eos_id).to(dtype=torch.long) |
| 78 | mask = mask + torch.cat([torch.zeros_like(mask[:, 0:1]), mask[:, :-1]], |
| 79 | dim=1) |
| 80 | return mask.cumsum(dim=1) - 1 |
| 81 | |
| 82 | |
| 83 | def get_token_type_mask(input_ids, cls_id=CLS_ID, eos_id=EOS_ID): |