trainer.py

from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
import numpy as np
import time
import torch
import collections
from packaging import version
from torch.distributions import Categorical
import torch.nn as nn
from loss_func.repnoise_loss import rep_noise_loss
from transformers import Trainer
from transformers import logging
# from transformers.file_utils import is_torch_tpu_available
from transformers.trainer_pt_utils import (
    get_parameter_names,
)
import torch.nn.functional as F
from transformers.utils import (
    is_sagemaker_mp_enabled
)
from utils import prune_wanda_outlier,SupervisedDataset,prune_with_FI
import copy 
from transformers.models.llama.modeling_llama import LlamaAttention,LlamaMLP
from transformers.models.opt.modeling_opt import OPTAttention
from transformers.models.mistral.modeling_mistral import MistralAttention
from transformers.models.gemma.modeling_gemma import GemmaAttention
from transformers.models.gemma2.modeling_gemma2 import Gemma2Attention
from transformers.models.qwen2.modeling_qwen2 import Qwen2Attention
from transformers.models.qwen3.modeling_qwen3 import Qwen3Attention
from transformers.models.granite.modeling_granite import GraniteAttention
# from transformers.models.falcon.modeling_falcon import FalconAttention
# from transformers.models.mistral.modeling_mistral import MistralAttentionWWWW
if version.parse(torch.__version__) >= version.parse("1.6"):
    from torch.cuda.amp import autocast
import torch.nn.functional as F

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

logger = logging.get_logger(__name__)
IGNORE_INDEX = -100

class VlguardTrainer(Trainer):
    def get_alignment_dataloader(self,alignment_dataset) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
        LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(alignment_dataset)

        dataloader_params = {
            "batch_size": 4,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(alignment_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker

        return self.accelerator.prepare(DataLoader(alignment_dataset, **dataloader_params))


    def init(self,  alignment_datast):
        self.clock = 0
        self.steps = 0
        if self.args.guide_data_num>0:
            self.alignment_dataloader = self.get_alignment_dataloader(alignment_datast)
            self.alignment_data_iter = iter(self.alignment_dataloader)

    def sample_from_alignment(self):
        # Get a  batch
        try:
            batch = next(self.alignment_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.alignment_data_iter = iter(self.alignment_dataloader)
            batch = next(self.alignment_data_iter)
        return batch

    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> torch.Tensor:
        # may change input due to mode change
        model.train()
        inputs = self._prepare_inputs(inputs)
        alignment_inputs = self.sample_from_alignment()
        alignment_inputs = self._prepare_inputs(alignment_inputs)
        def step():
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs) + self.args.lamb* self.compute_loss(model, alignment_inputs)
            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
                # print("gere2")
            return loss
        loss = step()
        return loss.detach() / self.args.gradient_accumulation_steps


class BoosterAlignmentTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )

        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm
    

    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        harmful_inputs = self._prepare_inputs(harmful_inputs)

        def step():
             # calculate the alignment grad
            # torch.cuda.empty_cache()
            with self.compute_loss_context_manager():
                loss3 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss3, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss3)

            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            model.zero_grad()


            # then backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()

            with torch.no_grad():
                grad_norm = self._grad_norm(harmful_grads)

            backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}

            # perturb the weights
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.data -= self.args.alpha * harmful_grads[name] / grad_norm

            # backward the gradient after harmful perturbation
            with self.compute_loss_context_manager():
                loss2 = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)
 
            # recover the weights
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])
            del backup 
            torch.cuda.empty_cache()
            print("harmful grad norm {}".format(grad_norm))
            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data=alignment_grads[name]  + self.args.lamb*harmful_grads[name] - self.args.lamb* param.grad.data

            del harmful_grads
            del alignment_grads
            # torch.cuda.empty_cache()
            print("adv loss: {:.4f}  align loss: {:.4f}".format(loss.item() - loss2.item(),loss3.item()))
            return loss3

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


    def init(self,  harmful_datast):
        self.clock = 0
        self.steps = 0
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_datast,batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)



