def train()

in src/jobs/tune_gpt2.py [0:0]
81 lines of code
4 McCabe index (conditional complexity)

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

        os.environ["WANDB_LOG_MODEL"] = "end"  # save the model to WandB

        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_training_dataset,
            tokenizer=self.tokenizer
        )

        trainer.train()
        results_labels = []
        results_output = []

        for item in tokenized_eval_dataset:
            input_ids = self.tokenizer(f"{item['input_text']} {self.break_string}", return_tensors="pt", max_length=512, return_attention_mask=True).input_ids.to("cuda:0")
            label = item['target_text']
            outputs = self.model.generate(input_ids, max_new_tokens=30, num_return_sequences=1)
            # Remove the input prompt to get only the generated text
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            prompt_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
            response = response[len(prompt_text):].strip()
            results_output.append(response)
            results_labels.append(label)
        self.compute_metrics_text(results_output, results_labels)
        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 Set": validation_table})
        self.model.generation_config.update(bad_words_ids=get_bad_word_ids())
        self.model.save_pretrained("./gpt2-finetuned-topic")
        self.tokenizer.save_pretrained("./gpt2-finetuned-topic")

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

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