| |
| |
| |
| |
| |
| |
| |
|
|
| 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"] |
|
|
| |
| 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): |
| |
| |
| 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 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) |
| |
| |
| 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)]) |
| |
| |
| parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) |
| model_args, data_args, training_args = parser.parse_args_into_dataclasses(args=cmd_args) |
| else: |
| |
| print("Loading configuration from command line arguments") |
| parser = transformers.HfArgumentParser((ModelArguments, DataArguments, TrainingArguments)) |
| model_args, data_args, training_args = parser.parse_args_into_dataclasses() |
|
|
| |
| |
| |
| 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"] |
|
|
| |
| 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) |
|
|
| |
| |
| |
| 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)] |
|
|
| |
| sources_id = [ |
| torch.tensor(x.numpy().tolist() + [bot_id], dtype=torch.long) for x in sources_id |
| ] |
| |
| 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, |
| cot_end_position=cot_end_position, |
| ) |
|
|
| 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') |
| |
| 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: |
| 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: |
| 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] |
| |
| |
|
|
| answer = example['answer'].split(' ')[-1] |
| |
| |
| 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: |
| |
| question = example["problem"].strip() + "\n" |
| cot = example["solution"].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) |
| 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", |
| ) |
| ) |
|
|
| |
| 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) |
|
|
| |
| 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), |
| cot_end_position=torch.tensor(cot_end_position, dtype=torch.long), |
| ) |
|
|
| 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: |
| |
| |
| 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 |
| ) |
|
|
| |
| 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}") |
|
|
| |
| 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() |
|
|
