Blame for code/pmps/train.py

# Modified from codi_ori/train.py
# Key changes:
#   1. Import CODI_CT (now PMPS-enhanced) instead of CODI
#   2. Preprocess now computes cot_start_position and cot_end_position
#   3. DataCollator passes CoT-position fields to the model so the PMPS
#      loss can extract teacher-side CoT hidden states.
#   4. CustomTrainer logs pmps_loss and related diagnostics

import copy
import logging
import os
import re
import random
from dataclasses import dataclass, field
from typing import Dict, Optional, Sequence
import torch
import json
import transformers
from torch.utils.data import Dataset
from transformers import Trainer
from safetensors.torch import load_file
from tqdm import tqdm
from math import ceil
from peft import PeftModel, LoraConfig, TaskType, get_peft_model
from datasets import load_dataset
from functools import partial
import yaml
import argparse
import wandb

from src.model import (
    CODI_CT as CODI_PMPS,
    ModelArguments,
    DataArguments,
    TrainingArguments,
    freeze_model,
)

IGNORE_INDEX = -100
def read_json(file_path):
    """
    从指定路径读取JSON文件并返回对应的Python对象。
    """
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            data = json.load(file)
            return data
    except Exception as e:
        print(f"读取JSON文件时出错: {e}")
        return None

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)


class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, num_items_in_batch):
        step = self.state.global_step

        batch_size = self.args.per_device_train_batch_size
        gradient_accumulation_steps = self.args.gradient_accumulation_steps
        num_epochs = self.args.num_train_epochs
        dataset_size = len(self.train_dataset)

        effective_batch_size = batch_size * self.args.world_size * gradient_accumulation_steps
        total_steps = ceil(dataset_size / effective_batch_size) * num_epochs

        inputs["step_ratio"] = step / total_steps
        inputs["step"] = step
        outputs = model(**inputs)
        
        if self.state.epoch is None:
            raise RuntimeError("self.state.epoch is not set")
        current_epoch = self.state.epoch
        if current_epoch >= self.args.pmps_start_epoch:
            loss = outputs["loss"]
        else:
            loss = outputs["no_pmps_loss"]

        # Store custom metrics for logging at every step
        self._custom_metrics = {
            "train/ce_loss": outputs["ce_loss"],
            "train/distill_loss": outputs["distill_loss"],
            "train/ref_ce_loss": outputs["ref_ce_loss"],
            "train/pmps_loss": outputs.get("pmps_loss", 0),
            "train/pmps_loss_c2l": outputs.get("pmps_loss_c2l", 0),
            "train/pmps_loss_l2c": outputs.get("pmps_loss_l2c", 0),
            "train/code_entropy_latent": outputs.get("code_entropy_latent", 0),
            "train/code_entropy_cot": outputs.get("code_entropy_cot", 0),
            "train/match_entropy_c2l": outputs.get("match_entropy_c2l", 0),
            "train/match_order_correlation": outputs.get("match_order_correlation", 0),
            "train/position_bias_strength": outputs.get("position_bias_strength", 0),
            "train/diversity_loss": outputs.get("diversity_loss", 0),
            "train/latent_avg_cosine_sim": outputs.get("latent_avg_cosine_sim", 0),
            "train/total_loss": loss.item(),
            "train/no_pmps_loss": outputs["no_pmps_loss"].item() if isinstance(outputs["no_pmps_loss"], torch.Tensor) else outputs["no_pmps_loss"],
        }
        return loss

    def log(self, logs, start_time=None):
        # Merge custom metrics into the logs dict so they are all sent
        # to wandb (and other loggers) in a single call
        if hasattr(self, "_custom_metrics") and self._custom_metrics:
            logs.update(self._custom_metrics)
            self._custom_metrics = {}
        super().log(logs)


