# coding=utf-8
# Copyright 2020-present the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
The Trainer class, to easily train a 🤗 Transformers from scratch or finetune it on a new task.
"""

import sys
sys.path.append('./')

import contextlib
import functools
import glob
import inspect
import math
import os
import random
import re
import shutil
import sys
import time
import warnings
from collections.abc import Mapping
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
import copy
from .metrics import f1
import numpy as np
from collections import OrderedDict

from tqdm.auto import tqdm
from transformers import Trainer
from sklearn.linear_model import LinearRegression, LogisticRegression, LogisticRegressionCV

# Integrations must be imported before ML frameworks:
from transformers.integrations import (  # isort: split
    # default_hp_search_backend,
    get_reporting_integration_callbacks,
    hp_params,
    # is_fairscale_available,
    is_optuna_available,
    is_ray_tune_available,
    is_sigopt_available,
    is_wandb_available,
    run_hp_search_optuna,
    run_hp_search_ray,
    run_hp_search_sigopt,
    run_hp_search_wandb,
)

import numpy as np
import torch
import torch.distributed as dist
from packaging import version
from torch import nn
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
from torch.utils.data.distributed import DistributedSampler
import importlib.metadata
from huggingface_hub import Repository

from transformers import __version__
from transformers.configuration_utils import PretrainedConfig
from transformers.data.data_collator import DataCollator, DataCollatorWithPadding, default_data_collator
from transformers.debug_utils import DebugOption, DebugUnderflowOverflow
from transformers.deepspeed import deepspeed_init, is_deepspeed_zero3_enabled
from transformers.dependency_versions_check import dep_version_check
from transformers.modelcard import TrainingSummary
from transformers.modeling_utils import PreTrainedModel, load_sharded_checkpoint, unwrap_model
from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES, MODEL_MAPPING_NAMES
from transformers.optimization import Adafactor, get_scheduler
# from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS, is_torch_greater_or_equal_than_1_10, is_torch_less_than_1_11
from transformers.tokenization_utils_base import PreTrainedTokenizerBase
from transformers.trainer_callback import (
    CallbackHandler,
    DefaultFlowCallback,
    PrinterCallback,
    ProgressCallback,
    TrainerCallback,
    TrainerControl,
    TrainerState,
)
from transformers.trainer_pt_utils import (
    DistributedLengthGroupedSampler,
    DistributedSamplerWithLoop,
    DistributedTensorGatherer,
    IterableDatasetShard,
    LabelSmoother,
    LengthGroupedSampler,
    SequentialDistributedSampler,
    ShardSampler,
    distributed_broadcast_scalars,
    distributed_concat,
    find_batch_size,
    get_module_class_from_name,
    get_parameter_names,
    nested_concat,
    nested_detach,
    nested_numpify,
    nested_truncate,
    nested_xla_mesh_reduce,
    reissue_pt_warnings,
)
from transformers.trainer_utils import (
    PREFIX_CHECKPOINT_DIR,
    BestRun,
    EvalLoopOutput,
    EvalPrediction,
    FSDPOption,
    HPSearchBackend,
    HubStrategy,
    IntervalStrategy,
    PredictionOutput,
    RemoveColumnsCollator,
    # ShardedDDPOption,
    TrainerMemoryTracker,
    TrainOutput,
    default_compute_objective,
    denumpify_detensorize,
    enable_full_determinism,
    find_executable_batch_size,
    get_last_checkpoint,
    has_length,
    number_of_arguments,
    seed_worker,
    set_seed,
    speed_metrics,
)
from transformers.training_args import OptimizerNames, ParallelMode, TrainingArguments
from transformers.utils import (
    CONFIG_NAME,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
    find_labels,
    get_full_repo_name,
    is_apex_available,
    is_datasets_available,
    is_in_notebook,
    is_peft_available,
    is_ipex_available,
    is_sagemaker_dp_enabled,
    is_sagemaker_mp_enabled,
    is_torch_tensorrt_fx_available,
    is_torch_tpu_available,
    is_torchdynamo_available,
    logging,
)
from transformers.utils.generic import ContextManagers

#
from automatic_prompt_engineer import evaluate


# _is_native_cpu_amp_available = is_torch_greater_or_equal_than_1_10

DEFAULT_CALLBACKS = [DefaultFlowCallback]
DEFAULT_PROGRESS_CALLBACK = ProgressCallback

if is_in_notebook():
    from .utils.notebook import NotebookProgressCallback

    DEFAULT_PROGRESS_CALLBACK = NotebookProgressCallback

if is_apex_available():
    from apex import amp

if is_datasets_available():
    import datasets

if is_torch_tpu_available(check_device=False):
    import torch_xla.core.xla_model as xm
    import torch_xla.debug.metrics as met
    import torch_xla.distributed.parallel_loader as pl

# if is_fairscale_available():
#     dep_version_check("fairscale")
#     import fairscale
#     from fairscale.nn.data_parallel import FullyShardedDataParallel as FullyShardedDDP
#     from fairscale.nn.data_parallel import ShardedDataParallel as ShardedDDP
#     from fairscale.nn.wrap import auto_wrap
#     from fairscale.optim import OSS
#     from fairscale.optim.grad_scaler import ShardedGradScaler


if is_sagemaker_mp_enabled():
    import smdistributed.modelparallel.torch as smp
    from smdistributed.modelparallel import __version__ as SMP_VERSION

    IS_SAGEMAKER_MP_POST_1_10 = version.parse(SMP_VERSION) >= version.parse("1.10")

    from .trainer_pt_utils import smp_forward_backward, smp_forward_only, smp_gather, smp_nested_concat
else:
    IS_SAGEMAKER_MP_POST_1_10 = False

if is_peft_available():
    from peft import PeftModel
    
def _is_peft_model(model):
    if is_peft_available():
        classes_to_check = (PeftModel,) if is_peft_available() else ()
        # Here we also check if the model is an instance of `PeftMixedModel` introduced in peft>=0.7.0: https://github.com/huggingface/transformers/pull/28321
        if version.parse(importlib.metadata.version("peft")) >= version.parse("0.7.0"):
            from peft import PeftMixedModel

            classes_to_check = (*classes_to_check, PeftMixedModel)
        return isinstance(model, classes_to_check)
    return False

if TYPE_CHECKING:
    import optuna

logger = logging.get_logger(__name__)


# Name of the files used for checkpointing
TRAINING_ARGS_NAME = "training_args.bin"
TRAINER_STATE_NAME = "trainer_state.json"
OPTIMIZER_NAME = "optimizer.pt"
SCHEDULER_NAME = "scheduler.pt"
SCALER_NAME = "scaler.pt"


torch.autograd.set_detect_anomaly(True)

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

class AIO_Trainer(Trainer):

    from transformers.trainer_pt_utils import _get_learning_rate, log_metrics, metrics_format, save_metrics, save_state

    def _inner_training_loop(
        self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
    ):
        """
        We overload the original training loop to add linear probing and MeZO. Search key word "MeZO added"
        for those updates.
        """
        self._train_batch_size = batch_size
        # Data loader and number of training steps
        train_dataloader = self.get_train_dataloader()

        # MeZO added: Linear probing
        if self.args.linear_probing:

            def _get_token_prediction_layer(model):
                if model.config.model_type == "opt":
                    return model.lm_head
                else:
                    raise NotImplementedError(model.config.model_type)

            def _extract_features(model, *args, **kwargs):
                """some magic for getting features pre last layer"""
                features = {}
                def __hook(model_, input_, output_):
                    features["features"] = input_[0].detach()

                _get_token_prediction_layer(model).register_forward_hook(__hook)
                model.forward(*args, **kwargs)
                return features["features"]

            logger.info("Linear probing")
            logger.info("Starting to get features for training dataset")
            targets = []
            features = []
            with torch.inference_mode():
                for step, inputs in enumerate(tqdm(train_dataloader)):
                    for k, v in inputs.items():
                        if isinstance(v, torch.Tensor):
                            inputs[k] = v.to(self.model.device)
                        
                    feature = _extract_features(self.model, **inputs)
                    target = inputs["labels"]

                    # Shift the target (bc it's autoregressive LM) and add the corresponding part
                    assert not self.args.train_as_classification and self.args.only_train_option
                    feature, target = feature[:, :-1], target[:, 1:]
                    for _i, _len in enumerate(inputs["option_len"]):
                        features.append(feature[_i, -_len:])
                        targets.append(target[_i, -_len:])

            logger.info("Finished getting features for training dataset")

            features = torch.cat(features, dim=0).cpu().numpy()
            targets = torch.cat(targets, dim=0).cpu().numpy()
            # Whether to use bias
            if self.model.config.model_type in ["opt", "gpt2"]:
                use_bias = False
            else:
                raise NotImplementedError
            # Set early stopping
            tol = 0.01 if self.args.lp_early_stopping else 1e-4 # 1e-4 is scipy default
            max_iter = 1000 if self.args.lp_early_stopping else 5000

            logger.info("Fitting logistic regression...")
            reg = LogisticRegressionCV(max_iter=max_iter, fit_intercept=use_bias, multi_class="multinomial", random_state=0, tol=tol, n_jobs=-1).fit(features, targets)
            logger.info("Done")

            logger.info("Assigning weights to model")
            decoder = _get_token_prediction_layer(self.model)
            coef_torch = torch.tensor(reg.coef_, device=decoder.weight.device, dtype=decoder.weight.dtype)
            if use_bias:
                bias_torch = torch.tensor(reg.intercept_, device=decoder.weight.device, dtype=decoder.weight.dtype)
            if coef_torch.shape[0] == 1: # The regressor only detects two classes
                assert len(reg.classes_) == 2
                coef_torch = torch.cat([-coef_torch / 2, coef_torch / 2], dim=0)
                if use_bias:
                    bias_torch = torch.cat([-bias_torch / 2, bias_torch / 2], dim=0)

            for _i, token_id in enumerate(reg.classes_):
                decoder.weight.data[token_id] = coef_torch[_i]
                if use_bias:
                    decoder.bias.data[token_id] = bias_torch[_i]

            return None

        # 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 = args.train_batch_size * args.gradient_accumulation_steps * args.world_size

        len_dataloader = 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)
            # args.max_steps is default as -1
            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
            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
        #
        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
        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"
                    " (torch.distributed.launch)."
                )
            else:
                debug_overflow = DebugUnderflowOverflow(self.model)  # noqa

        delay_optimizer_creation = (
            # self.sharded_ddp is not None
            # and self.sharded_ddp != ShardedDDPOption.SIMPLE
            # or is_sagemaker_mp_enabled()
            is_sagemaker_mp_enabled()
            # or self.fsdp is not None
        )
        if args.deepspeed:
            deepspeed_engine, optimizer, lr_scheduler = deepspeed_init(
                self, num_training_steps=max_steps, resume_from_checkpoint=resume_from_checkpoint
            )
            self.model = deepspeed_engine.module
            self.model_wrapped = deepspeed_engine
            self.deepspeed = deepspeed_engine
            self.optimizer = optimizer
            self.lr_scheduler = lr_scheduler
        
        elif 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

        # Activate gradient checkpointing if needed
        if args.gradient_checkpointing:
            self.model.gradient_checkpointing_enable()

        model = self._wrap_model(self.model_wrapped)

        if is_sagemaker_mp_enabled() and resume_from_checkpoint is not None:
            self._load_from_checkpoint(resume_from_checkpoint, 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

        if delay_optimizer_creation:
            self.create_optimizer_and_scheduler(num_training_steps=max_steps)

        # 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), etc.

        # Train!
        logger.info("***** Running training *****")
        logger.info(f"  Num examples = {num_examples}")
        logger.info(f"  Num Epochs = {num_train_epochs}")
        logger.info(f"  Instantaneous batch size per device = {args.per_device_train_batch_size}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size}")
        logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
        logger.info(f"  Total optimization stesps = {max_steps}")
        # logger.info(
        #     f"  Number of parameters that require grad = {sum(p.numel() for _, p in self.named_parameters if p.requires_grad)}"
        # )

        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))
            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

            logger.info("  Continuing training from checkpoint, will skip to saved global_step")
            logger.info(f"  Continuing training from epoch {epochs_trained}")
            logger.info(f"  Continuing training from global step {self.state.global_step}")
            if not args.ignore_data_skip:
                logger.info(
                    f"  Will skip the first {epochs_trained} epochs then the first {steps_trained_in_current_epoch} "
                    "batches in the first epoch. If this takes a lot of time, you can add the `--ignore_data_skip` "
                    "flag to your launch command, but you will resume the training on data already seen by your model."
                )
                if self.is_local_process_zero() and not args.disable_tqdm:
                    steps_trained_progress_bar = tqdm(total=steps_trained_in_current_epoch)
                    steps_trained_progress_bar.set_description("Skipping the first batches")

        # 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()

        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):
                is_random_sampler = hasattr(train_dataloader, "sampler") and isinstance(
                    train_dataloader.sampler, RandomSampler
                )
                if is_torch_less_than_1_11 or not is_random_sampler:
                    # We just need to begin an iteration to create the randomization of the sampler.
                    # That was before PyTorch 1.11 however...
                    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!
                    _ = list(train_dataloader.sampler)

        ################################################  Iterating over epochs ################################################
        for epoch in range(epochs_trained, num_train_epochs):
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
                train_dataloader.sampler.set_epoch(epoch)
            elif hasattr(train_dataloader, "dataset") and isinstance(train_dataloader.dataset, IterableDatasetShard):
                train_dataloader.dataset.set_epoch(epoch)

            if is_torch_tpu_available():
                parallel_loader = pl.ParallelLoader(train_dataloader, [args.device]).per_device_loader(args.device)
                epoch_iterator = parallel_loader
            else:
                epoch_iterator = train_dataloader

            # 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)

            ##########################################  Iterating over steps ############################################
            step = -1
            for step, inputs in enumerate(epoch_iterator):

                # 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)

                # MeZO added: estimate gradient
                if args.trainer == "zo":
                    self.update_random_seeds()
                    tr_loss_step = self.zo_step(model, inputs, epoch, step)

                # ---------------------------------------------------------------------------
                
                elif args.trainer == "aio_hybrid":
                    # self.optimizer = torch.optim.SGD(model.parameters(), lr = self.args.learning_rate)
                    # self.optimizer, self.lr_scheduler = None, None
                    # self.create_optimizer_and_scheduler(num_training_steps=max_steps)
                    #
                    if (
                        ((step + 1) % args.gradient_accumulation_steps != 0)
                        and args.local_rank != -1
                        and args._no_sync_in_gradient_accumulation
                    ):
                        # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
                        with model.no_sync():
                            tr_loss_step = self.AIO_training_step(model, inputs, epoch, step)
                    # Perform backward()
                    else:
                        tr_loss_step = self.AIO_training_step(model, inputs, epoch, step)
                
                    #
                    tr_loss_step = torch.tensor(tr_loss_step.item()).to(args.device)

                # ---------------------------------------------------------------------------
                
                else:
                    if (
                        ((step + 1) % args.gradient_accumulation_steps != 0)
                        and args.local_rank != -1
                        and args._no_sync_in_gradient_accumulation
                    ):
                        # Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
                        with model.no_sync():
                            tr_loss_step = self.training_step(model, inputs)
                    # Perform backward()
                    else:
                        tr_loss_step = self.training_step(model, inputs)

                # ---------------------------------
                if (
                    args.logging_nan_inf_filter
                    and not is_torch_tpu_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:
                    tr_loss += tr_loss_step

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

                # Optimizer step for deepspeed must be called on every step regardless of the value of gradient_accumulation_steps
                if self.deepspeed:
                    self.deepspeed.step()

                #
                self.do_grad_scaling = False

                if (step + 1) % args.gradient_accumulation_steps == 0 or (
                    # last step in epoch but step is always smaller than gradient_accumulation_steps
                    steps_in_epoch <= args.gradient_accumulation_steps
                    and (step + 1) == steps_in_epoch
                ):
                    # MeZO added: update model with the estimated gradient
                    if args.trainer == "zo":
                        self.zo_update(model, epoch, step)
                    else:
                        # Gradient clipping
                        if args.max_grad_norm is not None and args.max_grad_norm > 0 and not self.deepspeed:

                            ###
                            self.do_grad_scaling = self.args.do_grad_scaling
                            # deepspeed does its own clipping

                            if self.do_grad_scaling:
                                # Reduce gradients first for XLA
                                if is_torch_tpu_available():
                                    gradients = xm._fetch_gradients(self.optimizer)
                                    xm.all_reduce("sum", gradients, scale=1.0 / xm.xrt_world_size())
                                # AMP: gradients need unscaling
                                self.scaler.unscale_(self.optimizer)

                            if is_sagemaker_mp_enabled() and args.fp16:
                                self.optimizer.clip_master_grads(args.max_grad_norm)
                            elif hasattr(self.optimizer, "clip_grad_norm"):
                                # Some optimizers (like the sharded optimizer) have a specific way to do gradient clipping
                                self.optimizer.clip_grad_norm(args.max_grad_norm)
                            elif hasattr(model, "clip_grad_norm_"):
                                # Some models (like FullyShardedDDP) have a specific way to do gradient clipping
                                model.clip_grad_norm_(args.max_grad_norm)
                            else:
                                # Revert to normal clipping otherwise, handling Apex or full precision
                                nn.utils.clip_grad_norm_(
                                    amp.master_params(self.optimizer) if self.use_apex else model.parameters(),
                                    args.max_grad_norm,
                                )

                        # Optimizer step
                        optimizer_was_run = True
                        ###
                        if self.deepspeed:
                            pass  # called outside the loop
                        elif is_torch_tpu_available():
                            if self.do_grad_scaling:
                                self.scaler.step(self.optimizer)
                                self.scaler.update()
                            else:
                                xm.optimizer_step(self.optimizer)
                        elif self.do_grad_scaling:
                            scale_before = self.scaler.get_scale()
                            self.scaler.step(self.optimizer)
                            self.scaler.update()
                            scale_after = self.scaler.get_scale()
                            optimizer_was_run = scale_before <= scale_after
                        else:
                            # for param in model.parameters():
                            #     assert torch.isfinite(param.grad).all()
                            #     assert torch.isfinite(param).all()
                            for name, param in model.named_parameters():
                                if param.requires_grad:
                                    if param.grad is not None:
                                        print(f"Before. Name: {name}, grad_norm: {param.grad.data.norm(2)}, param_norm:{param.data.norm(2)}. Dtype: {param.dtype}")
                                    else:
                                        print(f"Before. Name: {name}, None GRAD! Dtype: {param.dtype}")
                            
                            # Update parameters
                            self.optimizer.step()
                            
                            #
                            for name, param in model.named_parameters():
                                if param.requires_grad:
                                    if param.grad is not None:
                                        print(f"After. Name: {name}, grad_norm: {param.grad.data.norm(2)}, param_norm:{param.data.norm(2)}. Dtype: {param.dtype}")
                                    else:
                                        print(f"After. Name: {name}, None GRAD! Dtype: {param.dtype}")
                            print(f"Perform optim. step: epoch {epoch}, step {step}")

                        if optimizer_was_run and not self.deepspeed:
                            self.lr_scheduler.step()
                        model.zero_grad()

                    ######################### Evaluation after training step #########################
                    this_sample_seed = int(str(epoch) + str(step))
                    this_eval_loss, instruction = self.AIO_forward_for_loss(model, inputs=None, sample_seed=this_sample_seed, evaluation_flag=True)
                    print("#"*30)
                    print(f"--- EVALUATION --- [Epoch]: {epoch}, [Step]: {step} --- Loss Eval: {this_eval_loss}")
                    print("#"*30)

                    ####
                    self.AIO_forward_model.check_if_update_best_training_instruction(this_eval_loss.item(), instruction, epoch, step)
                    
                    # -------------------------------------------------------------------------------------
                    self.state.global_step += 1
                    self.state.epoch = epoch + (step + 1) / steps_in_epoch
                    self.control = self.callback_handler.on_step_end(args, self.state, self.control)

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

                if self.control.should_epoch_stop or self.control.should_training_stop:
                    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=tr_loss, grad_norm=None, model=model, trial=trial, epoch=epoch, 
                                            ignore_keys_for_eval=ignore_keys_for_eval)

            if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
                if is_torch_tpu_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

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

        logger.info("\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 sur the model has been saved by process 0.
            if is_torch_tpu_available():
                xm.rendezvous("load_best_model_at_end")
            elif args.local_rank != -1:
                dist.barrier()
            elif is_sagemaker_mp_enabled():
                smp.barrier()

            self._load_best_model()

        # add remaining tr_loss
        self._total_loss_scalar += tr_loss.item()
        train_loss = self._total_loss_scalar / self.state.global_step

        metrics = speed_metrics("train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps)
        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.
        if self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
            for checkpoint in checkpoints_sorted:
                if checkpoint != self.state.best_model_checkpoint:
                    logger.info(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)

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


    ################################################################################################################################################
    ################################################ Original MeZO Funcs ###########################################################################
    ################################################################################################################################################

    def update_random_seeds(self):
        # Sample the random seed for sampling z
        if self.args.zo_approx_steps > 1:
            self.zo_random_seed_list = np.random.randint(1000000000, size=(self.args.zo_approx_steps, ))
        else:
            self.zo_random_seed = np.random.randint(1000000000)

    def zo_perturb_parameters(self, random_seed=None, scaling_factor=1):
        """
        Perturb the parameters with random vector z.
        Input: 
        - random_seed: random seed for MeZO in-place perturbation (if it's None, we will use self.zo_random_seed)
        - scaling_factor: theta = theta + scaling_factor * z * eps
        """

        # Set the random seed to ensure that we sample the same z for perturbation/update
        torch.manual_seed(random_seed if random_seed is not None else self.zo_random_seed)
        
        for name, param in self.named_parameters_to_optim:
            z = torch.normal(mean=0, std=1, size=param.data.size(), device=param.data.device, dtype=param.data.dtype)
            param.data = param.data + scaling_factor * z * self.args.zo_eps


    def zo_forward(self, model, inputs):
        """
        Get (no gradient) loss from the model. Dropout is turned off too.
        """
        model.eval()
        if self.args.non_diff:
            # Non-differentiable objective (may require autoregressive generation)
            return self.zo_forward_nondiff(model, inputs)

        with torch.inference_mode():
            inputs = self._prepare_inputs(inputs)
            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)
            if self.args.n_gpu > 1:
                # Warning: this is copied from the original Huggingface Trainer. Untested.
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
        return loss.detach()


    def zo_forward_nondiff(self, model, inputs):
        """
        Get (no gradient) non-diffiable loss from the model.
        """
        model.eval()
        assert self.args.task_name == "SQuAD", "Non differentiable objective only supports SQuAD for now."

        with torch.inference_mode():
            inputs = self._prepare_inputs(inputs)
            args = self.args
            outputs = self.model.generate(
                inputs["input_ids"], do_sample=args.sampling, temperature=args.temperature, 
                num_beams=args.num_beams, top_p=args.top_p, top_k=args.top_k, max_new_tokens=min(args.max_new_tokens, args.max_length - inputs["input_ids"].size(1)), 
                num_return_sequences=1, eos_token_id=[self.tokenizer.encode(args.eos_token, add_special_tokens=False)[-1], self.tokenizer.eos_token_id],
            )
            output_text = []
            for i in range(len(outputs)):
                output_text.append(self.tokenizer.decode(outputs[i][inputs["input_ids"].size(1):], skip_special_tokens=True).strip())
            f1s = [f1(output_text[i], inputs['gold'][i]) for i in range(len(output_text))]
        
        return -torch.tensor(np.mean(f1s), dtype=torch.float32)


    def zo_step(self, model, inputs, epoch, step):
        """
        Estimate gradient by MeZO. Return the loss from f(theta + z)
        """
        # Sample seed to ensure consistency between loss1 and loss2
        sample_seed = int(str(epoch) + str(step))

        # What parameters to optimize 
        self.named_parameters_to_optim = []
        #
        if self.args.only_tune_linear_header:
            for name, param in model.named_parameters():
                if param.requires_grad and name == 'lm_head.weight':
                    self.named_parameters_to_optim.append((name, param))
        #
        else:
            for name, param in model.named_parameters():
                if param.requires_grad:
                    # print("[Name]: ", name)
                    self.named_parameters_to_optim.append((name, param))
        logger.warning(f"Number of trainable parameters = {sum(p.numel() for _, p in self.named_parameters_to_optim if p.requires_grad)}")

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

        if self.args.zo_approx_steps > 1:
            self.self_projected_grad_array = torch.zeros(self.args.zo_approx_steps, dtype=torch.float32)
            #
            for this_step in range(self.args.zo_approx_steps):
                # set random seed
                self.zo_random_seed = self.zo_random_seed_list[this_step]

                # First function evaluation
                self.zo_perturb_parameters(scaling_factor=1)
                
                # Get Loss
                loss1, _ = self.AIO_forward_for_loss(model, inputs=None, sample_seed=sample_seed)
                print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss 1: {loss1}")

                # Second function evaluation
                self.zo_perturb_parameters(scaling_factor=-2)
                # Get loss
                loss2, _ = self.AIO_forward_for_loss(model, inputs=None, sample_seed=sample_seed)
                print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss 2: {loss2}")

                #
                this_projected_grad = ((loss1 - loss2) / (2 * self.args.zo_eps)).item()
                self.self_projected_grad_array[this_step] = this_projected_grad

                # No gradient accumulation support
                assert self.args.gradient_accumulation_steps == 1

                # Reset model back to its parameters at start of step
                self.zo_perturb_parameters(scaling_factor=1)
        
            return loss1
        #
        else:
            # First function evaluation
            self.zo_perturb_parameters(scaling_factor=1)
            
            # Get Loss
            # loss1 = self.zo_forward(model, inputs)
            #
            loss1, _ = self.AIO_forward_for_loss(model, inputs=None, sample_seed=sample_seed)
            print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss 1: {loss1}")

            # Second function evaluation
            self.zo_perturb_parameters(scaling_factor=-2)
            # Get loss
            # loss2 = self.zo_forward(model, inputs)
            #
            loss2, _ = self.AIO_forward_for_loss(model, inputs=None, sample_seed=sample_seed)
            print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss 2: {loss2}")

            self.projected_grad = ((loss1 - loss2) / (2 * self.args.zo_eps)).item()

            # No gradient accumulation support
            assert self.args.gradient_accumulation_steps == 1

            # Reset model back to its parameters at start of step
            self.zo_perturb_parameters(scaling_factor=1)
        
            return loss1


    def zo_update(self, model, epoch, step):
        """
        Update the parameters with the estimated gradients.
        """
        args = self.args
  
        #
        this_lr = self._get_learning_rate()
        print(f"[Epoch]: {epoch}, [Step]: {step} --- Learning rate: {this_lr}")

        if self.args.zo_approx_steps > 1:
            for random_seed_idx in range(self.args.zo_approx_steps):
                #
                this_seed = self.zo_random_seed_list[random_seed_idx]
                # Reset the random seed for sampling zs
                torch.manual_seed(this_seed)
                #
                for name, param in self.named_parameters_to_optim:
                    # Resample z
                    z = torch.normal(mean=0, std=1, size=param.data.size(), device=param.data.device, dtype=param.data.dtype)
                    this_projected_grad = self.self_projected_grad_array[random_seed_idx]
                    #
                    if "bias" not in name and "layer_norm" not in name and "layernorm" not in name:
                        param.data = param.data - (this_lr / self.args.zo_approx_steps) * (this_projected_grad * z + args.weight_decay * param.data)
                    else:
                        param.data = param.data - (this_lr / self.args.zo_approx_steps) * (this_projected_grad * z)
        else:
            # Reset the random seed for sampling zs
            torch.manual_seed(self.zo_random_seed)  
            for name, param in self.named_parameters_to_optim:
                # Resample z
                z = torch.normal(mean=0, std=1, size=param.data.size(), device=param.data.device, dtype=param.data.dtype)
                if "bias" not in name and "layer_norm" not in name and "layernorm" not in name:
                    param.data = param.data - this_lr * (self.projected_grad * z + args.weight_decay * param.data)
                else:
                    param.data = param.data - this_lr * (self.projected_grad * z)

        self.lr_scheduler.step()

    ################################################################################################################################################
    ################################################ AIO Funcs #####################################################################################
    ################################################################################################################################################

    def add_AIO_forward_model_attribute(self, AIO_forward_model):
        self.AIO_forward_model = AIO_forward_model

    def AIO_forward_for_loss(self, model, inputs=None, sample_seed=None, evaluation_flag=False):
        forward_model = self.AIO_forward_model
        #
        input_ids = forward_model.white_box_LLM_tokenizer(forward_model.init_token, return_tensors="pt").input_ids.cuda()
        input_embed = forward_model.white_box_LLM_embedding[input_ids]
        
        #
        if self.args.soft_prompt_tuning or self.args.prefix_tuning:
            output_instance = model.generate(input_ids=input_ids, max_new_tokens=2000, 
                                         pad_token_id=forward_model.white_box_LLM_tokenizer.eos_token_id, 
                                         return_dict_in_generate=True, output_hidden_states=True, 
                                         output_scores=True, output_logits=True)
        else:
            output_instance = model.generate(inputs_embeds=input_embed, max_new_tokens=2000, 
                                            pad_token_id=forward_model.white_box_LLM_tokenizer.eos_token_id, 
                                            return_dict_in_generate=True, output_hidden_states=True, 
                                            output_scores=True, output_logits=True)
        generated_output = output_instance.sequences
        
        #
        instruction = forward_model.white_box_LLM_tokenizer.batch_decode(generated_output, skip_special_tokens=True)
        
        ################################################

        print('[Init Instruction]: {}'.format(instruction))
        #
        matches = re.findall(forward_model.re_pattern, instruction[0], re.DOTALL)
        instruction[0] = ' '.join(matches)
        #
        print('[Pre-re Instruction]: {}'.format(instruction))

        instruction[0] = forward_model.added_instruction_header + instruction[0]
        start = instruction[0].find(forward_model.added_instruction_header)
        end = instruction[0].find('Comment: ')
        #
        if end == -1:
            instruction[0] = instruction[0][start:]
        else:
            instruction[0] = instruction[0][start: end]
        print('[Post-re Instruction]: {}'.format(instruction))
        
        ###############
        _, scores = evaluate.evaluate_prompts(instruction, forward_model.eval_template, forward_model.eval_data, 
                                                forward_model.demos_template, forward_model.few_shot_data, 
                                                forward_model.conf['evaluation']['method'], forward_model.conf['evaluation'], 
                                                eval_LLM=forward_model.eval_LLM_model, 
                                                eval_LLM_tokenizer=forward_model.eval_LLM_tokenizer,
                                                sample_seed=sample_seed,
                                                print_details_flag=True)
        # : Label smooth?
        loss = -1 * torch.tensor(np.mean(scores), dtype=torch.float32)

        ####
        if self.args.include_perturbation_instructions_flag:
            self.AIO_forward_model.check_if_update_best_training_instruction(loss.item(), instruction)

        ##### Evaluate and update TS model checkpoints #####
        if evaluation_flag and self.args.TS_aided_grad_approx:
            with torch.no_grad():
                inputs = {'input_ids': input_ids}
                outputs = model(**inputs)

                #
                if isinstance(outputs, dict) and "logits" not in outputs:
                    raise ValueError(
                        "The model did not return a loss from the inputs, only the following keys: "
                        f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
                    )
                logits = outputs["logits"] if isinstance(outputs, dict) else outputs[0]
                #
                mean_logits = logits.mean(dim=0).mean(dim=0).cpu().detach().view(-1, )
                
                # 
                checkpoint_res = self.AIO_forward_model.TS_model.check_and_update_checkpoint(these_parameters=mean_logits)
                if checkpoint_res:
                    print("="*30)
                    print("REFRESH TS CHECKPOINT!!")
                    print("="*30)

        return loss, instruction
    
    #######################################################################################################
    # ------------------------- Functions for Hybrid AIO Training -----------------------------------------

    def AIO_training_step(self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],
                          epoch: int, step: int) -> torch.Tensor:
        """
        Perform a training step on a batch of inputs.

        Subclass and override to inject custom behavior.

        Args:
            model (`nn.Module`):
                The model to train.
            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
                The inputs and targets of the model.

                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
                argument `labels`. Check your model's documentation for all accepted arguments.

        Return:
            `torch.Tensor`: The tensor with training loss on this batch.

        ### After injecting the proxy weight matrix, used for run the back propagation. ###
        """
        
        #
        model.eval()
        #
        this_random_seed = np.random.randint(1000000000)
        white_box_LLM = self.AIO_forward_model.white_box_LLM
        epsilon_val = self.args.zo_eps

        ####################################
        if self.args.zo_approx_steps > 1:
            loss_diff = torch.zeros((self.args.zo_approx_steps, ), dtype=torch.float32)
            #
            white_box_LLM.refresh_gaussian_vec(random_seed=this_random_seed)
            #
            for approx_step_idx in range(self.args.zo_approx_steps):
                #################### Positive perturbation + Get the evaluation loss
                white_box_LLM.manually_update_gaussian_vec(gaussian_vec=white_box_LLM.random_gaussian_m[approx_step_idx, :])
                
                # With positive perturbation
                white_box_LLM.set_epsilon_val_and_direction(epsilon_val=epsilon_val, perturb_direction=1)
                #
                white_box_LLM.set_logits_perturbation_flag(True)
                white_box_LLM.set_apply_black_box_proxy_m_flag(False)
                #
                loss_pos, _ = self.AIO_forward_for_loss(model, inputs, sample_seed=this_random_seed)
                print(f"[Idx]: {approx_step_idx}. [Epoch]: {epoch}, [Step]: {step} --- Loss Pos: {loss_pos}")

                #################### Negative perturbation + Get the evaluation loss
                
                # With negative perturbation
                white_box_LLM.set_epsilon_val_and_direction(epsilon_val=epsilon_val, perturb_direction=-1)
                #
                white_box_LLM.set_logits_perturbation_flag(True)
                white_box_LLM.set_apply_black_box_proxy_m_flag(False)
                #
                loss_neg, _ = self.AIO_forward_for_loss(model, inputs, sample_seed=this_random_seed)
                print(f"[Idx]: {approx_step_idx}. [Epoch]: {epoch}, [Step]: {step} --- Loss Neg: {loss_neg}")

                #
                this_loss_diff_val = (loss_pos - loss_neg).item()
                loss_diff[approx_step_idx] = this_loss_diff_val
                print(f"[Idx]: {approx_step_idx}. [Loss_diff Val]: ", this_loss_diff_val)

                # Update records for TS-aided model
                if self.args.TS_aided_grad_approx:
                    hidden_size = self.AIO_forward_model.white_box_LLM.vocab_size
                    # this_arm_context = white_box_LLM.random_gaussian_m[approx_step_idx, :hidden_size].view(-1, ).cpu().detach().numpy()
                    this_arm_context = white_box_LLM.random_gaussian_m[approx_step_idx, :hidden_size].view(1, -1).detach().cpu()
                    #
                    if self.args.TS_single_direction_reward and self.AIO_forward_model.last_eval_loss_val is not None:
                        prev_loss = self.AIO_forward_model.last_eval_loss_val
                        pos_rwd = (loss_pos.item() - prev_loss) / self.args.zo_eps
                        neg_rwd = (loss_neg.item() - prev_loss) / self.args.zo_eps
                        #
                        TS_model = self.AIO_forward_model.TS_model
                        #
                        TS_model.store_info(x=this_arm_context, y=pos_rwd)
                        TS_model.store_info(x=(-1*this_arm_context), y=neg_rwd)
                        #
                        # TS_model.store_info(x=(-1*this_arm_context), y=pos_rwd)
                        # TS_model.store_info(x=this_arm_context, y=neg_rwd)
                    else:
                        this_reward = (loss_pos - loss_neg).item() / (2 * self.args.zo_eps)
                        #
                        TS_model = self.AIO_forward_model.TS_model
                        TS_model.store_info(x=this_arm_context, y=this_reward)
        else:
            #### Positive perturbation + Get the evaluation loss
            #
            white_box_LLM.refresh_gaussian_vec(random_seed=this_random_seed)
            # With positive perturbation
            white_box_LLM.set_epsilon_val_and_direction(epsilon_val=epsilon_val, perturb_direction=1)
            #
            white_box_LLM.set_logits_perturbation_flag(True)
            white_box_LLM.set_apply_black_box_proxy_m_flag(False)
            #
            loss_pos, _ = self.AIO_forward_for_loss(model, inputs, sample_seed=this_random_seed)
            print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss Pos: {loss_pos}")

            #### Negative perturbation + Get the evaluation loss
            # With negative perturbation
            white_box_LLM.set_epsilon_val_and_direction(epsilon_val=epsilon_val, perturb_direction=-1)
            #
            white_box_LLM.set_logits_perturbation_flag(True)
            white_box_LLM.set_apply_black_box_proxy_m_flag(False)
            #
            loss_neg, _ = self.AIO_forward_for_loss(model, inputs, sample_seed=this_random_seed)
            print(f"[Epoch]: {epoch}, [Step]: {step} --- Loss Neg: {loss_neg}")

            #
            loss_diff = (loss_pos - loss_neg).item()
            print("[Loss_diff Val]: ", loss_diff)

            ### : check the radius and refresh contexts
            # Update records for TS-aided model
            if self.args.TS_aided_grad_approx:
                hidden_size = self.AIO_forward_model.white_box_LLM.vocab_size
                # this_arm_context = white_box_LLM.random_gaussian_vec[:hidden_size].view(-1, ).cpu().detach().numpy()
                this_arm_context = white_box_LLM.random_gaussian_m[approx_step_idx, :hidden_size].view(1, -1).detach().cpu()

                if self.args.TS_single_direction_reward and self.AIO_forward_model.last_eval_loss_val is not None:
                    prev_loss = self.AIO_forward_model.last_eval_loss_val
                    pos_rwd = (loss_pos.item() - prev_loss) / self.args.zo_eps
                    neg_rwd = (loss_neg.item() - prev_loss) / self.args.zo_eps
                    #
                    TS_model = self.AIO_forward_model.TS_model
                    #
                    TS_model.store_info(x=this_arm_context, y=pos_rwd)
                    TS_model.store_info(x=(-1*this_arm_context), y=neg_rwd)
                    #
                    # TS_model.store_info(x=(-1*this_arm_context), y=pos_rwd)
                    # TS_model.store_info(x=this_arm_context, y=neg_rwd)
                else:
                    this_reward = (loss_pos - loss_neg).item() / (2 * self.args.zo_eps)
                    #
                    TS_model = self.AIO_forward_model.TS_model
                    TS_model.store_info(x=this_arm_context, y=this_reward)

        ################## Get the forward logits after INJECTING the black-box gradient approx. ##################
        white_box_LLM.inject_data_into_proxy_weight_m(epsilon_val=epsilon_val, loss_diff=loss_diff, min_coef=self.args.hybrid_min_coef)
        ###
        white_box_LLM.set_logits_perturbation_flag(False)
        white_box_LLM.set_apply_black_box_proxy_m_flag(True)
        
        ###
        model.train()
        #
        model.zero_grad()
        self.optimizer.zero_grad()
        #
        inputs = self._prepare_inputs(inputs)

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

        # Compute the scalar logit for calculating gradients.
        with self.compute_loss_context_manager():
            loss = self.AIO_compute_loss(model, inputs)

        print("[AIO_compute_loss Loss 1]: ", loss)

        del inputs
        torch.cuda.empty_cache()

        if self.args.n_gpu > 1:
            loss = loss.mean()  # mean() to average on multi-gpu parallel training

        print("[AIO_compute_loss Loss 2]: ", loss)

        ###################################################################################
        # Perform back propagation on the loss
        if self.use_apex:
            with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                scaled_loss.backward()
        else:
            self.accelerator.backward(loss)

        ###################################################################################
        white_box_LLM.set_logits_perturbation_flag(False)
        white_box_LLM.set_apply_black_box_proxy_m_flag(False)

        return loss.detach() / self.args.gradient_accumulation_steps


    def AIO_compute_loss(self, model, inputs, return_outputs=False):
        """
        How the loss is computed by Trainer. By default, all models return the loss in the first element.

        Subclass and override for custom behavior.
        """

        # Used for computing the output for back-propagation. Note the loss here will actually be the forward logit (scalar).
        ###
        forward_model = self.AIO_forward_model
        # : the "inputs" passed into this function are not working.
        input_ids = forward_model.white_box_LLM_tokenizer(forward_model.init_token, return_tensors="pt").input_ids.cuda()
        inputs = {'input_ids': input_ids}
        outputs = model(**inputs)
        
        # Save past state if it exists
        # : this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        #
        if isinstance(outputs, dict) and "logits" not in outputs:
            raise ValueError(
                "The model did not return a loss from the inputs, only the following keys: "
                f"{','.join(outputs.keys())}. For reference, the inputs it received are {','.join(inputs.keys())}."
            )
        # We don't use .loss here since the model may return tuples instead of ModelOutput.
        logits = outputs["logits"] if isinstance(outputs, dict) else outputs[0]

        return (logits, outputs) if return_outputs else logits
    
