def train()

in src/jobs/tune_bart.py [0:0]

• 97 lines of code
• 4 McCabe index (conditional complexity)

    def train(self):
        torch.cuda.empty_cache()

        # Ensure W&B logs the model at the end of training
        os.environ["WANDB_LOG_MODEL"] = "end"

        # Initialize W&B run
        wandb.init(
            project="tab_grouping",
            config={
                "learning_rate": self.learning_rate,
                "batch_size": self.batch_size,
                "model_name": self.model_name,
                "label_column": self.label_column,
                "use_keywords": self.use_keywords,
                "learning_rate_decay": self.learning_rate_decay,
                "single_tab_handling": self.single_tab_handling,
                "input_prompt_id": INPUT_PROMPT_ID,
                "filename": self.filename
            }
        )
        print(f"W&B Run ID: {wandb.run.id}")
        print(f"W&B Run Name: {wandb.run.name}")

        tokenized_training_dataset = self.train_dataset.map(self.preprocess_function, batched=True)
        tokenized_eval_dataset = self.eval_dataset.map(self.preprocess_function, batched=True)

        if self.learning_rate_decay:
            training_args = TrainingArguments(
                output_dir="./results",
                evaluation_strategy="epoch",
                learning_rate=self.learning_rate,
                per_device_train_batch_size=self.batch_size,
                per_device_eval_batch_size=1,
                num_train_epochs=3,
                weight_decay=0.01,
                save_total_limit=1,
                save_strategy="epoch",
                lr_scheduler_type="cosine",
                warmup_ratio=0.1
            )
        else:
            training_args = TrainingArguments(
                output_dir="./results",
                evaluation_strategy="epoch",
                learning_rate=self.learning_rate,
                per_device_train_batch_size=self.batch_size,
                per_device_eval_batch_size=1,
                num_train_epochs=3,
                weight_decay=0.01,
                save_total_limit=1,
                save_strategy="epoch",
            )

        trainer = Trainer(
            model=self.model,
            args=training_args,
            train_dataset=tokenized_training_dataset,
            eval_dataset=tokenized_eval_dataset
        )

        trainer.train()

        # Generate predictions on the evaluation set
        results_labels = []
        results_output = []
        for item in tokenized_eval_dataset:
            input_ids = self.tokenizer(
                item["input_text"],
                return_tensors="pt"
            ).input_ids.to(self.model.device)

            label = item["target_text"]

            outputs = self.model.generate(
                input_ids,
                max_length=30,
                num_return_sequences=1
            )
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            results_output.append(response)
            results_labels.append(label)

        self.compute_metrics_text(results_output, results_labels)

        # Log sample predictions in a W&B Table
        validation_table = wandb.Table(
            columns=["input", "label", "prediction"],
            data=list(
                zip(
                    [d["input_text"] for d in tokenized_eval_dataset],
                    results_labels,
                    results_output
                )
            ),
        )
        wandb.log({"Validation Data": validation_table})

        self.model.config.bad_words_ids = get_bad_word_ids()

        output_dir = "./bart-finetuned-topic"
        self.model.save_pretrained(output_dir)
        self.tokenizer.save_pretrained(output_dir)

        current_date = datetime.now()
        date_string = current_date.isoformat().replace(":", "_")
        upload_directory(
            output_dir,
            "stage-fx-tab-grouping",
            f"topic/models/{date_string}/",
            depth=1
        )

        wandb.finish()

        self.model = None
        torch.cuda.empty_cache()