utils.py

import json
import torch
import random
import math
import time
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from collections import OrderedDict
import os

from transformers import Trainer
from transformers.utils import (
    logging,
    is_sagemaker_mp_enabled,
    is_torch_xla_available,
    is_accelerate_available
)
from transformers.trainer_utils import (
    has_length, 
    speed_metrics,
    TrainOutput,
)
from transformers.debug_utils import DebugOption
from transformers.trainer_callback import TrainerState
from transformers.trainer_pt_utils import get_model_param_count
from transformers.integrations.deepspeed import deepspeed_init, deepspeed_load_checkpoint, is_deepspeed_available
from torch.nn.utils.rnn import pad_sequence

from accelerate.utils import DistributedType



logger = logging.get_logger(__name__)

class BaseActor(torch.nn.Module):
    def __init__(self, size):
        super(BaseActor, self).__init__()
        self.size = size
        self.bias = torch.nn.Linear(size, 1, bias=False)
        for p in self.bias.parameters():
            p.data.fill_(1.)

    def forward(self, feature):
        logits = self.bias.weight * feature
        return logits

def balance_probabilities_by_temperature(probabilities, temperature):
    """
    probabilities: dictionary of probabilities
    temperature: temperature to apply to the probabilities
    """
    temperature = float(temperature)
    for k, v in probabilities.items():
        probabilities[k] = v ** (1 / temperature)

    sum_prob = sum(probabilities.values())

    for k, v in probabilities.items():
        probabilities[k] = v / sum_prob
    
    return probabilities

def read_jsonl(path):
    with open(path, "r") as f:
        data = [json.loads(line) for line in f]
    return data