class LisaTrainer(Trainer):

    def get_alignment_dataloader(self,alignment_dataset) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
        LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(alignment_dataset)

        dataloader_params = {
            "batch_size": self._train_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(alignment_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker

        return self.accelerator.prepare(DataLoader(alignment_dataset, **dataloader_params))


    def init(self,  alignment_dataset):
        if self.args.alignment_step!=0 and self.args.guide_data_num>0:
            self.status = "alignment"
        else:
            self.status = "finetune"
        self.alignment_weights ={}
        self.finetune_weights ={}
        # self.gamma ={}
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.alignment_weights[name] = param.data.detach().clone()
                self.finetune_weights[name] = param.data.detach().clone()
                # self.gamma[name]= torch.zeros_like(param)
        self.clock = 0
        self.steps = 0
        if self.args.guide_data_num>0:
            self.alignment_dataloader = self.get_alignment_dataloader(alignment_dataset)
            self.data_iter = iter(self.alignment_dataloader)

    def end_training(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                if self.status == "alignment":
                    self.alignment_weights[name] = param.data.detach().clone()
                else:
                    self.finetune_weights[name] = param.data.detach().clone()





    def switch_model(self):
        sum_drift =0
        if self.status == "alignment":
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    self.finetune_weights[name] = param.data.detach().clone()
                    sum_drift += torch.norm(self.finetune_weights[name] - self.alignment_weights[name])**2
            print("finetuning drift to consensus{}".format(sum_drift))
        else:
            for name, param in self.model.named_parameters():
                if param.requires_grad:
                    self.alignment_weights[name] = param.data.detach().clone()
                    sum_drift += torch.norm(self.finetune_weights[name] - self.alignment_weights[name])**2
            print("alignment drift to consensus{}".format(sum_drift))



    def sample_from_alignment(self):
        # Get a  batch
        try:
            batch = next(self.data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.data_iter = iter(self.alignment_dataloader)
            batch = next(self.data_iter)
        return batch


    def check_mode(self, inputs):
        if self.status == "alignment":
            if self.clock% (self.args.alignment_step )  ==  0 and self.steps!=0 and self.args.finetune_step!=0:
                self.status ="finetune"
                self.switch_model()
                # print("swith from alignment to finetune {}".format(self.steps))
                self.clock=0

            else:
                # alignment need another input
                inputs = self.sample_from_alignment()
        else:
            if  self.clock% (  self.args.finetune_step  )  ==  0 and self.steps!=0 and self.args.alignment_step!=0 and self.args.guide_data_num>0:
                self.status ="alignment"
                self.switch_model()
                 # alignment need another input

                inputs = self.sample_from_alignment()
                # print("swith from finetune to alignment {}".format(self.steps))
                self.clock=0
        return inputs


    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch
    ) -> torch.Tensor:
        # may change input due to mode change
        inputs = self.check_mode(inputs)
        model.train()

        inputs = self._prepare_inputs(inputs)


        def step():
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)
            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            if self.status =="alignment":
                # print("alignment_loss_prev: {}".format(loss.item()))
                if self.steps>0.1* len(self.get_train_dataloader()) * self.args.num_train_epochs:
                    for name, param in model.named_parameters():
                        if param.requires_grad and self.args.rho>0:
                            # loss +=torch.sum(self.gamma[name] *  param)+ self.args.rho/2* torch.norm( param- self.finetune_weights[name])**2
                            loss += self.args.rho/2* torch.norm( param- self.finetune_weights[name])**2
                # print("alignment_loss: {}".format(loss.item()))
            else:
                # print("finetune_loss_prev: {}".format(loss.item()))

                if self.steps>0.1* len(self.get_train_dataloader()) * self.args.num_train_epochs:
                    for name, param in model.named_parameters():
                        # we observe that for Gsm8k, proximal term will hurt convergence. Don't do proximal for the first few rounds.
                        if param.requires_grad and self.args.rho>0:
                            # loss += (- torch.sum(self.gamma[name] *  param )) + self.args.rho/2* torch.norm( param- self.alignment_weights[name])**2
                            loss +=  self.args.rho/2* torch.norm( param- self.alignment_weights[name])**2
                # print("finetune_loss: {}".format(loss.item()))
            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
                # print("gere2")
            return loss
        loss = step()
        self.steps+=1
        self.clock+=1
        return loss.detach() / self.args.gradient_accumulation_steps




class LDIFSTrainer(Trainer):
    
    
    def init(self, model):
        import copy
       

        # Deep copy the object
        self.alignment_model = copy.deepcopy(model)

        # Ensure all tensors are in half precision
        # self.alignment_model = self.alignment_model.half()
        # self.alignment_model.eval()
        # Verifying if the parameters are in half precision
        # for param in model.parameters():
        #     print(param.dtype)  # Should print torch.float16 for all parameters
       
        self.steps = 0
    
    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> torch.Tensor:
        # may change input due to mode change
        model.train()
        import copy
        inputs = self._prepare_inputs(inputs)
        
        def step():
            def register_activation_hook(model):
                activations = {}
                hooks = []
                i=0
                for name, param in model.named_modules():
                    if name == f'base_model.model.model.layers.{i}.mlp':
                        param.name = name
                        def _hook(module, __, val):
                            activations[module.name] = val
                            # print(val)
                        hooks += [param.register_forward_hook(_hook)]
                        i+=1
                        # print(name)
                    
                return activations, hooks 
            
            activations, hooks = register_activation_hook(model)
    
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)
            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training
            # if self.steps>=0* len(self.get_train_dataloader()) * self.args.num_train_epochs:
            # if self.steps>0.1* len(self.get_train_dataloader()) * self.args.num_train_epochs:
            
            def compare_models(model1, model2):
                for param1, param2 in zip(model1.parameters(), model2.parameters()):
                    if not torch.equal(param1, param2):
                        print("Mismatch found")
                        return False
                return True
            
            
            alignment_activations, alignment_model_hooks = register_activation_hook(self.alignment_model)            
            self.alignment_model(inputs['input_ids'], attention_mask=inputs['attention_mask'])  
            
            # if compare_models(model, self.alignment_model):
            #     print("Models are identical")
            # else:
            #     print("Models differ")
            proximal_loss=0
            for name in alignment_activations:
                # print(alignment_activations[name])
                # print(alignment_activations[name].shape)
                
                # in some layers the proximal loss will be NAN, drop those overflow loss
                proximal_loss = self.args.rho/2* torch.norm( activations [name]- alignment_activations[name])**2
                if proximal_loss<0.1:
                    # print(name)
                    # print(proximal_loss)
                    loss += proximal_loss
            # print(loss)    
            # clean up before leaving
            for hook in hooks:
                hook.remove()
            hooks = []
            activations =  {}
            
            for hook in alignment_model_hooks:
                hook.remove()
            alignment_model_hooks = []
            alignment_activations = {}
    
            
            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
                # print("gere2")
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
            return loss 
        loss = step()    
        
        self.steps+=1
        return loss.detach() / self.args.gradient_accumulation_steps
   
def get_leaf_modules_with_grad(module):
    # # print([name for name,param  in module.named_parameters()])
    # if len(list(module.children())) == 0 and any(p.requires_grad for p in module.parameters()) and "lora_B" in module._get_name():
    #     return [module]
    # else:
    #     return [submodule for child in module.children() for submodule in get_leaf_modules_with_grad(child)]
    module_list= []
    for name, module in module.named_modules():
    #     if "lora_B" in name and "v_proj" in name and len(list(module.children())) == 0:
    #         module_list+= [module]
        if isinstance(module,LlamaAttention) or isinstance(module, OPTAttention) or isinstance(module, MistralAttention) or isinstance(module, GemmaAttention) or isinstance(module, Qwen2Attention)or isinstance(module, Gemma2Attention) or isinstance(module, Qwen3Attention) or isinstance(module, GraniteAttention):
        # if isinstance(module,LlamaAttention) or isinstance(module, OPTAttention) or isinstance(module, MistralAttention):
            module_list+= [module]
    # # print(module_list)
    return module_list
            
            
class VaccineTrainer(Trainer):
    def init(self, mask_ratio):
        self.mask_ratio=mask_ratio
        self.round = 0
        # self.warm_up_round = 11999
        
    
    
    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:
        model.train()
        inputs = self._prepare_inputs(inputs)
        def step():
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)
            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
                # print("gere2")
            return loss 

        self.vaccine_state = {}
        self.vaccine_state ["hooks"] = []
        self.vaccine_state ["gradient"] = {}
        self.pre_first_step(model)
        step()
        self.after_first_step(model)
        model.zero_grad()
        self.pre_second_step(model)
        loss = step()
        self.after_second_step(model)

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

    
    
    
    @torch.no_grad()
    def pre_first_step(self, model ):
        def track_gradient_hook(module, grad_input, grad_output):
            # Store the gradients for the current layer
            self.vaccine_state["gradient"][module] = grad_output[0].detach().clone()/self.args.gradient_accumulation_steps
            # print(grad_output[0])
            
        def apply_backward_hooks_recursive(module, hook_fn, hooks):
            hook = module.register_backward_hook(hook_fn)
            hooks.append(hook)  # Append the hook to the list
            
        # Call the function with the initial empty hooks list
        leaf_modules_with_grad = get_leaf_modules_with_grad(model)
        for layer in leaf_modules_with_grad:
            self.vaccine_state["gradient"][layer] = 0
            apply_backward_hooks_recursive(layer, track_gradient_hook, self.vaccine_state["hooks"])
            
    
    
    @torch.no_grad()
    def pre_second_step(self, model):
        def purturbation_hook(module, input, output):
            # Modify the output, for example, by adding a perturbatio
            perturbation = self.vaccine_state["gradient"][module]
            # print(perturbation[0,1,:])
            # # print(output.shape)
            # print(output[0,1,:])
            output[0].data =output[0] + perturbation
            # print(output.shape)
            return output
           
        
        # Register forward hooks for adding perturbation
        def apply_purturbation_hooks_recursive(module, hook_fn, hooks):
            hook = module.register_forward_hook(hook_fn)
            hooks.append(hook)
    
        
        leaf_modules_with_grad = get_leaf_modules_with_grad(model)
        for layer in leaf_modules_with_grad:
            # print(layer._get_name())
            apply_purturbation_hooks_recursive(layer, purturbation_hook, self.vaccine_state["hooks"])
        
    @torch.no_grad()
    def after_first_step(self, model):
        for hook in self.vaccine_state["hooks"]:
            hook.remove()
        self.vaccine_state["hooks"] = []
        
        # print(self.vaccine_state["gradient"].items())
        grad_norm = self._grad_norm(self.vaccine_state["gradient"])
        # logging.info(grad_norm)
        # logging.info("norm{}".format(grad_norm))
        for module in self.vaccine_state["gradient"]:
            # grad_norm = self._grad_norm(self.vaccine_state["gradient"][module])
            grad = self.vaccine_state["gradient"][module]
            scale = self. args. rho  / (grad_norm +1e-7) 
            e_r =  (grad)* scale
            self.vaccine_state["gradient"][module] = e_r.detach().clone()
   
    @torch.no_grad()
    def after_second_step(self, model):
        # disable hook here
        # for module in self.vaccine_state["e_r"]:
        #     module.weight.data -= self.vaccine_state["e_r"][module]
        for hook in self.vaccine_state["hooks"]:
            hook.remove()
        self.vaccine_state["hooks"] = []
        # torch.nn.utils.clip_grad_norm_(model.parameters(), 10)



    @torch.no_grad()
    def _grad_norm(self,poison_grads_representation):
        norm = torch.norm(
                torch.stack([
                    #original sam 
                    ( poison_grads_representation[name] ).norm(p=2)
                    #asam 
                    # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                    for name in poison_grads_representation
                ]),
                p=2
               )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm


