mirror of
https://github.com/huggingface/transformers.git
synced 2025-07-03 21:00:08 +06:00
1b57de8dcf
* make output_dir optional * inintaied a basic testing module to validate and verify the changes * Test output_dir default to 'tmp_trainer' when unspecified. * test existing functionality of output_dir. * test that output dir only created when needed * final check * added doc string and changed the tmp_trainer to trainer_output * amke style fixes to test file. * another round of fixup --------- Co-authored-by: sambhavnoobcoder <indosambahv@gmail.com>
43 lines
1.6 KiB
Python
43 lines
1.6 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from transformers import TrainingArguments
|
|
|
|
|
|
class TestTrainingArguments(unittest.TestCase):
|
|
def test_default_output_dir(self):
|
|
"""Test that output_dir defaults to 'tmp_trainer' when not specified."""
|
|
args = TrainingArguments(output_dir=None)
|
|
self.assertEqual(args.output_dir, "tmp_trainer")
|
|
|
|
def test_custom_output_dir(self):
|
|
"""Test that output_dir is respected when specified."""
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
args = TrainingArguments(output_dir=tmp_dir)
|
|
self.assertEqual(args.output_dir, tmp_dir)
|
|
|
|
def test_output_dir_creation(self):
|
|
"""Test that output_dir is created only when needed."""
|
|
with tempfile.TemporaryDirectory() as tmp_dir:
|
|
output_dir = os.path.join(tmp_dir, "test_output")
|
|
|
|
# Directory should not exist before creating args
|
|
self.assertFalse(os.path.exists(output_dir))
|
|
|
|
# Create args with save_strategy="no" - should not create directory
|
|
args = TrainingArguments(
|
|
output_dir=output_dir,
|
|
do_train=True,
|
|
save_strategy="no",
|
|
report_to=None,
|
|
)
|
|
self.assertFalse(os.path.exists(output_dir))
|
|
|
|
# Now set save_strategy="steps" - should create directory when needed
|
|
args.save_strategy = "steps"
|
|
args.save_steps = 1
|
|
self.assertFalse(os.path.exists(output_dir)) # Still shouldn't exist
|
|
|
|
# Directory should be created when actually needed (e.g. in Trainer)
|