disable saving in HF Trainer integration

logix-project · May 31, 2024 · 1871908 · 1871908
1 parent 2038014
commit 1871908
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/logix/huggingface/__init__.py b/logix/huggingface/__init__.py
@@ -1,2 +1,2 @@
 from .patch import patch_trainer
-from .arguments import LogIXArgument
+from .arguments import LogIXArguments
diff --git a/logix/huggingface/arguments.py b/logix/huggingface/arguments.py
@@ -5,7 +5,7 @@
 
 
 @dataclass
-class LogIXArgument:
+class LogIXArguments:
     project: str = field(
         default="tmp_logix", metadata={"help": "The name of the project."}
     )

diff --git a/logix/huggingface/callback.py b/logix/huggingface/callback.py
@@ -3,15 +3,15 @@
 from transformers.trainer import TrainerCallback
 
 from logix import LogIX, LogIXScheduler
-from logix.huggingface.arguments import LogIXArgument
+from logix.huggingface.arguments import LogIXArguments
 
 
 class LogIXCallback(TrainerCallback):
     def __init__(
         self,
         logix: LogIX,
         logix_scheduler: LogIXScheduler,
-        args: LogIXArgument,
+        args: LogIXArguments,
     ):
         self.logix = logix
         self.logix_scheduler = iter(logix_scheduler)

diff --git a/logix/huggingface/patch.py b/logix/huggingface/patch.py
@@ -5,7 +5,7 @@
 from logix import LogIX, LogIXScheduler
 from logix.utils import DataIDGenerator
 from logix.huggingface.callback import LogIXCallback
-from logix.huggingface.arguments import LogIXArgument
+from logix.huggingface.arguments import LogIXArguments
 
 
 def patch_trainer(TrainerClass):
@@ -19,9 +19,9 @@ def patch_trainer(TrainerClass):
     class PatchedTrainer(TrainerClass):
         def __init__(
             self,
-            logix_args: Optional[LogIXArgument] = None,
+            logix_args: LogIXArguments,
             model: Union[PreTrainedModel, nn.Module] = None,
-            args: TrainingArguments = None,
+            args: Optional[TrainingArguments] = None,
             data_collator: Optional[DataCollator] = None,
             train_dataset: Optional[Dataset] = None,
             eval_dataset: Optional[Union[Dataset, Dict[str, Dataset]]] = None,
@@ -59,6 +59,7 @@ def __init__(
                 args = TrainingArguments(output_dir=output_dir)
             args.num_train_epochs = len(self.logix_scheduler)
             args.report_to = []
+            args.save_strategy = "no"
 
             super().__init__(
                 model,