(src, tgt, pad_idx, device)
| 102 | |
| 103 | # Create masks for input into model |
| 104 | def create_mask(src, tgt, pad_idx, device): |
| 105 | |
| 106 | # Get sequence length |
| 107 | src_seq_len = src.shape[0] |
| 108 | tgt_seq_len = tgt.shape[0] |
| 109 | |
| 110 | # Generate the mask |
| 111 | tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device) |
| 112 | src_mask = torch.zeros((src_seq_len, src_seq_len),device=device).type(torch.bool) |
| 113 | |
| 114 | # Overlay the mask over the original input |
| 115 | src_padding_mask = (src == pad_idx).transpose(0, 1) |
| 116 | tgt_padding_mask = (tgt == pad_idx).transpose(0, 1) |
| 117 | return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask |
| 118 | |
| 119 | # A small test to make sure our data loasd in correctly |
| 120 | if __name__=="__main__": |
no test coverage detected