| 25 | ) |
| 26 | |
| 27 | def load_and_preprocess_data(task, tokenizer, args): |
| 28 | |
| 29 | if "mnli" in task: |
| 30 | dataset = load_dataset("glue", "mnli") |
| 31 | else: |
| 32 | dataset = load_dataset("glue", task) |
| 33 | |
| 34 | def tokenize_function(examples): |
| 35 | |
| 36 | # Handle different input formats |
| 37 | if "premise" in examples and "hypothesis" in examples: |
| 38 | # MNLI and similar tasks |
| 39 | return tokenizer( |
| 40 | examples["premise"], |
| 41 | examples["hypothesis"], |
| 42 | truncation=True, |
| 43 | padding="max_length", |
| 44 | max_length=args.max_seq_length, |
| 45 | ) |
| 46 | elif "question" in examples and "sentence" in examples: |
| 47 | # QNLI and similar tasks |
| 48 | return tokenizer( |
| 49 | examples["question"], |
| 50 | examples["sentence"], |
| 51 | truncation=True, |
| 52 | padding="max_length", |
| 53 | max_length=args.max_seq_length, |
| 54 | ) |
| 55 | elif "sentence1" in examples and "sentence2" in examples: |
| 56 | # MRPC, STS-B |
| 57 | return tokenizer( |
| 58 | examples["sentence1"], |
| 59 | examples["sentence2"], |
| 60 | truncation=True, |
| 61 | padding="max_length", |
| 62 | max_length=args.max_seq_length, |
| 63 | ) |
| 64 | elif "question1" in examples and "question2" in examples: |
| 65 | # QQP |
| 66 | return tokenizer( |
| 67 | examples["question1"], |
| 68 | examples["question2"], |
| 69 | truncation=True, |
| 70 | padding="max_length", |
| 71 | max_length=args.max_seq_length, |
| 72 | ) |
| 73 | elif "sentence" in examples: |
| 74 | # CoLA, SST-2 |
| 75 | return tokenizer( |
| 76 | examples["sentence"], |
| 77 | truncation=True, |
| 78 | padding="max_length", |
| 79 | max_length=args.max_seq_length, |
| 80 | ) |
| 81 | else: |
| 82 | raise ValueError(f"Unexpected format for task {task}") |
| 83 | |
| 84 | tokenized_datasets = dataset.map(tokenize_function, batched=True) |