class RepNoiseTrainer(Trainer):
    def init(self, harmful_dataset):
        # reploss needs standard dataset, load alpaca here
        self.step_size = 0
        from transformers.trainer_utils import (seed_worker)
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator
        sampler = RandomSampler(harmful_dataset)
        dataloader_params = {
            "batch_size": self._train_batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }
        if not isinstance(harmful_dataset, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker
        self.harmful_dataloader = self.accelerator.prepare(DataLoader(harmful_dataset, **dataloader_params))

    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch
    ) -> torch.Tensor:
        model.train()
        inputs = self._prepare_inputs(inputs)
        # Get an iterator from the DataLoader
        data_iter = iter(self.harmful_dataloader)
        # Get the next batch
        harmful_inputs = next(data_iter)
        harmful_inputs = self._prepare_inputs(harmful_inputs)

        def step():
            with self.compute_loss_context_manager():
                # loss = self.compute_loss(model, inputs)
                loss = rep_noise_loss(model, harmful_inputs, inputs, beta=self.args.beta, alpha=self.args.alpha)

            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
                # print("gere2")
            return loss

        loss = step()
        self.step_size+=1

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


class FITrainer(Trainer):

    def init(self, model ):
        self.initial_weights = {}
        for name, module in model.named_modules():
            if "lora" in name  and len(list(module.children()))==0 and isinstance(module, torch.nn.Linear):
                self.initial_weights[module] = module.weight.data.detach().clone()
        self.round = 0



    def training_step(
        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> torch.Tensor:
        model.train()



        inputs = self._prepare_inputs(inputs)
        def step():
            if is_sagemaker_mp_enabled():
                loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulation_steps)
                return loss_mb.reduce_mean().detach().to(self.args.device)

            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, inputs)

            reg = 0
            for name, module in model.named_modules():
                if "lora" in name and len(list(module.children()))==0 and isinstance(module, torch.nn.Linear):
                    reg += self.args.lamb * torch.sum(self.fisher_vector[module]* torch.square(module.weight -self.initial_weights[module] ))
                    # reg += self.args.lamb * torch.sum(torch.square(module.weight -self.initial_weights[module] ))
            # print(reg)
            loss +=reg
            if self.args.n_gpu > 1:
                loss = loss.mean()  # mean() to average on multi-gpu parallel training

            if self.do_grad_scaling:
                self.scaler.scale(loss).backward()
            elif self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
            return loss

        if self.round==0:
            self. fisher_vector = {module : 0  for name, module in model.named_modules() if "lora" in name  and len(list(module.children()))==0 and isinstance(module, torch.nn.Linear)}
            eval_dataloader = self.get_eval_dataloader(self.eval_dataset)
            for stepsize, old_inputs in enumerate(eval_dataloader):
                # Update the observed num examples
                # print(inputs)
                model.zero_grad()
                old_inputs = self._prepare_inputs(old_inputs)
                with self.compute_loss_context_manager():
                    loss = self.compute_loss(model, old_inputs)
                self.accelerator.backward(loss)
                for name, module in model.named_modules():
                    if "lora" in name  and len(list(module.children()))==0 and isinstance(module, torch.nn.Linear):
                        self.fisher_vector[module] += torch.square(module.weight.grad.data.detach().clone())
                        # print(self.fisher_vector[module])
                print(loss)


        loss = step()
        # print( sum([torch.norm(self.vaccine_state ["gradient"][module]) for module in self.vaccine_state ["gradient"]  ]))
        # leaf_modules_with_grad = get_leaf_modules_with_grad(model)
        # for module in leaf_modules_with_grad:
        #     # print(module.q_proj.lora_A["default"])
        #     module.weight.grad*= (1-self.masks[index])
        #     index+=1
        self.round+=1
        return loss.detach() / self.args.gradient_accumulation_steps

class CTRAPTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=10) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            dataloader_params["worker_init_fn"] = seed_worker

        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init(self, harmful_datast, helpful_datast):
        self.clock = 0
        self.steps = 0
        if self.args.guide_data_num > 0:
            self.harmful_dataloader = self.get_harmful_dataloader(harmful_datast)
            self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0

        self.helpful_dataloader = self.get_harmful_dataloader(helpful_datast, batch_size=self._train_batch_size)
        self.helpful_data_iter = iter(self.helpful_dataloader)
        # self.random_int_tensor = torch.randint(low, high, (50000,))
        
    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    def sample_from_helpful(self):
        # Get a  batch
        try:
            batch = next(self.helpful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.helpful_data_iter = iter(self.helpful_dataloader)
            batch = next(self.helpful_data_iter)
        return batch


    def pad_and_concatenate(self, tensor1, tensor2, pad_type):
        if pad_type == "input_ids":
            pad_value = self.tokenizer.pad_token_id
        elif pad_type == "attention_mask":
            pad_value = 0
        elif pad_type == "labels":
            pad_value = IGNORE_INDEX

        b1, d1 = tensor1.size()
        b2, d2 = tensor2.size()
        max_d = max(d1, d2)

        if d1 < max_d:
            padding_a = (0, max_d - d1)
            a_padded = torch.nn.functional.pad(tensor1, padding_a, "constant", pad_value)
        else:
            a_padded = tensor1

        if d2 < max_d:
            padding_b = (0, max_d - d2)
            b_padded = torch.nn.functional.pad(tensor2, padding_b, "constant", pad_value)
        else:
            b_padded = tensor2

        result = torch.cat((a_padded, b_padded), dim=0)
        return result

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm


    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]]
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        harmful_inputs = self._prepare_inputs(harmful_inputs)

        helpful_inputs = self.sample_from_helpful()
        retain_inputs = {"input_ids":self.pad_and_concatenate(alignment_inputs["input_ids"],helpful_inputs["input_ids"],pad_type='input_ids'),
                         "attention_mask":self.pad_and_concatenate(alignment_inputs["attention_mask"],helpful_inputs["attention_mask"],pad_type='attention_mask'),
                         "labels":self.pad_and_concatenate(alignment_inputs["labels"],helpful_inputs["labels"],pad_type="labels")}
        retain_inputs = self._prepare_inputs(retain_inputs)

        error_token = self.tokenizer('an error')['input_ids'][-1]
        # 2. Generate a random token ID
        # The range is usually [0, vocab_size - 1]

        # collapse = {
        #     k: torch.where(v != IGNORE_INDEX, torch.tensor(error_token, dtype=v.dtype), v) if k == "labels" else v
        #     for k, v in retain_inputs.items()
        # }
        # import copy 
        # collapse_input = copy.deepcopy(collapse_input)
        collapse_inputs = {"input_ids":self.pad_and_concatenate(retain_inputs["input_ids"],harmful_inputs["input_ids"],pad_type='input_ids'),
                         "attention_mask":self.pad_and_concatenate(retain_inputs["attention_mask"],harmful_inputs["attention_mask"],pad_type='attention_mask'),
                         "labels":self.pad_and_concatenate(retain_inputs["labels"],harmful_inputs["labels"],pad_type="labels")}
        answer_mask = collapse_inputs["labels"] != IGNORE_INDEX
        collapse_inputs["labels"] = torch.where(answer_mask, torch.tensor(error_token, dtype=collapse_inputs["labels"].dtype), collapse_inputs["labels"])
        
      
        
        # print(collapse_inputs)
        # collapse_inputs = self._prepare_inputs(collapse)

        def step():
            # first backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)
            stored_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()

            backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # Take step with the harmful perturbation
            with torch.no_grad():
                grad_norm = self._grad_norm(stored_grads) + 1e-7
            # perturb the weights
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.data -= self.args.alpha * stored_grads[name]/grad_norm

            # backward the collapse gradient after harmful perturbation
            with self.compute_loss_context_manager():
                loss2 = self.compute_loss(model, collapse_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)
            perturb_grads = {name: param.grad.clone() for name, param in model.named_parameters() if
                             param.requires_grad}

            model.zero_grad()

            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])

            # backward the alignment gradient
            with self.compute_loss_context_manager():
                loss3 = self.compute_loss(model, retain_inputs)
            if self.use_apex:
                with amp.scale_loss(loss3, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss3)

            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data = param.grad.data + self.args.lamb * perturb_grads[name]
        
            print("Harmful loss: {:.4f}  after perturb collapse loss: {:.4f}  retain loss: {:.4f}".format(loss.item(), loss2.item(),loss3.item()))
            return loss3+loss2
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps





class AntibodyAlignmentTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=4) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )

        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm
    

    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], num_items_in_batch
    ) -> torch.Tensor:
        # may change input due to mode change
        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        harmful_inputs = self._prepare_inputs(harmful_inputs)

        def step():

            # first calculate the alignment grad
            with self.compute_loss_context_manager():
                loss4 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss4, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss4)

            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            model.zero_grad()


            # then backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                loss = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()

            # Take step with the harmful perturbation
            with torch.no_grad():
                grad_norm = self._grad_norm(harmful_grads) + 1e-7

            backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # perturb the weights
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.data -= self.args.alpha * harmful_grads[name] / grad_norm

            # backward the gradient after harmful perturbation
            with self.compute_loss_context_manager():
                loss2 = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)
            perturb_grads = {name: param.grad.clone() for name, param in model.named_parameters() if
                             param.requires_grad}

            model.zero_grad()
        
            # now caculate the harmful perturb safety alignment loss 
            with self.compute_loss_context_manager():
                loss3 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss3, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss3)

            # recover the weights
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])


    
            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data= alignment_grads[name]  + (self.args.lamb)*harmful_grads[name] -self.args.lamb* perturb_grads[name] +self.args.rho *  param.grad.data

            print("adv loss: {:.4f} perturb_safety loss {:.4f}  align loss: {:.4f}".format(loss.item() - loss2.item(),loss3.item(), loss4.item()))
            return loss3

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


    def init(self,  harmful_datast):
        self.clock = 0
        self.steps = 0
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_datast,batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)


class InterpolateTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )

        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init(self, harmful_datast, helpful_datast):
        self.clock = 0
        self.steps = 0
        if self.args.guide_data_num > 0:
            self.harmful_dataloader = self.get_harmful_dataloader(harmful_datast, batch_size=self._train_batch_size)
            self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0
        self.helpful_dataloader = self.get_harmful_dataloader(helpful_datast, batch_size=self._train_batch_size)
        self.helpful_data_iter = iter(self.helpful_dataloader)
        self.average_similarity=0
        # self.initial = {name: param.data.clone() for name, param in self.model.named_parameters() if param.requires_grad}

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    def sample_from_helpful(self):
        # Get a  batch
        try:
            batch = next(self.helpful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.helpful_data_iter = iter(self.helpful_dataloader)
            batch = next(self.helpful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm
    def _standard_cross_entropy_loss(self, model, inputs):
        # min -log p(y | x)
        outputs = model(**inputs)

        # Upcast to float if we need to compute the loss to avoid potential precision issues
        logits = outputs["logits"].float()
        labels = inputs["labels"]

        # Shift so that tokens < n predict n
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        # Flatten the tokens
        shift_logits = shift_logits.view(-1, shift_logits.shape[-1])
        shift_labels = shift_labels.view(-1)
        shift_labels = shift_labels.to(shift_logits.device)
        valid_counts = (shift_labels != -100).sum().float()

        loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
        loss = loss_fct(shift_logits, shift_labels)
        loss = loss.sum() / valid_counts
        return loss
    
    def _unlikelihood_loss(self, model, inputs):
        # min -log (1-p(y | x))
        import torch.nn.functional as F
        outputs = model(**inputs)

        # Upcast to float if we need to compute the loss to avoid potential precision issues
        logits = outputs["logits"].float()
        labels = inputs["labels"]

        # Shift so that tokens < n predict n
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        # Flatten the tokens
        shift_logits = shift_logits.view(-1, shift_logits.shape[-1])
        shift_labels = shift_labels.view(-1)
        shift_labels = shift_labels.to(shift_logits.device)

        # create a mask
        valid_mask = shift_labels != -100
        valid_logits = shift_logits[valid_mask]
        valid_labels = shift_labels[valid_mask]

        # compute the loss
        probs = F.softmax(valid_logits, dim=-1)
        gold_probs = probs[torch.arange(valid_labels.shape[0]), valid_labels]
        ## -log(1-p(y|x))
        unlikelihood_loss = -torch.log(1 - gold_probs + 1e-8)
        loss = unlikelihood_loss.mean()

        return loss
    
    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        harmful_inputs = self._prepare_inputs(harmful_inputs)
        helpful_inputs = self.sample_from_helpful()
        helpful_inputs = self._prepare_inputs(helpful_inputs)
        
        def step():
            # first backward gradient for safety dataset
            # backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                loss2 = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()
            grad_norm2 = self._grad_norm(harmful_grads) + 1e-7
            print("harmful grad norm {}".format(grad_norm2))
            # backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # with torch.no_grad():
            #     for name, param in model.named_parameters():
            #         if param.requires_grad:
            #             param.data.copy_(self.initial[name])
            # print(alignment_inputs)
            for index in range(alignment_inputs["labels"].shape[0]):
                input_vector= alignment_inputs["labels"][index]
                input_vector[input_vector == 128001] = -100
            # print(alignment_inputs)
            with self.compute_loss_context_manager():
                if self.args.beta == 0: 
                    loss1 = -self.compute_loss(model, alignment_inputs)
                    # loss1 = -1.0* self._standard_cross_entropy_loss(model, alignment_inputs)
                elif self.args.beta == 2: 
                    # loss1 = -self.compute_loss(model, alignment_inputs)
                    loss1 = self.compute_loss(model, helpful_inputs)
                    # loss1 = -1.0* self._standard_cross_entropy_loss(model, alignment_inputs)
                else:
                    loss1 = self.compute_loss(model, harmful_inputs)
                    
                if self.use_apex:
                    with amp.scale_loss(loss1, self.optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    self.accelerator.backward(loss1)
            
            self.project_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # with torch.no_grad():
            #     for name, param in model.named_parameters():
            #         if param.requires_grad:
            #             param.data.copy_(backup[name])
            with torch.no_grad():
                grad_norm = self._grad_norm(self.project_grads) 
            print("alignment grad norm {}".format(grad_norm))
            # model.zero_grad()
            # with self.compute_loss_context_manager():
            #     loss3 = self.compute_loss(model, helpful_inputs)
            #     self.accelerator.backward(loss3)
            # helpful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # grad_norm5= self._grad_norm(helpful_grads) 
            # global_dot_product = torch.tensor(0.0, device=self.accelerator.device)
            # global_align_grad_norm_sq = torch.tensor(0.0, device=self.accelerator.device)
            # projection_scalar = torch.tensor(0.0, device=self.accelerator.device)
        
            # for name in self.project_grads:
            #     # Convert the gradient to float32 for accumulatiozn to prevent
            #     # precision loss or overflow during the accumulation process.
            #     global_dot_product += torch.sum(harmful_grads[name].to(torch.float32) * self.project_grads[name].to(torch.float32))
            #     global_align_grad_norm_sq += torch.sum(self.project_grads[name].to(torch.float32) * self.project_grads[name].to(torch.float32))
            
            # if global_dot_product > 0:
            #     if global_align_grad_norm_sq > 1e-9: 
            #         projection_scalar = global_dot_product / global_align_grad_norm_sq
            # print(projection_scalar)
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data =  self.project_grads[name]
            
            # for name, param in model.named_parameters():
            #     if param.requires_grad:
            #         param.grad.data = self.args.alpha * (self.project_grads[name])/grad_norm  + (1-self.args.alpha)* harmful_grads[name]/grad_norm2
            
            del harmful_grads
            interpolate_grad = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            grad_norm4 = self._grad_norm(interpolate_grad) 
            
            with torch.no_grad():
                similarity = 0
                for name in interpolate_grad:
                    similarity += torch.sum(self.project_grads[name]* (interpolate_grad[name]) /(grad_norm*grad_norm4))
                self.average_similarity+=similarity
            
            del self.project_grads
            del interpolate_grad
            # del helpful_grads
            # del alignment_grads
            self.steps+=1
            # torch.cuda.empty_cache()
            print("alignment loss:  {:.4f} harmful loss: {:.4f} similarity {:.4f} avg_similarity {:.4f}".format(loss1.item(), loss2.item(), similarity, self.average_similarity/self.steps))
            return loss1
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps
    



class TcellTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }
        
            
        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )
        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init2(self, harmful_dataset, helpful_datast):
        self.clock = 0
        self.steps = 0
        self.harmful_dataset=harmful_dataset
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_dataset, batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0
        self.helpful_dataloader = self.get_harmful_dataloader(helpful_datast, batch_size=self._train_batch_size)
        self.helpful_data_iter = iter(self.helpful_dataloader)
     

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm

    
    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        print("alignment index {}".format(inputs["indices"]))
        print("harmful index {}".format(harmful_inputs["indices"]))
        harmful_inputs = self._prepare_inputs(harmful_inputs)
        
      
      
        def step():
            # first backward gradient for safety dataset
            with self.compute_loss_context_manager():
                loss1 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss1, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss1)
            
            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            model.zero_grad()
            
             
            with torch.no_grad():
                grad_norm = self._grad_norm(alignment_grads) 
            print("alignment grad norm {}".format(grad_norm))
            

            # then backward gradient for harmful dataset  
            with self.compute_loss_context_manager():
                # correct one
                loss2 = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            model.zero_grad()
            grad_norm2 = self._grad_norm(harmful_grads) 
            print("harmful grad norm {}".format(grad_norm2))

            
    
            backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            with torch.no_grad():
                similarity = 0
                for name in harmful_grads:
                    similarity += torch.sum(harmful_grads[name]* (alignment_grads[name]) /(grad_norm*grad_norm2))
        
        
            
            # ablation
            if self.args.ablation !="True":
                # take step with the safety gradient 
                with torch.no_grad():
                    for name, param in model.named_parameters():
                        if param.requires_grad:
                            # correct one
                            param.data -=  self.args.beta* alignment_grads[name]/grad_norm
                        
            # compute the safety gradient over the safety perturb model
            with self.compute_loss_context_manager():
                loss3= self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss3, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss3)
            perturb_safety_grads = {name: param.grad for name, param in model.named_parameters() if param.requires_grad}
            model.zero_grad()
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])
            

            # take step with the harmful gradient 
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        # correct one
                        param.data -=  self.args.beta* harmful_grads[name]/grad_norm2
            del harmful_grads
            
            # compute the harmful gradient over the harmful perturb model
            with self.compute_loss_context_manager():
                loss4 = self.compute_loss(model, alignment_inputs)
            similarity_loss = self.args.rho*torch.norm(loss3.detach()-loss4)**2
            self.accelerator.backward(similarity_loss)
                
            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])
            del backup
            # torch.cuda.empty_cache()

            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data = alignment_grads[name]+perturb_safety_grads[name] +   param.grad.data
                      
            # del alignment_grads
            # del perturb_safety_grads
            # torch.cuda.empty_cache()
            print("harmful loss: {:.4f} alignment loss {:.7f}  after safety perturb safety loss  {:.7f} after harmful perturb safety loss {:.7f}   similarity loss {:.6f}   similarity {:.4f}".format(loss2.item(),loss1.item(), loss3.item(), loss4.item(),  similarity_loss, similarity))
            return loss1
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps

class TcellTestTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }
        
            
        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )
        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init2(self, harmful_dataset):
        self.clock = 0
        self.steps = 0
        self.harmful_dataset=harmful_dataset
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_dataset, batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0
        self.similarity=0
        self.loss=0

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name].double()).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm

    
    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        # print(inputs["indices"])
        harmful_inputs = self.sample_from_harmful()
        print("alignment index {}".format(inputs["indices"]))
        print("harmful index {}".format(harmful_inputs["indices"]))
        harmful_inputs = self._prepare_inputs(harmful_inputs)
      
        def step():
            # first backward gradient for safety dataset
            with self.compute_loss_context_manager():
                loss1 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss1, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss1)
            
            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()
            
             
            with torch.no_grad():
                grad_norm = self._grad_norm(alignment_grads) 
            print("alignment grad norm {}".format(grad_norm))
                
            # backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                # correct one
                loss2 = self.compute_loss(model, harmful_inputs)
 
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            # model.zero_grad()
            grad_norm2 = self._grad_norm(harmful_grads) 
            print("harmful grad norm {}".format(grad_norm2))

            # perturb the weights (ablation here)
            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data = harmful_grads[name]/grad_norm2
            
            final_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            grad_norm3 = self._grad_norm(final_grads) 
            with torch.no_grad():
                similarity = 0
                for name in final_grads:
                    similarity += torch.sum(final_grads[name].double()* (alignment_grads[name].double()) /(grad_norm3*grad_norm))
            self.steps+=1
            self.similarity+=similarity
            self.loss += loss1.item()
            avg_similarity = self.similarity/self.steps
            avg_loss = self.loss/self.steps
            del harmful_grads
            del final_grads
            # torch.cuda.empty_cache()
            # # now backward the helful grad 
            loss6=torch.tensor(0)
            print("harmful loss: {:.4f} alignment loss {:.7f}  similarity {:.4f} avg similarity {:.4f} avg safety loss {:.7f}".format(loss2.item(),loss1.item(), similarity,avg_similarity, avg_loss))
            return loss1
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps


class gMixerTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }
        
            
        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )
        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init2(self, harmful_dataset):
        self.clock = 0
        self.steps = 0
        self.harmful_dataset=harmful_dataset
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_dataset, batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0
        self.similarity = 0
        self.loss=0
    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch

    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name].double()).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm

    
    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        # print(inputs["indices"])
        harmful_inputs = self.sample_from_harmful()
        print("alignment index {}".format(inputs["indices"]))
        print("harmful index {}".format(harmful_inputs["indices"]))
        harmful_inputs = self._prepare_inputs(harmful_inputs)
      
        def step():
            # first backward gradient for safety dataset
            with self.compute_loss_context_manager():
                loss1 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss1, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss1)
            
            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()
            
             
            with torch.no_grad():
                grad_norm = self._grad_norm(alignment_grads) 
            print("alignment grad norm {}".format(grad_norm))
                
            # backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                # correct one
                loss2 = self.compute_loss(model, harmful_inputs)
 
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            # model.zero_grad()
            grad_norm2 = self._grad_norm(harmful_grads) 
            print("harmful grad norm {}".format(grad_norm2))

            # perturb the weights (ablation here)
            
        
        
            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.grad.data = (1-self.args.alpha) * harmful_grads[name]/grad_norm2 + (self.args.alpha) * alignment_grads[name]/grad_norm
            
            final_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            grad_norm3 = self._grad_norm(final_grads) 
            with torch.no_grad():
                similarity = 0
                for name in final_grads:
                    similarity += torch.sum(final_grads[name].double()* (alignment_grads[name].double()) /(grad_norm3*grad_norm))

            self.steps+=1
            self.similarity+=similarity
            self.loss += loss1.item()
            avg_similarity = self.similarity/self.steps
            avg_loss = self.loss/self.steps
            del harmful_grads
            del final_grads
            torch.cuda.empty_cache()
            # # now backward the helful grad 
            loss6=torch.tensor(0)

            print("harmful loss: {:.4f} alignment loss {:.7f}  similarity {:.4f} avg_similarity {:.4f} avg_loss {:.4f}".format(loss2.item(),loss1.item(), similarity, avg_similarity, avg_loss))
            return loss1
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps

class TARTrainer(Trainer):

    def get_harmful_dataloader(self, harmful_datast, batch_size=5) -> DataLoader:
        """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """

        from transformers.trainer_utils import (
            seed_worker
        )
        from transformers.trainer_pt_utils import (
            LengthGroupedSampler,
        )
        from torch.utils.data import DataLoader, RandomSampler
        data_collator = self.data_collator

        sampler = RandomSampler(harmful_datast)

        dataloader_params = {
            "batch_size": batch_size,
            "collate_fn": data_collator,
            "num_workers": self.args.dataloader_num_workers,
            "pin_memory": self.args.dataloader_pin_memory,
        }

        if not isinstance(harmful_datast, torch.utils.data.IterableDataset):
            dataloader_params["sampler"] = sampler
            dataloader_params["drop_last"] = self.args.dataloader_drop_last
            from functools import partial
            dataloader_params["worker_init_fn"] = partial(
                seed_worker, num_workers=self.args.dataloader_num_workers, rank=self.args.process_index
            )

        return self.accelerator.prepare(DataLoader(harmful_datast, **dataloader_params))

    def init(self, harmful_datast):
        self.clock = 0
        self.steps = 0
        self.harmful_dataloader = self.get_harmful_dataloader(harmful_datast, batch_size=self._train_batch_size)
        self.harmful_data_iter = iter(self.harmful_dataloader)
        self.statistic = 0

    def sample_from_harmful(self):
        # Get a  batch
        try:
            batch = next(self.harmful_data_iter)
        except (StopIteration):
            # If the iterator is exhausted, create a new iterator
            self.harmful_data_iter = iter(self.harmful_dataloader)
            batch = next(self.harmful_data_iter)
        return batch


    @torch.no_grad()
    def _grad_norm(self, poison_grads_representation):
        norm = torch.norm(
            torch.stack([
                # original sam
                (poison_grads_representation[name]).norm(p=2)
                # asam
                # ((torch.abs(p) if group["adaptive"] else 1.0) * p.grad).norm(p=2).to(shared_device)
                for name in poison_grads_representation
            ]),
            p=2
        )
        # norm = ( poison_grads_representation ).norm(p=2)
        return norm


    @torch.no_grad()
    def layer_grad_norm(self, poison_grads_representation):
        norm = (poison_grads_representation).norm(p=2)
        return norm
    
    def obj_max_entropy_next_token(self, 
        model: torch.nn.Module, inputs: Dict[str, torch.Tensor]
    ):
    
        def max_entropy_loss(logits: torch.Tensor) -> torch.Tensor:
            softmax = F.softmax(logits, dim=-1)
            log_softmax = F.log_softmax(logits, dim=-1)
            entropy = torch.sum(-softmax * log_softmax, dim=-1).mean()
            return entropy.mean() * -1

        outputs = model(**inputs, output_hidden_states=False)
        logits = outputs.logits
        me_loss = max_entropy_loss(logits)
        return me_loss
        
    def training_step(
            self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]],num_items_in_batch
    ) -> torch.Tensor:

        model.train()
        alignment_inputs = self._prepare_inputs(inputs)
        harmful_inputs = self.sample_from_harmful()
        harmful_inputs = self._prepare_inputs(harmful_inputs)
      
        def step():
            # first backward gradient for safety dataset
            with self.compute_loss_context_manager():
                loss1 = self.compute_loss(model, alignment_inputs)
            if self.use_apex:
                with amp.scale_loss(loss1, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss1)
            alignment_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            model.zero_grad()
    
                
            # backward gradient for harmful dataset
            with self.compute_loss_context_manager():
                # correct one
                loss2 = self.compute_loss(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss2, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss2)

            harmful_grads = {name: param.grad.data.clone() for name, param in model.named_parameters() if
                            param.requires_grad}
            with torch.no_grad():
                grad_norm = self._grad_norm(harmful_grads) + 1e-7
            model.zero_grad()
            backup = {name: param.data.clone() for name, param in model.named_parameters() if param.requires_grad}
            # perturb the weights
            for name, param in model.named_parameters():
                if param.requires_grad:
                    param.data -= self.args.beta * harmful_grads[name]/grad_norm

            # backward the harmful gradient after harmful perturbation
            with self.compute_loss_context_manager():
                loss3 = self.obj_max_entropy_next_token(model, harmful_inputs)
            if self.use_apex:
                with amp.scale_loss(loss3, self.optimizer) as scaled_loss:
                    scaled_loss.backward()
            else:
                self.accelerator.backward(loss3)
            perturb_grads = {name: param.grad.clone() for name, param in model.named_parameters() if
                             param.requires_grad}

            with torch.no_grad():
                for name, param in model.named_parameters():
                    if param.requires_grad:
                        param.data.copy_(backup[name])

            # Finally, sum the grad
            for name, param in model.named_parameters():
                if param.requires_grad:
                    # correct one
                    param.grad.data = alignment_grads[name]  + self.args.rho * perturb_grads[name] 

            print("harmful loss: {:.4f} alignment loss {:.4f}".format(loss3.item(),loss1.item()))
            return loss1+loss3
        loss = step()
        # torch.cuda.empty_cache()
        return loss.detach() / self.args.gradient_accumulation_steps