def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer) -> Dict:
    """Tokenize a list of strings."""
    tokenized_list = [
        tokenizer(
            text,
            return_tensors="pt",
            padding="longest",
            max_length=256,
            truncation=True,
            return_attention_mask=False,
        )
        for text in strings
    ]
    input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
    input_ids_lens = labels_lens = [
        tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item() for tokenized in tokenized_list
    ]
    return dict(
        input_ids=input_ids,
        labels=labels,
        input_ids_lens=input_ids_lens,
        labels_lens=labels_lens,
    )


def extract_answer_number(sentence: str) -> float:
    sentence = sentence.replace(",", "")
    pred = [s for s in re.findall(r"-?\d+\.?\d*", sentence)]
    if not pred:
        return float("inf")
    segment = [sentence]
    if len(segment) > 1:
        pred_answer = segment[1]
        pred_answer = [s for s in re.findall(r"-?\d+\.?\d*", pred_answer)]
        if len(pred_answer) > 0:
            pred_answer = pred_answer[0]
        else:
            pred_answer = float(pred[-1])
    else:
        pred_answer = float(pred[-1])

    if isinstance(pred_answer, str):
        try:
            pred_answer = float(pred_answer)
        except ValueError:
            pred_answer = float("inf")
    return pred_answer


def train():
    arg_parser = argparse.ArgumentParser(description="CODI Training with YAML config support")
    arg_parser.add_argument("--config_file", type=str, default=None, help="Path to YAML config file")
    args, remaining_args = arg_parser.parse_known_args()
    
    # If config_file is provided, load from YAML
    if args.config_file:
        print(f"Loading configuration from YAML file: {args.config_file}")
        with open(args.config_file, 'r') as f:
            config = yaml.safe_load(f)
        
        # Convert YAML config to command line arguments
        cmd_args = []
        for key, value in config.items():
            if value is None:
                continue
            if isinstance(value, bool):
                if value:
                    cmd_args.extend([f"--{key}", "True"])
                else:
                    cmd_args.extend([f"--{key}", "False"])
            else:
                cmd_args.extend([f"--{key}", str(value)])
        
        # Parse with HfArgumentParser
        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args)
    else:
        # Original behavior: parse from command line
        print("Loading configuration from command line arguments")
        parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
        model_args, data_args, training_args = parser.parse_args_into_dataclasses()

    ##########################
    #       Peft Model       #
    ##########################
    if model_args.lora_init:
        task_type = TaskType.CAUSAL_LM
        if any(name in model_args.model_name_or_path.lower() for name in ["llama", "mistral", "falcon", "qwen"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj", "gate_proj"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["phi"]):
            target_modules = ["q_proj", "k_proj", "v_proj", "dense", "fc1", "fc2"]
        elif any(name in model_args.model_name_or_path.lower() for name in ["gpt2"]):
            target_modules = ["c_attn", "c_proj", "c_fc"]
        else:
            raise ValueError(f"Only support LLAMA, Mistral, Falcon, Phi-2, but got {model_args.model_name_or_path}.")

        lora_config = LoraConfig(
            task_type=task_type,
            inference_mode=False,
            r=model_args.lora_r,
            lora_alpha=model_args.lora_alpha,
            lora_dropout=0.1,
            target_modules=target_modules,
            init_lora_weights=True,
        )

    model = CODI_PMPS(model_args, training_args, lora_config)
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_args.model_name_or_path,
        token=model_args.token,
        cache_dir=training_args.cache_dir,
        model_max_length=training_args.model_max_length,
        padding_side="right",
    )

    if tokenizer.pad_token_id is None:
        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
        tokenizer.pad_token_id = model.pad_token_id
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids("[PAD]")

    def get_answer_token_position(tokens, answer_prompts, tokenizer):
        try:
            match_indices = (
                (tokens.unfold(0, len(answer_prompts[0]), 1) == answer_prompts[0])
                .all(dim=1)
                .nonzero(as_tuple=True)[0]
                .item()
            )
            answer_token_id = match_indices + len(answer_prompts[0])
            return answer_token_id
        except Exception:
            breakpoint()

    def preprocess(
        sources: Sequence[str],
        targets: Sequence[str],
        answers: Sequence[str],
        tokenizer: transformers.PreTrainedTokenizer,
        bot_id: int,
        eot_id: int,
    ) -> Dict:
        """
        Preprocess data for CODI_CT training.

        Key difference from codi_ori: we also compute cot_start_position and
        cot_end_position to identify the CoT token range within ref_input_ids.
        These positions are used to extract CoT hidden states for CT loss.
        """
        print("Tokenizing inputs... This may take some time...")
        sources_id = _tokenize_fn(sources, tokenizer)["input_ids"]
        cot_id = _tokenize_fn(targets, tokenizer)["input_ids"]
        answers_id = _tokenize_fn(answers, tokenizer)["input_ids"]

        # Add eos token to accommodate pretrained model's format
        if not training_args.remove_eos:
            sources_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in sources_id]
            cot_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in cot_id]
        answers_id = [torch.tensor(x.numpy().tolist() + [tokenizer.eos_token_id], dtype=torch.long) for x in answers_id]
        

        if cot_id[0][0] == tokenizer.bos_token_id:
            cot_id = [x[1:] for x in cot_id]
            answers_id = [x[1:] for x in answers_id]

        ref_input_ids = [
            torch.cat([x, y, z]).to(torch.long) for x, y, z in zip(sources_id, cot_id, answers_id)
        ]
        ref_labels = []
        for x, y in zip(ref_input_ids, sources_id):
            z = x.clone()
            z[: len(y)] = -100
            ref_labels.append(z)

        # ===== NEW: Compute CoT position range in ref_input_ids =====
        # ref_input_ids = [question_tokens | cot_tokens | answer_tokens]
        # cot_start = len(question_tokens), cot_end = len(question_tokens) + len(cot_tokens)
        cot_start_position = [len(src) for src in sources_id]
        cot_end_position = [len(src) + len(cot) for src, cot in zip(sources_id, cot_id)]

        # Add eot to source
        sources_id = [
            torch.tensor(x.numpy().tolist() + [bot_id], dtype=torch.long) for x in sources_id
        ]
        # Add eot and eos
        if training_args.remove_eos:
            answers_id = [
                torch.tensor([eot_id] + x.numpy().tolist(), dtype=torch.long) for x in answers_id
            ]
        else:
            answers_id = [
                torch.tensor([eot_id, tokenizer.eos_token_id] + x.numpy().tolist(), dtype=torch.long)
                for x in answers_id
            ]

        answer_prompts = [
            torch.tensor(tokenizer.encode("The answer is:")),
            torch.tensor(tokenizer.encode("The next step result is:")),
        ]
        if answer_prompts[0][0] == tokenizer.bos_token_id:
            answer_prompts[0] = answer_prompts[0][1:]
            answer_prompts[1] = answer_prompts[1][1:]

        ref_answer_position = [
            get_answer_token_position(x, answer_prompts, tokenizer) for i, x in enumerate(ref_input_ids)
        ]
        model_answer_position = [
            get_answer_token_position(x, answer_prompts, tokenizer) for x in answers_id
        ]

        ref_eos_position = [len(x) - 1 for x in ref_input_ids]
        model_eos_position = [len(x) - 1 for x in answers_id]

        return dict(
            encoder_input_ids=sources_id,
            decoder_input_ids=answers_id,
            ref_input_ids=ref_input_ids,
            labels=answers_id,
            ref_answer_position=ref_answer_position,
            model_answer_position=model_answer_position,
            ref_eos_position=ref_eos_position,
            model_eos_position=model_eos_position,
            ref_labels=ref_labels,
            cot_start_position=cot_start_position,  # NEW
            cot_end_position=cot_end_position,  # NEW
        )

    class SupervisedDataset(Dataset):
        QUESTION_PROMPT = "\nAnswer the above question. First think step by step and then answer the final number.\n"
        QUESTION_DA_PROMPT = "\nAnswer the above question. Answer the final number directly in one number.\n"

        def __init__(self, data_name, raw_data, tokenizer, bot, eot):
            super(SupervisedDataset, self).__init__()
            logging.warning("Formatting inputs...")

            self.data_name = data_name
            questions, cots, answers = [], [], []
            if "math" in self.data_name:
                print("loading MATH training data")
                raw_data = read_json('/data/oss_bucket_0/liqi/dataset/MATH/train_clean.json')
                # raw_data is already loaded from data_path in make_supervised_data_module
            elif "icot-full" in self.data_name:
                print("loading gsm8k-aug-nl training data")
                raw_data = read_json('/data/oss_bucket_0/liqi/dataset/gsm8k-aug-nl/train.json')
            else:
                print("loading gsm9k-aug training data")
                raw_data = read_json('/data/oss_bucket_0/liqi/gsk8k-aug/gsm_train.json')   
               
            for num_iter, example in enumerate(raw_data):
                if training_args.exp_mode and num_iter > training_args.exp_data_num:
                    break
                
                if "icot" in self.data_name and "full" in self.data_name: # icot-full (GSM8k-Aug-NL)
                    question = f"{example['question']}"
                    if example["answer"] is None:
                        continue
                    example['cot'] = ' '.join(example['steps'])
                    token_num = len(tokenizer.encode(example["question"] + example["cot"] + example["answer"]))
                    if token_num > training_args.max_token_num:
                        continue
                    cot = f"{example['cot']}".split(". ")
                    if not training_args.include_last_cot:
                        cot = cot[:-1]
                    answer = example["answer"].split(" ")[-1]
                    if not answer[0].isdigit():
                        continue
                    answer = f"The answer is: {answer}"
                    answer = answer.replace("####", "")
                    questions.append(question)
                    if cot:
                        cot = ". ".join(cot) + ".\n"
                    else:
                        cot = ""
                    cots.append(cot)
                    answers.append(answer)

                elif "icot" in self.data_name: # icot (GSM8k-Aug)
                    question = f"{example['question']}"
                    if "cot" not in example:
                        example["cot"] = example["steps"]
                        example['cot'] = ' '.join(example['cot'])
                    token_num = len(tokenizer.encode(example["question"] + example["cot"] + example["answer"]))
                    if token_num > training_args.max_token_num:
                        continue

                    cot = f"{example['cot']}".split(" ")
                    if not training_args.include_last_cot:
                        cot = cot[:-1]
                    # if len(cot) == 0:
                    #     continue

                    answer = example['answer'].split(' ')[-1]
                    
                    # some answers startwith the negative sign (-), bringing distillation problems for LLaMA
                    if not answer[0].isdigit():
                        continue
                    answer = f"The answer is: {answer}" 
                    answer = answer.replace("####", "")
                    questions.append(question)
                    cots.append(" ".join(cot))
                    answers.append(answer)

                elif "commonsense" in self.data_name or "strategy" in self.data_name:
                    question = example["question"].strip() + "\n"
                    cot = example["cot"].strip() + "\n"
                    answer = f"The answer is: {str(example['answer']).strip()}"
                    token_num = len(tokenizer.encode(question + " " + cot + " " + answer))
                    if token_num > training_args.max_token_num:
                        continue
                    questions.append(question)
                    cots.append(cot)
                    answers.append(answer)

                elif "prontoqa" in data_args.data_name:
                    question = example["question"].strip() + "\n"
                    cot = "\n".join(example["steps"][:-1]) + "\n"
                    answer = f"The answer is: {str(example['answer']).strip()}"
                    token_num = len(tokenizer.encode(question + " " + cot + " " + answer))
                    if token_num > training_args.max_token_num:
                        continue
                    questions.append(question)
                    cots.append(cot)
                    answers.append(answer)

                elif "math" in self.data_name:
                    # MATH dataset: problem -> question, solution -> full CoT, answer -> final answer
                    question = example["problem"].strip() + "\n"
                    cot = example["solution"].strip() + "\n"  # Keep full CoT, no step deletion
                    answer = f"The answer is: {str(example['answer']).strip()}"
                    token_num = len(tokenizer.encode(question + " " + cot + " " + answer))
                    if token_num > training_args.max_token_num:
                        continue
                    questions.append(question)
                    cots.append(cot)
                    answers.append(answer)
                else:
                    raise NotImplementedError

            if training_args.exp_mode:
                questions = questions[: training_args.exp_data_num]
                cots = cots[: training_args.exp_data_num]
                answers = answers[: training_args.exp_data_num]

            print(f"{len(cots)} data in total...")
            logging.warning("Tokenizing inputs... This may take some time...")

            self.data_dict = preprocess(questions, cots, answers, tokenizer, bot, eot)
            self.keys = list(self.data_dict.keys())

        def __len__(self):
            return len(self.data_dict["encoder_input_ids"])

        def __getitem__(self, i) -> Dict[str, torch.Tensor]:
            return {key: self.data_dict[key][i] for key in self.keys}

    @dataclass
    class DataCollatorForSupervisedDataset(object):
        """Collate examples for supervised fine-tuning with CT support."""

        tokenizer: transformers.PreTrainedTokenizer

        def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
            (
                encoder_input_ids,
                decoder_input_ids,
                ref_input_ids,
                labels,
                ref_answer_position,
                model_answer_position,
                ref_labels,
                cot_start_position,
                cot_end_position,
            ) = tuple(
                [instance[key] for instance in instances]
                for key in (
                    "encoder_input_ids",
                    "decoder_input_ids",
                    "ref_input_ids",
                    "labels",
                    "ref_answer_position",
                    "model_answer_position",
                    "ref_labels",
                    "cot_start_position",
                    "cot_end_position",
                )
            )

            # Pad left for encoder
            reversed_input_ids = [seq.flip(0) for seq in encoder_input_ids]
            encoder_input_ids = torch.nn.utils.rnn.pad_sequence(
                reversed_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
            ).flip(1)

            # Pad right for ref and decoder
            ref_input_ids = torch.nn.utils.rnn.pad_sequence(
                ref_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
            )
            ref_labels = torch.nn.utils.rnn.pad_sequence(
                ref_labels, batch_first=True, padding_value=IGNORE_INDEX
            )
            decoder_input_ids = torch.nn.utils.rnn.pad_sequence(
                decoder_input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
            )
            labels = torch.nn.utils.rnn.pad_sequence(
                labels, batch_first=True, padding_value=IGNORE_INDEX
            )

            return dict(
                encoder_input_ids=encoder_input_ids,
                decoder_input_ids=decoder_input_ids,
                ref_input_ids=ref_input_ids,
                labels=labels,
                encoder_attention_mask=encoder_input_ids.ne(self.tokenizer.pad_token_id),
                ref_answer_position=torch.tensor(ref_answer_position, dtype=torch.long),
                model_answer_position=torch.tensor(model_answer_position, dtype=torch.long),
                ref_attention_mask=ref_input_ids.ne(self.tokenizer.pad_token_id),
                ref_labels=ref_labels,
                cot_start_position=torch.tensor(cot_start_position, dtype=torch.long),  # NEW
                cot_end_position=torch.tensor(cot_end_position, dtype=torch.long),  # NEW
            )

    def make_supervised_data_module(tokenizer, data_args) -> Dict:
        """Make dataset and collator for supervised fine-tuning."""
        logging.warning("Downloading Data")
        if "icot" in data_args.data_name:
            if "full" in data_args.data_name:
                dataset = None
            else:
                dataset = None
            train_dataset = SupervisedDataset(
                data_name=data_args.data_name,
                raw_data=dataset,
                tokenizer=tokenizer,
                bot=model.bot_id,
                eot=model.eot_id,
            )
            data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
            return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
        elif "strategy" in data_args.data_name:
            dataset = load_dataset("zen-E/StrategyQA_CoT_GPT4o")["train"]
            train_dataset = SupervisedDataset(
                data_name=data_args.data_name,
                raw_data=dataset,
                tokenizer=tokenizer,
                bot=model.bot_id,
                eot=model.eot_id,
            )
            data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
            return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
        elif "commonsense" in data_args.data_name:
            dataset = load_dataset("zen-E/CommonsenseQA-GPT4omini")["train"]
            train_dataset = SupervisedDataset(
                data_name=data_args.data_name,
                raw_data=dataset,
                tokenizer=tokenizer,
                bot=model.bot_id,
                eot=model.eot_id,
            )
            data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
            return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
        elif "prontoqa" in data_args.data_name:
            with open("/home/ubuntu/coconut/data/prontoqa_train.json") as f:
                dataset = json.load(f)
            train_dataset = SupervisedDataset(
                data_name=data_args.data_name,
                raw_data=dataset,
                tokenizer=tokenizer,
                bot=model.bot_id,
                eot=model.eot_id,
            )
            data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
            return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
        elif "math" in data_args.data_name:
            # assert data_args.data_path is not None, "data_path must be provided for MATH dataset"
            # dataset = read_json(data_args.data_path)
            dataset = None
            train_dataset = SupervisedDataset(
                data_name=data_args.data_name,
                raw_data=dataset,
                tokenizer=tokenizer,
                bot=model.bot_id,
                eot=model.eot_id,
            )
            data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
            return dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
        else:
            raise NotImplementedError(f"Dataset {data_args.data_name} is not supported.")

    training_args.output_dir = os.path.join(
        training_args.output_dir,
        training_args.expt_name
    )

    # ===== Save all input arguments to a JSON file =====
    os.makedirs(training_args.output_dir, exist_ok=True)
    from dataclasses import asdict
    all_args = {
        "model_args": asdict(model_args),
        "data_args": asdict(data_args),
        "training_args": {k: v for k, v in asdict(training_args).items()
                         if not k.startswith("_")},
    }
    args_save_path = os.path.join(training_args.output_dir, "training_arguments.json")
    with open(args_save_path, "w", encoding="utf-8") as f:
        json.dump(all_args, f, ensure_ascii=False, indent=4, default=str)
    print(f"All training arguments saved to {args_save_path}")

    # ===== Wandb Initialization =====
    if training_args.use_wandb and int(os.environ.get("RANK", 0)) == 0:
        if training_args.wandb_api_key:
            wandb.login(key=training_args.wandb_api_key)
        run_name = training_args.wandb_run_name or training_args.expt_name
        wandb.init(
            project=training_args.wandb_project,
            name=run_name,
            config={
                "model_name_or_path": model_args.model_name_or_path,
                "lora_r": model_args.lora_r,
                "lora_alpha": model_args.lora_alpha,
                "num_latent": training_args.num_latent,
                "learning_rate": training_args.learning_rate,
                "per_device_train_batch_size": training_args.per_device_train_batch_size,
                "gradient_accumulation_steps": training_args.gradient_accumulation_steps,
                "num_train_epochs": training_args.num_train_epochs,
                "pmps_loss_factor": training_args.pmps_loss_factor,
                "pmps_temperature": training_args.pmps_temperature,
                "pmps_match_temperature": training_args.pmps_match_temperature,
                "pmps_nmb_prototypes": training_args.pmps_nmb_prototypes,
                "pmps_use_projector": training_args.pmps_use_projector,
                "pmps_layer": training_args.pmps_layer,
                "pmps_start_epoch": training_args.pmps_start_epoch,
                "pmps_position_bias_max": training_args.pmps_position_bias_max,
                "pmps_position_bias_schedule": training_args.pmps_position_bias_schedule,
                "diversity_loss_factor": training_args.diversity_loss_factor,
                "distill_loss_factor": training_args.distill_loss_factor,
            },
        )
        training_args.report_to = ["wandb"]
        print("Wandb initialized successfully.")
    elif not training_args.use_wandb:
        os.environ["WANDB_DISABLED"] = "true"

    data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
    trainer = CustomTrainer(model=model, tokenizer=tokenizer, args=training_args, **data_module)
    trainer.train()

    trainer.save_state()
    trainer.save_model(output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()