add experiment tracking integration to text, image classification & LLM finetuning

src/autotrain/cli/run_image_classification.py

@@ -203,6 +203,13 @@ def register_subcommand(parser: ArgumentParser):
                 "required": False,
                 "type": str,
             },
+            {
+                "arg": "--log",
+                "help": "Use experiment tracking",
+                "required": False,
+                "type": str,
+                "default": "none",
+            },
         ]
         run_text_classification_parser = parser.add_parser(
             "image-classification", description="✨ Run AutoTrain Image Classification"
@@ -291,6 +298,7 @@ def run(self):
                 fp16=self.args.fp16,
                 push_to_hub=self.args.push_to_hub,
                 repo_id=self.args.repo_id,
+                log=self.args.log,
             )
             params.save(output_dir=self.args.project_name)
             if self.num_gpus == 1:

src/autotrain/cli/run_llm.py

@@ -340,6 +340,13 @@ def register_subcommand(parser: ArgumentParser):
                 "action": "store_true",
                 "alias": ["--use-flash-attention-2", "--use-fa2"],
             },
+            {
+                "arg": "--log",
+                "help": "Use experiment tracking",
+                "required": False,
+                "type": str,
+                "default": "none",
+            },
             {
                 "arg": "--disable_gradient_checkpointing",
                 "help": "Disable gradient checkpointing",
@@ -488,6 +495,7 @@ def run(self):
                 merge_adapter=self.args.merge_adapter,
                 username=self.args.username,
                 use_flash_attention_2=self.args.use_flash_attention_2,
+                log=self.args.log,
                 rejected_text_column=self.args.rejected_text_column,
                 disable_gradient_checkpointing=self.args.disable_gradient_checkpointing,
             )

src/autotrain/cli/run_text_classification.py

@@ -232,6 +232,13 @@ def register_subcommand(parser: ArgumentParser):
                 "required": False,
                 "type": str,
             },
+            {
+                "arg": "--log",
+                "help": "Use experiment tracking",
+                "required": False,
+                "type": str,
+                "default": "none",
+            },
         ]
         run_text_classification_parser = parser.add_parser(
             "text-classification", description="✨ Run AutoTrain Text Classification"
@@ -326,6 +333,7 @@ def run(self):
                 repo_id=self.args.repo_id,
                 token=self.args.token,
                 username=self.args.username,
+                log=self.args.log,
             )
 
             if self.args.backend.startswith("spaces"):

src/autotrain/trainers/clm/__main__.py

@@ -304,7 +304,7 @@ def train(config):
 
     else:
         logging_steps = config.logging_steps
-
+    
     training_args = dict(
         output_dir=config.project_name,
         per_device_train_batch_size=config.batch_size,
@@ -316,7 +316,7 @@ def train(config):
         save_total_limit=config.save_total_limit,
         save_strategy=config.save_strategy,
         gradient_accumulation_steps=config.gradient_accumulation,
-        report_to="tensorboard",
+        report_to=config.log,
         auto_find_batch_size=config.auto_find_batch_size,
         lr_scheduler_type=config.scheduler,
         optim=config.optimizer,

src/autotrain/trainers/clm/params.py

@@ -46,6 +46,7 @@ class LLMTrainingParams(BaseModel):
     merge_adapter: bool = Field(False, title="Merge adapter")
     username: str = Field(None, title="Hugging Face Username")
     use_flash_attention_2: bool = Field(False, title="Use flash attention 2")
+    log: str = Field("none", title="Logging using experiment tracking")
     disable_gradient_checkpointing: bool = Field(False, title="Gradient checkpointing")
 
     def save(self, output_dir):

src/autotrain/trainers/image_classification.py

@@ -233,7 +233,7 @@ def train(co2_tracker, payload, huggingface_token, model_path):
     fp16 = True
     if device == "cpu":
         fp16 = False
-
+    
     training_args = dict(
         output_dir=model_path,
         per_device_train_batch_size=job_config.train_batch_size,
@@ -248,7 +248,7 @@ def train(co2_tracker, payload, huggingface_token, model_path):
         save_strategy="epoch",
         disable_tqdm=not bool(os.environ.get("ENABLE_TQDM", 0)),
         gradient_accumulation_steps=job_config.gradient_accumulation_steps,
-        report_to="none",
+        report_to=job_config.log,
         auto_find_batch_size=True,
         lr_scheduler_type=job_config.scheduler,
         optim=job_config.optimizer,

src/autotrain/trainers/image_classification/__main__.py

@@ -98,7 +98,7 @@ def train(config):
 
     else:
         logging_steps = config.logging_steps
-
+    
     training_args = dict(
         output_dir=config.project_name,
         per_device_train_batch_size=config.batch_size,
@@ -111,7 +111,7 @@ def train(config):
         save_total_limit=config.save_total_limit,
         save_strategy=config.save_strategy,
         gradient_accumulation_steps=config.gradient_accumulation,
-        report_to="tensorboard",
+        report_to=config.log,
         auto_find_batch_size=config.auto_find_batch_size,
         lr_scheduler_type=config.scheduler,
         optim=config.optimizer,

src/autotrain/trainers/image_classification/params.py

@@ -30,6 +30,7 @@ class ImageClassificationParams(BaseModel):
     evaluation_strategy: str = Field("epoch", title="Evaluation strategy")
     image_column: str = Field("image", title="Image column")
     target_column: str = Field("target", title="Target column")
+    log: str = Field("none", title="Logging using experiment tracking")
 
     def __str__(self):
         data = self.dict()

src/autotrain/trainers/lm_trainer.py

@@ -465,4 +465,4 @@ def group_texts(examples):
     model_repo.git_pull()
     model_repo.git_add()
     model_repo.git_commit(commit_message="Commit From AutoTrain")
-    model_repo.git_push()
+    model_repo.git_push()

src/autotrain/trainers/text_classification.py

@@ -274,4 +274,4 @@ def train(co2_tracker, payload, huggingface_token, model_path):
     model_repo.git_pull()
     model_repo.git_add()
     model_repo.git_commit(commit_message="Commit From AutoTrain")
-    model_repo.git_push()
+    model_repo.git_push()

src/autotrain/trainers/text_classification/__main__.py

@@ -113,7 +113,7 @@ def train(config):
 
     else:
         logging_steps = config.logging_steps
-
+    
     training_args = dict(
         output_dir=config.project_name,
         per_device_train_batch_size=config.batch_size,
@@ -126,7 +126,7 @@ def train(config):
         save_total_limit=config.save_total_limit,
         save_strategy=config.save_strategy,
         gradient_accumulation_steps=config.gradient_accumulation,
-        report_to="tensorboard",
+        report_to=config.log,
         auto_find_batch_size=config.auto_find_batch_size,
         lr_scheduler_type=config.scheduler,
         optim=config.optimizer,

src/autotrain/trainers/text_classification/params.py

@@ -34,6 +34,7 @@ class TextClassificationParams(BaseModel):
     repo_id: str = Field(None, title="Repo id")
     evaluation_strategy: str = Field("epoch", title="Evaluation strategy")
     username: str = Field(None, title="Hugging Face Username")
+    log: str = Field("none", title="Logging using experiment tracking")
 
     def __str__(self):
         data = self.dict()