| 109 | |
| 110 | |
| 111 | class RandomTokenDataset(torch.utils.data.Dataset): |
| 112 | def __init__(self, *, num_samples: int, seq_len: int, vocab_size: int, pad_token_id: int): |
| 113 | self.num_samples = int(num_samples) |
| 114 | self.seq_len = int(seq_len) |
| 115 | self.vocab_size = int(vocab_size) |
| 116 | self.pad_token_id = int(pad_token_id) |
| 117 | |
| 118 | def __len__(self): |
| 119 | return self.num_samples |
| 120 | |
| 121 | def __getitem__(self, idx): |
| 122 | del idx |
| 123 | input_ids = torch.randint(0, self.vocab_size, (self.seq_len,), dtype=torch.long) |
| 124 | attention_mask = torch.ones_like(input_ids) |
| 125 | return {"input_ids": input_ids, "attention_mask": attention_mask} |
| 126 | |
| 127 | |
| 128 | def main(): |
no outgoing calls
no test coverage detected
searching dependent graphs…