delete

Author: ydshieh

tests/trainer/test_trainer.py

@@ -163,6 +163,22 @@
 PATH_SAMPLE_TEXT = f"{get_tests_dir()}/fixtures/sample_text.txt"
 
 
+def get_dataset(file_path, tokenizer, max_len):
+    dataset = datasets.load_dataset("text", file_path)
+
+    # Define tokenization function
+    def tokenize_function(examples):
+        tokenized = tokenizer(examples["text"], add_special_tokens=True, truncation=True, max_length=max_len)
+        # Add labels as a copy of input_ids
+        tokenized["labels"] = tokenized["input_ids"].copy()
+        return tokenized
+
+    # Apply tokenization and remove original text column
+    tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
+
+    return tokenized_dataset
+
+
 class StoreLossCallback(TrainerCallback):
     """
     Simple callback to store the loss.
@@ -1528,13 +1544,7 @@ def test_multiple_peft_adapters(self):
         tiny_model = get_peft_model(tiny_model, peft_config, "adapter1")
         tiny_model.add_adapter("adapter2", peft_config)
 
-        train_dataset = LineByLineTextDataset(
-            tokenizer=tokenizer,
-            file_path=PATH_SAMPLE_TEXT,
-            block_size=tokenizer.max_len_single_sentence,
-        )
-        for example in train_dataset.examples:
-            example["labels"] = example["input_ids"]
+        train_dataset = get_dataset(PATH_SAMPLE_TEXT, tokenizer, tokenizer.max_len_single_sentence)
 
         tokenizer.pad_token = tokenizer.eos_token
 
@@ -3754,13 +3764,7 @@ def test_trainer_eval_multiple(self):
         MODEL_ID = "openai-community/gpt2"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
         model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
-        dataset = LineByLineTextDataset(
-            tokenizer=tokenizer,
-            file_path=PATH_SAMPLE_TEXT,
-            block_size=tokenizer.max_len_single_sentence,
-        )
-        for example in dataset.examples:
-            example["labels"] = example["input_ids"]
+        dataset = get_dataset(PATH_SAMPLE_TEXT, tokenizer, tokenizer.max_len_single_sentence)
         with tempfile.TemporaryDirectory() as tmp_dir:
             training_args = TrainingArguments(
                 output_dir=tmp_dir,
@@ -3784,11 +3788,7 @@ def test_trainer_eval_multiple(self):
     def test_trainer_eval_lm(self):
         MODEL_ID = "distilbert/distilroberta-base"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        dataset = LineByLineTextDataset(
-            tokenizer=tokenizer,
-            file_path=PATH_SAMPLE_TEXT,
-            block_size=tokenizer.max_len_single_sentence,
-        )
+        dataset = get_dataset(PATH_SAMPLE_TEXT, tokenizer, tokenizer.max_len_single_sentence)
         self.assertEqual(len(dataset), 31)
 
     def test_training_iterable_dataset(self):
@@ -4975,13 +4975,7 @@ def test_trainer_works_without_model_config(self):
         model = BasicTextGenerationModel(vocab_size=tokenizer.vocab_size, hidden_size=32)
         # Note that this class does not have a config attribute
 
-        train_dataset = LineByLineTextDataset(
-            tokenizer=tokenizer,
-            file_path=PATH_SAMPLE_TEXT,
-            block_size=tokenizer.max_len_single_sentence,
-        )
-        for example in train_dataset.examples:
-            example["labels"] = example["input_ids"]
+        train_dataset = get_dataset(PATH_SAMPLE_TEXT, tokenizer, tokenizer.max_len_single_sentence)
 
         with tempfile.TemporaryDirectory() as tmpdir:
             training_args = TrainingArguments(