| 237 | |
| 238 | |
| 239 | def create_dummy_data(num_samples, seq_length, vocab_size, model_type): |
| 240 | input_ids = torch.randint(0, vocab_size, (num_samples, seq_length)) |
| 241 | attention_mask = torch.ones((num_samples, seq_length)) |
| 242 | if model_type == ModelType.mlm: |
| 243 | labels = torch.randint(0, vocab_size, (num_samples, seq_length)) |
| 244 | mask = torch.rand(num_samples, seq_length) < 0.7 |
| 245 | labels[mask] = -100 |
| 246 | elif model_type == ModelType.seqcls: |
| 247 | labels = torch.randint(0, 5, (num_samples, 1)) |
| 248 | else: |
| 249 | raise ValueError(f"Invalid model type: {model_type}") |
| 250 | return TensorDataset(input_ids, attention_mask, labels) |
| 251 | |
| 252 | |
| 253 | def tile_list_to_length(lst, length): |