def write_jsonl(data, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w") as f:
        for d in data:
            f.write(json.dumps(d, ensure_ascii=False) + "\n")


class MultiDataset:
    def __init__(
        self, 
        datasets=None,
        category_probabilities=None,
        tokenizer=None,
        prob_log_path=None,
        model_id=None,
        ema_alpha=None,
    ):  

        self.datasets = datasets
        self.category_probabilities = category_probabilities
        self.ema_alpha = ema_alpha

        print("Category Probabilities")
        for k, v in self.category_probabilities.items():
            print(f"{k}: {v}")

        self.sizes = [len(dataset) for dataset in self.datasets.values()]

        self.tokenizer = tokenizer
        self.eos_token = tokenizer.eos_token
        self.eos_token_id = tokenizer.eos_token_id
        self.pad_token = tokenizer.pad_token
        self.pad_token_id = tokenizer.pad_token_id
        self.max_length = min(2048, tokenizer.model_max_length)
        self.prepare_dataset()

        self.category_occurrences = OrderedDict([(k, 0) for k, v in self.datasets.items()])

        self.last_step = 0

        self.prob_log_path = prob_log_path
        
        self.write_one_log(self.last_step, self.category_probabilities)

        self.model_id = model_id

        # ema
        self.prev_rewards = OrderedDict([(k, 1) for k, v in self.datasets.items()])



    def __len__(self):
        return sum([len(dataset) for dataset in self.datasets.values()])

    def __getitem__(self, idx):
        category = random.choices(
            list(self.category_probabilities.keys()), 
            weights=list(self.category_probabilities.values()), 
            k=1
        )[0]
        self.category_occurrences[category] += 1
        dataset = self.datasets[category]
        return random.choice(dataset)
    
    def prepare_dataset(self):
        for category, dataset in self.datasets.items():
            data = self.datasets[category]
            print(f"Formatting {category} ...")
            data = data.map(self.format_text, num_proc=4)
            print(f"Tokenizing {category} ...")
            data = data.map(self.tokenize_fn, num_proc=4)
            print(data)

            self.datasets[category] = data

    def tokenize_fn(self, example):
        inputs = self.tokenizer(example["text"], truncation=True, max_length=2048)
        example["input_ids"] = inputs["input_ids"]
        example["attention_mask"] = inputs["attention_mask"]
        return example

    def format_text(self, example):

        conv = example["conversations"]
        text = self.tokenizer.apply_chat_template(conv, tokenize=False)
        example["text"] = text.strip()

        return example

    def collate_fn(self, batch):

        input_ids = pad_sequence([torch.tensor(x["input_ids"]) for x in batch], batch_first=True, padding_value=self.eos_token_id)
        labels = input_ids.clone().detach()
        attention_mask = pad_sequence([torch.tensor(x["attention_mask"]) for x in batch], batch_first=True, padding_value=0)

        results = {
            "input_ids": input_ids,
            "labels": labels,
            "attention_mask": attention_mask,
        }
        return results

    def sample_batch(self, batch_size, category=None):
        if category is not None:
            dataset = self.datasets[category]
            batch = random.choices(dataset, k=batch_size)
        else:
            batch = [self.__getitem__(i) for i in range(batch_size)]

        return batch

    def update_sampling_probabilities(self, new_probabilities, step, new_rewards=None):

        if step == self.last_step:
            return None

        print("***** Updating Sampling Probabilities *****")
        for x, y, z, k in zip(self.category_probabilities.keys(), self.category_probabilities.values(), new_probabilities.values(), self.category_occurrences.values()):
            print("Category: %20s | Old Probability =  %6.3f%% | New Probability = %6.3f%% | Interval Occurrences = %7d" % (x, y*100, z*100, k))
        self.category_probabilities = new_probabilities
        self.category_occurrences = OrderedDict([(k, 0) for k, v in self.datasets.items()])
        self.last_step = step
        self.write_one_log(step, self.category_probabilities)

        print("***** Update Sampling Probabilities Done *****")

    def write_one_log(self, step, probabilities):

        if self.prob_log_path is not None:
            new_dict = OrderedDict()
            new_dict["step"] = step
            for k, v in probabilities.items():
                new_dict[k] = float(v)
            
            with open(self.prob_log_path, "a") as f:
                f.write(json.dumps(new_dict, ensure_ascii=False) + "\n")


class DynaTrainer(Trainer):

    def __init__(
        self,
        model=None, args=None, data_collator=None, train_dataset=None, eval_dataset=None,
        tokenizer=None, model_init=None, compute_metrics=None,
        callbacks=None, optimizers=(None, None), preprocess_logits_for_metrics=None,
        ######################
        sampling_prob_update_steps=None,
        reward_type=None,
        model_id=None,
        use_ema=None,
        specialist=None,
    ):
        super().__init__(
            model=model, args=args, data_collator=data_collator, train_dataset=train_dataset, 
            eval_dataset=eval_dataset, tokenizer=tokenizer, model_init=model_init, compute_metrics=compute_metrics, 
            callbacks=callbacks, optimizers=optimizers, preprocess_logits_for_metrics=preprocess_logits_for_metrics
        )

        self.num_categories = len(list(train_dataset.category_probabilities.keys()))
        self.data_actor = BaseActor(size=self.num_categories)
        self.data_optimizer = torch.optim.Adam([p for p in self.data_actor.parameters() if p.requires_grad], lr=1e-4)
        self.data_optim_steps = 200
        self.sampling_prob_update_steps=sampling_prob_update_steps
        self.reward_type = reward_type
        self.model_id = model_id
        self.use_ema = use_ema
        self.specialist = specialist


    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):
        self.accelerator.free_memory()
        self._train_batch_size = batch_size
        if self.args.auto_find_batch_size:
            if self.state.train_batch_size != self._train_batch_size:
                from accelerate.utils import release_memory

                (self.model_wrapped,) = release_memory(self.model_wrapped)
                self.model_wrapped = self.model

                # Check for DeepSpeed *after* the intial pass and modify the config
                if self.is_deepspeed_enabled:
                    # Temporarily unset `self.args.train_batch_size`
                    original_bs = self.args.per_device_train_batch_size
                    self.args.per_device_train_batch_size = self._train_batch_size // max(1, self.args.n_gpu)
                    self.propagate_args_to_deepspeed(True)
                    self.args.per_device_train_batch_size = original_bs
            self.state.train_batch_size = self._train_batch_size
        logger.debug(f"Currently training with a batch size of: {self._train_batch_size}")
        # Data loader and number of training steps
        train_dataloader = self.get_train_dataloader()
        if self.is_fsdp_xla_v2_enabled:
            train_dataloader = tpu_spmd_dataloader(train_dataloader)

        # Setting up training control variables:
        # number of training epochs: num_train_epochs
        # number of training steps per epoch: num_update_steps_per_epoch
        # total number of training steps to execute: max_steps
        total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size

        len_dataloader = None
        num_train_tokens = None
        if has_length(train_dataloader):
            len_dataloader = len(train_dataloader)
            num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
            num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
            num_examples = self.num_examples(train_dataloader)
            if args.max_steps > 0:
                max_steps = args.max_steps
                num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
                    args.max_steps % num_update_steps_per_epoch > 0
                )
                # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
                # the best we can do.
                num_train_samples = args.max_steps * total_train_batch_size
                if args.include_tokens_per_second:
                    num_train_tokens = (
                        self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
                    )
            else:
                max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
                num_train_epochs = math.ceil(args.num_train_epochs)
                num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
                if args.include_tokens_per_second:
                    num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs
        elif args.max_steps > 0:  # Rely on max_steps when dataloader does not have a working size
            max_steps = args.max_steps
            # Setting a very large number of epochs so we go as many times as necessary over the iterator.
            num_train_epochs = sys.maxsize
            num_update_steps_per_epoch = max_steps
            num_examples = total_train_batch_size * args.max_steps
            num_train_samples = args.max_steps * total_train_batch_size
            if args.include_tokens_per_second:
                num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps
        else:
            raise ValueError(
                "args.max_steps must be set to a positive value if dataloader does not have a length, was"
                f" {args.max_steps}"
            )

        if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
            if self.args.n_gpu > 1:
                # nn.DataParallel(model) replicates the model, creating new variables and module
                # references registered here no longer work on other gpus, breaking the module
                raise ValueError(
                    "Currently --debug underflow_overflow is not supported under DP. Please use DDP"
                    " (torchrun or torch.distributed.launch (deprecated))."
                )
            else:
                debug_overflow = DebugUnderflowOverflow(self.model)  # noqa

        delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled

        # We need to reset the scheduler, as its parameters may be different on subsequent calls
        if self._created_lr_scheduler:
            self.lr_scheduler = None
            self._created_lr_scheduler = False

        if self.is_deepspeed_enabled:
            self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)

        if not delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        self.state = TrainerState()
        self.state.is_hyper_param_search = trial is not None
        self.state.train_batch_size = self._train_batch_size

        # Compute absolute values for logging, eval, and save if given as ratio
        if args.logging_steps is not None:
            if args.logging_steps < 1:
                self.state.logging_steps = math.ceil(max_steps * args.logging_steps)
            else:
                self.state.logging_steps = args.logging_steps
        if args.eval_steps is not None:
            if args.eval_steps < 1:
                self.state.eval_steps = math.ceil(max_steps * args.eval_steps)
            else:
                self.state.eval_steps = args.eval_steps
        if args.save_steps is not None:
            if args.save_steps < 1:
                self.state.save_steps = math.ceil(max_steps * args.save_steps)
            else:
                self.state.save_steps = args.save_steps

        # Activate gradient checkpointing if needed
        if args.gradient_checkpointing:
            if args.gradient_checkpointing_kwargs is None:
                gradient_checkpointing_kwargs = {}
            else:
                gradient_checkpointing_kwargs = args.gradient_checkpointing_kwargs

            self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

        model = self._wrap_model(self.model_wrapped)

        # as the model is wrapped, don't use `accelerator.prepare`
        # this is for unhandled cases such as
        # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX
        use_accelerator_prepare = True if model is self.model else False

        if delay_optimizer_creation:
            if use_accelerator_prepare:
                self._fsdp_qlora_plugin_updates()
                self.model = self.accelerator.prepare(self.model)
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        # prepare using `accelerator` prepare
        if use_accelerator_prepare:
            self.model.train()
            if hasattr(self.lr_scheduler, "step"):
                if self.use_apex:
                    model = self.accelerator.prepare(self.model)
                else:
                    model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
            else:
                # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config.
                model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
                    self.model, self.optimizer, self.lr_scheduler
                )

        if self.is_fsdp_enabled:
            self.model = self.model_wrapped = model

        # for the rest of this function `model` is the outside model, whether it was wrapped or not
        if model is not self.model:
            self.model_wrapped = model

        # backward compatibility
        if self.is_deepspeed_enabled:
            self.deepspeed = self.model_wrapped

        # ckpt loading
        if resume_from_checkpoint is not None:
            if self.is_deepspeed_enabled:
                deepspeed_load_checkpoint(
                    self.model_wrapped, resume_from_checkpoint, load_module_strict=not _is_peft_model(self.model)
                )
            elif is_sagemaker_mp_enabled() or self.is_fsdp_enabled:
                self._load_from_checkpoint(resume_from_checkpoint, self.model_wrapped)

        # Check if saved optimizer or scheduler states exist
        self._load_optimizer_and_scheduler(resume_from_checkpoint)

        # important: at this point:
        # self.model         is the Transformers Model
        # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model),
        # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc.

        # Train!
        print("***** Running training *****")
        print(f"  Num examples = {num_examples:,}")
        print(f"  Num Epochs = {num_train_epochs:,}")
        print(f"  Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
        if self.args.per_device_train_batch_size != self._train_batch_size:
            print(f"  Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}")
        print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}")
        print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        print(f"  Total optimization steps = {max_steps:,}")
        print(f"  Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}")

        ###########################################

        self.pretrain_data_actor()

        ###########################################


        self.state.epoch = 0
        start_time = time.time()
        epochs_trained = 0
        steps_trained_in_current_epoch = 0
        steps_trained_progress_bar = None

        # Check if continuing training from a checkpoint
        if resume_from_checkpoint is not None and os.path.isfile(
            os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
        ):
            self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
            self.compare_trainer_and_checkpoint_args(self.args, self.state)
            epochs_trained = self.state.global_step // num_update_steps_per_epoch
            if not args.ignore_data_skip:
                steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
                steps_trained_in_current_epoch *= args.gradient_accumulation_steps
            else:
                steps_trained_in_current_epoch = 0

            print("  Continuing training from checkpoint, will skip to saved global_step")
            print(f"  Continuing training from epoch {epochs_trained}")
            print(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                print(
                    f"  Will skip the first {epochs_trained} epochs then the first"
                    f" {steps_trained_in_current_epoch} batches in the first epoch."
                )

        # Update the references
        self.callback_handler.model = self.model
        self.callback_handler.optimizer = self.optimizer
        self.callback_handler.lr_scheduler = self.lr_scheduler
        self.callback_handler.train_dataloader = train_dataloader
        if self.hp_name is not None and self._trial is not None:
            # use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial
            # parameter to Train when using DDP.
            self.state.trial_name = self.hp_name(self._trial)
        if trial is not None:
            assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial
            self.state.trial_params = hp_params(assignments)
        else:
            self.state.trial_params = None
        # This should be the same if the state has been saved but in case the training arguments changed, it's safer
        # to set this after the load.
        self.state.max_steps = max_steps
        self.state.num_train_epochs = num_train_epochs
        self.state.is_local_process_zero = self.is_local_process_zero()
        self.state.is_world_process_zero = self.is_world_process_zero()

        # tr_loss is a tensor to avoid synchronization of TPUs through .item()
        tr_loss = torch.tensor(0.0).to(args.device)
        # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses
        self._total_loss_scalar = 0.0
        self._globalstep_last_logged = self.state.global_step
        model.zero_grad()
        grad_norm: Optional[float] = None

        self.control = self.callback_handler.on_train_begin(args, self.state, self.control)

        # Skip the first epochs_trained epochs to get the random state of the dataloader at the right point.
        if not args.ignore_data_skip:
            for epoch in range(epochs_trained):
                sampler = get_dataloader_sampler(train_dataloader)
                sampler_kinds = [RandomSampler]
                if version.parse(accelerate_version) > version.parse("0.23.0"):
                    sampler_kinds.append(SeedableRandomSampler)
                is_random_sampler = isinstance(sampler, tuple(sampler_kinds))
                if not is_random_sampler:
                    # We just need to begin an iteration to create the randomization of the sampler.
                    for _ in train_dataloader:
                        break
                else:
                    # Otherwise we need to call the whooooole sampler cause there is some random operation added
                    # AT THE VERY END!
                    sampler = sampler if sampler is not None else []
                    _ = list(sampler)

        total_batched_samples = 0
        for epoch in range(epochs_trained, num_train_epochs):
            epoch_iterator = train_dataloader
            if hasattr(epoch_iterator, "set_epoch"):
                epoch_iterator.set_epoch(epoch)

            # Reset the past mems state at the beginning of each epoch if necessary.
            if args.past_index >= 0:
                self._past = None

            steps_in_epoch = (
                len(epoch_iterator)
                if len_dataloader is not None
                else args.max_steps * args.gradient_accumulation_steps
            )
            self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)

            if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
                self._load_rng_state(resume_from_checkpoint)

            rng_to_sync = False
            steps_skipped = 0
            if steps_trained_in_current_epoch > 0:
                epoch_iterator = skip_first_batches(epoch_iterator, steps_trained_in_current_epoch)
                steps_skipped = steps_trained_in_current_epoch
                steps_trained_in_current_epoch = 0
                rng_to_sync = True

            step = -1
            for step, inputs in enumerate(epoch_iterator):
                total_batched_samples += 1

                if self.args.include_num_input_tokens_seen:
                    main_input_name = getattr(self.model, "main_input_name", "input_ids")
                    if main_input_name not in inputs:
                        logger.warning(
                            "Tried to track the number of tokens seen, however the current model is "
                            "not configured properly to know what item is the input. To fix this, add "
                            "a `main_input_name` attribute to the model class you are using."
                        )
                    else:
                        input_device = inputs[main_input_name].device
                        self.state.num_input_tokens_seen += torch.sum(
                            self.accelerator.gather(
                                torch.tensor(inputs[main_input_name].numel(), device=input_device, dtype=torch.int64)
                            )
                        ).item()
                if rng_to_sync:
                    self._load_rng_state(resume_from_checkpoint)
                    rng_to_sync = False

                # Skip past any already trained steps if resuming training
                if steps_trained_in_current_epoch > 0:
                    steps_trained_in_current_epoch -= 1
                    if steps_trained_progress_bar is not None:
                        steps_trained_progress_bar.update(1)
                    if steps_trained_in_current_epoch == 0:
                        self._load_rng_state(resume_from_checkpoint)
                    continue
                elif steps_trained_progress_bar is not None:
                    steps_trained_progress_bar.close()
                    steps_trained_progress_bar = None

                if step % args.gradient_accumulation_steps == 0:
                    self.control = self.callback_handler.on_step_begin(args, self.state, self.control)

                with self.accelerator.accumulate(model):
                    tr_loss_step = self.training_step(model, inputs)

                if (
                    args.logging_nan_inf_filter
                    and not is_torch_xla_available()
                    and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
                ):
                    # if loss is nan or inf simply add the average of previous logged losses
                    tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
                else:
                    if tr_loss.device != tr_loss_step.device:
                        raise ValueError(
                            f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}"
                        )
                    tr_loss += tr_loss_step

                self.current_flos += float(self.floating_point_ops(inputs))

                is_last_step_and_steps_less_than_grad_acc = (
                    steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch
                )

                if (
                    total_batched_samples % args.gradient_accumulation_steps == 0
                    or
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
                    is_last_step_and_steps_less_than_grad_acc
                ):
                    # the `or` condition of `is_last_step_and_steps_less_than_grad_acc` is not covered
                    # in accelerate. So, explicitly enable sync gradients to True in that case.
                    if is_last_step_and_steps_less_than_grad_acc:
                        self.accelerator.gradient_state._set_sync_gradients(True)

                    # Gradient clipping
                    if args.max_grad_norm is not None and args.max_grad_norm > 0:
                        # deepspeed does its own clipping

                        if is_sagemaker_mp_enabled() and args.fp16:
                            _grad_norm = self.optimizer.clip_master_grads(args.max_grad_norm)
                        elif self.use_apex:
                            # Revert to normal clipping otherwise, handling Apex or full precision
                            _grad_norm = nn.utils.clip_grad_norm_(
                                amp.master_params(self.optimizer),
                                args.max_grad_norm,
                            )
                        else:
                            _grad_norm = self.accelerator.clip_grad_norm_(
                                model.parameters(),
                                args.max_grad_norm,
                            )

                        if (
                            is_accelerate_available()
                            and self.accelerator.distributed_type == DistributedType.DEEPSPEED
                        ):
                            grad_norm = model.get_global_grad_norm()
                            # In some cases the grad norm may not return a float
                            if hasattr(grad_norm, "item"):
                                grad_norm = grad_norm.item()
                        else:
                            grad_norm = _grad_norm

                    # Optimizer step
                    self.optimizer.step()
                    optimizer_was_run = not self.accelerator.optimizer_step_was_skipped
                    if optimizer_was_run:
                        # Delay optimizer scheduling until metrics are generated
                        if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
                            self.lr_scheduler.step()

                    model.zero_grad()
                    self.state.global_step += 1
                    self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)

                    self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
                else:
                    self.control = self.callback_handler.on_substep_end(args, self.state, self.control)


                ##################################

                if (self.state.global_step != 0 and self.state.global_step % self.sampling_prob_update_steps == 0):
                    new_category_probabilities = self.update_dataset_probability(model, self.state.global_step)
                    self.train_dataset.update_sampling_probabilities(new_category_probabilities, self.state.global_step)

                ##################################

                if self.control.should_epoch_stop or self.control.should_training_stop:
                    # PyTorch/XLA relies on the data loader to insert the mark_step for
                    # each step. Since we are breaking the loop early, we need to manually
                    # insert the mark_step here.
                    if is_torch_xla_available():
                        xm.mark_step()
                    break
            if step < 0:
                logger.warning(
                    "There seems to be not a single sample in your epoch_iterator, stopping training at step"
                    f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
                    f" num_steps ({max_steps}) higher than the number of available samples."
                )
                self.control.should_training_stop = True

            self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
            self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)

            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
                if is_torch_xla_available():
                    # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                    xm.master_print(met.metrics_report())
                else:
                    logger.warning(
                        "You enabled PyTorch/XLA debug metrics but you don't have a TPU "
                        "configured. Check your training configuration if this is unexpected."
                    )
            if self.control.should_training_stop:
                break

        if args.past_index and hasattr(self, "_past"):
            # Clean the state at the end of training
            delattr(self, "_past")

        print("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
        if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
            # Wait for everyone to get here so we are sure the model has been saved by process 0.
            if is_torch_xla_available():
                xm.rendezvous("load_best_model_at_end")
            elif args.parallel_mode == ParallelMode.DISTRIBUTED:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            self._load_best_model()

        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
        effective_global_step = max(self.state.global_step, 0.001)  # Avoid ZeroDivisionError
        train_loss = self._total_loss_scalar / effective_global_step

        metrics = speed_metrics(
            "train",
            start_time,
            num_samples=num_train_samples,
            num_steps=self.state.max_steps,
            num_tokens=num_train_tokens,
        )
        self.store_flos()
        metrics["total_flos"] = self.state.total_flos
        metrics["train_loss"] = train_loss

        self.is_in_train = False

        self._memory_tracker.stop_and_update_metrics(metrics)

        self.log(metrics)

        run_dir = self._get_output_dir(trial)
        checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)

        # Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save.
        if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            for checkpoint in checkpoints_sorted:
                if not os.path.samefile(checkpoint, self.state.best_model_checkpoint):
                    print(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
                    shutil.rmtree(checkpoint)

        self.control = self.callback_handler.on_train_end(args, self.state, self.control)

        # Wait for the checkpoint to be uploaded.
        self._finish_current_push()

        # After training we make sure to retrieve back the original forward pass method
        # for the embedding layer by removing the forward post hook.
        if self.neftune_noise_alpha is not None:
            self._deactivate_neftune(self.model)

        return TrainOutput(self.state.global_step, train_loss, metrics)

    def pretrain_data_actor(self):
        print("***** Pretrain Data Actor *****")
        feature =  torch.ones(1, self.num_categories)
        target = torch.tensor([list(self.train_dataset.category_probabilities.values())])
        l = 100

        count = 0
        while l > 1e-10:
            self.data_actor.zero_grad()
            self.data_optimizer.zero_grad()
            a_logits = self.data_actor(feature)
            prob = torch.nn.functional.softmax(a_logits, dim=-1)
            loss = torch.nn.functional.mse_loss(prob, target)
            l = loss.item()
            if count % 1000 == 0 :
                print("Pretrain Data Actor | Loss = %.10f | num_updates = %10d" % (l, count))
            loss.backward()
            self.data_optimizer.step()
            count += 1

        with torch.no_grad():
            a_logits = self.data_actor(feature)
            prob = torch.nn.functional.softmax(a_logits, dim=-1)
            sim_list = [i for i in prob.data.view(-1).cpu().numpy()]

            for x, y, z in zip(self.train_dataset.category_probabilities.keys(), self.train_dataset.category_probabilities.values(), sim_list):
                print("Pretrained Data Actor | Category: %20s | Analytical Probability =  %6.3f%% | Numerical Probability = %6.3f%% " % (x, y*100, z*100))

        print("***** Pretrain Data Actor Done *****")

    def update_dataset_probability(self, model, step):

        if step == self.train_dataset.last_step:
            return self.train_dataset.category_probabilities

        torch.cuda.empty_cache()
        model.eval()

        all_reward_list = OrderedDict()
        if self.reward_type == "cossim":
            cate_embeds = OrderedDict()
            for category in self.train_dataset.category_probabilities.keys():
                torch.cuda.empty_cache()
                batch_embed = 0
                for i in range(8):
                    sample = self.train_dataset.sample_batch(8, category)
                    batch_embed += self.compute_batch_embedding(model, sample)
                cate_embeds[category] = batch_embed / 8
            
            for category in self.train_dataset.category_probabilities.keys():
                all_reward_list[category] = 0
                if self.specialist not in cate_embeds.keys():
                    print("Specialist not found in category embeddings")
                    for k, v in cate_embeds.items():
                        # if k == category:
                        #     continue
                        sim = self.compute_cossim(cate_embeds[category], v)
                        # print(category, k, sim)
                        all_reward_list[category] += sim
                    all_reward_list[category] /= (len(cate_embeds))
                else:
                    print(f"Specialist {self.specialist} found in category embeddings")
                    sim = self.compute_cossim(cate_embeds[category], cate_embeds[self.specialist])
                    all_reward_list[category] = sim

        elif self.reward_type == "diff":
            for category in self.train_dataset.category_probabilities.keys():
                r = 0
                for i in range(8):
                    torch.cuda.empty_cache()
                    sample = self.train_dataset.sample_batch(2, category)
                    if self.reward_type == "diff":
                        r += self.compute_learn_reward(model, sample)
                    else:
                        raise NotImplementedError(f"Reward type {self.reward_type} not implemented")
                
                all_reward_list[category] = r / 8
            if self.specialist in all_reward_list.keys():
                print(f"Specialist {self.specialist} found in category embeddings")
                all_reward_list[self.specialist ] = all_reward_list[self.specialist ] * 2

        else:
            raise NotImplementedError(f"Reward type {self.reward_type} not implemented")

        print("***** Reward List *****")
        if self.use_ema == "yes":
            for k, v in all_reward_list.items():
                self.train_dataset.prev_rewards[k] = self.train_dataset.ema_alpha * v + (1 - self.train_dataset.ema_alpha) * self.train_dataset.prev_rewards[k]
                print(f"{k} : before: {v} after : {self.train_dataset.prev_rewards[k]}")
            all_reward_list = self.train_dataset.prev_rewards

        print(all_reward_list)
        feature = torch.ones(1, self.num_categories)
        grad_scale = torch.tensor(list(all_reward_list.values())).view(1, -1)
        for _ in range(self.data_optim_steps):
            self.data_actor.zero_grad()
            self.data_optimizer.zero_grad()
            a_logits = self.data_actor(feature)
            loss = -torch.nn.functional.log_softmax(a_logits, dim=-1) * grad_scale
            loss = loss.sum()
            loss.backward()
            self.data_optimizer.step()

        with torch.no_grad():
            a_logits = self.data_actor(feature)
            probs = torch.nn.functional.softmax(a_logits, dim=-1)
            probs = [i for i in probs.data.view(-1).cpu().numpy()]
        
        category_probabilities = OrderedDict([(k, v) for k, v in zip(self.train_dataset.category_probabilities.keys(), probs)])
        
        model.train()
        return category_probabilities


    def compute_batch_embedding(self, model, sample):

        lst = []
        for example in sample:
            input_ids = torch.tensor([example["input_ids"]]).to(model.device)
            attention_mask = torch.tensor([example["attention_mask"]]).to(model.device)
            labels = input_ids.clone().detach().to(model.device)
            with torch.no_grad():
                outputs = model(
                    input_ids=input_ids, 
                    attention_mask=attention_mask, 
                    labels=labels, 
                    output_hidden_states=True
                )
                hidden_states = outputs.hidden_states
                last_hidden_state = hidden_states[-1]
                example_embed = last_hidden_state.mean(dim=1)
                if torch.isnan(example_embed).any():
                    # print("NaN detected")
                    continue
                lst.append(example_embed)
        
        batch_embed = torch.cat(lst, dim=0).mean(dim=0)
        return batch_embed


    def compute_cossim(self, embed_a, embed_b):
        cosine_sim = torch.nn.functional.cosine_similarity(embed_a, embed_b, dim=-1)
        return cosine_sim.item()

    def compute_ppl_reward(self, model, sample):
        lst = []
        for example in sample:
            ppl = self.compute_ppl(example, model)
            lst.append(ppl)
        return sum(lst) / len(lst)

    def compute_learn_reward(self, model, sample):
        lst = []
        for example in sample:
            ppl = self.compute_ppl(example, model)
            score = ppl / example[f"{self.model_id}_ppl"]
            lst.append(score)
        return sum(lst) / len(lst)

    def compute_ppl(self, example, model):
        input_ids = torch.tensor([example["input_ids"]]).long().to(model.device)
        attention_mask = torch.tensor([example["attention_mask"]]).long().to(model.device)
        labels = input_ids.clone().detach().to(model.device)
        with torch.no_grad():
            loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels).loss
        return float(torch.exp(loss).item())

# from transformers import AutoTokenizer, AutoModelForCausalLM
# import glob
# from datasets import load_dataset
# tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2")
# model = AutoModelForCausalLM.from_pretrained("microsoft/phi-2")

# paths = glob.glob("/home/minghaow/balancesft-project/datasets_mini/*.jsonl")
# raw_datasets = {}
# for p in paths:
#     print(f"Loading {p}")
#     data = load_dataset("json", data_files=p)["train"]
#     category = p.split("/")[-1].split(".")[0]
#     raw_datasets[category] = data
# init_probs = {k: len(v) / sum([len(v) for v in raw_datasets.values()]) for k, v in raw_datasets.items()}
# rebalanced_probs = balance_probabilities_by_temperature(init_probs, 1)
# train_dataset = MultiDataset(raw_datasets, rebalanced_probs, tokenizer)

# print(train_dataset[0])
# print(train_dataset.sample_batch(4))

# out = model(**train_dataset.sample_batch(4))
# print(out)