Project: 041_Unlocking_the_Power_of_Function_Vectors_for_Charac | Function: training_step (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Implements the core training step with function vector guided learning, combining language modeling loss with function vector consistency and KL divergence losses

**KEY TERMS:** {
  "dis_loss": "ℓFV + ℓKL",
  "pr_loss": "ℓFV consistency loss",
  "kl1, kl2": "ℓKL components",
  "loss": "ℓ total loss"
}

**CONTEXT CODE:**         kl_alpha2= 0.05,
        inter_scale: str = None,
        **kwargs
    ):  
        super().__init__(model, **kwargs)
        
        self.fv_kl = self.args.fv_kl
        self.fv_pr = self.args.fv_pr
        self.kl_alpha1 = self.args.kl_alpha1
        self.kl_alpha2 = self.args.kl_alpha2
        self.pr_alpha = self.args.pr_alpha
        
        self.use_fv = False        
        
        self.function_vector = None
        func_path = self.args.func_path
        if func_path is not None:
            if 'random' in func_path:
                function_vector = torch.load(func_path[len('random'):])
                self.function_vector = torch.sqrt(torch.var(function_vector))*torch.randn(1,self.model.config.hidden_size).to(function_vector.device)+torch.mean(function_vector)        
            else:
                function_vector = torch.load(func_path)
                self.function_vector = function_vector.reshape(1,self.model.config.hidden_size)
            self.use_fv = True
            
        
        if self.args.head_score is not None:
            indirect_effect = torch.load(head_score)
            mean_indirect_effect = - indirect_effect[:,:,:,1].mean(dim=0) # Subset to token class of interest

            h_shape = mean_indirect_effect.shape 
            topk_vals, topk_inds  = torch.topk(mean_indirect_effect.view(-1), k=10, largest=True)
            top_lh = list(zip(*np.unravel_index(topk_inds, h_shape), [round(x.item(),4) for x in topk_vals]))
            self.pos = top_lh[:10]
        else:
            if 'Llama-2-7b' in self.args.model_name_or_path:
                pos = [(14, 1, 0.0391), (11, 2, 0.0225), (9, 25, 0.02), (12, 15, 0.0196), (12, 28, 0.0191), (13, 7, 0.0171), (11, 18, 0.0152), (12, 18, 0.0113), (16, 10, 0.007), (14, 16, 0.007),]
            elif 'Llama-3.1-8B' in self.args.model_name_or_path:
                pos = [(27, 28, 0.1217), (13, 27, 0.1147), (15, 28, 0.1142), (17, 8, 0.0937), (21, 2, 0.0867), (10, 12, 0.0717), (15, 16, 0.067), (15, 2, 0.0668), (15, 1, 0.066), (31, 24, 0.0605),]
            elif 'Mistral-7B' in self.args.model_name_or_path:
                pos = [(14, 31, 0.1186), (26, 29, 0.0876), (12, 4, 0.0791), (12, 7, 0.0709), (30, 4, 0.0674), (30, 9, 0.0601), (22, 30, 0.0588), (14, 19, 0.0564), (11, 10, 0.047), (18, 1, 0.0429), ]

            self.pos = pos[:10]
        
        self.head_dim = self.model.config.hidden_size // self.model.config.num_attention_heads
        self.layer_hook_names = [f'module.base_model.model.model.layers.{layer}' for layer in range(32)]
        self.head_hook_names = [f'module.base_model.model.model.layers.{layer}.self_attn.o_proj' for layer in range(32)]

        self.cnt = 0

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

        raise NotImplementedError("Function implementation has been masked for testing")
    
    def 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.
        """
        if self.label_smoother is not None and "labels" in inputs:
            labels = inputs.pop("labels")
        else:
            labels = None
        
        outputs = model(**inputs, output_attentions=self.fv_pr)  
        
        # Save past state if it exists
        # TODO: this needs to be fixed and made cleaner later.
        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            unwrapped_model = unwrap_model(model)
            if is_peft_available() and isinstance(unwrapped_model, PeftModel):
                model_name = unwrapped_model.base_model.model._get_name()
            else:
                model_name = unwrapped_model._get_name()
            if model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
                loss = self.label_smoother(outputs, labels, shift_labels=True)
            else:
                loss = self.label_smoother(outputs, labels)
        else:
            if isinstance(outputs, dict) and "loss" 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.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
        
        if return_outputs:
            return loss, outputs
        else:
            return loss

    def KD_loss(self, new_logits, prev_logits, T):
        kd_loss = F.kl_div(F.log_softmax((prev_logits) / T, dim=-1),
                        F.softmax((new_logits) / T, dim=-1),
                        reduction='batchmean')
        
        return kd_loss
        
    def on_training_begin(self):
        return
    
    def on_training_end(self):
        return

    def _save_checkpoint(self, _, trial, metrics=None):
        """ Don't save base model, optimizer etc.
            but create checkpoint folder (needed for saving adapter) """
        checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.epoch}"

        run_dir = self._get_output_dir(trial=trial)
        output_dir = os.path.join(run_dir, checkpoint_folder)

        if metrics is not None and self.args.metric_for_best_model is not None:
            metric_to_check = self.args.metric_for_best_model
            if not metric_to_check.startswith("eval_"):
                metric_to_check = f"eval_{metric_to_check}"
            metric_value = metrics[metric_to_check]

            operator = np.greater if self.args.greater_is_better else np.less
            if (self.state.best_metric is None or self.state.best_model_checkpoint is None
                or operator(metric_value, self.state.best_metric)):
                self.state.best_metric = metric_value

                self.state.best_model_checkpoint = output_dir

        os.makedirs(output_dir, exist_ok=True)

        if self.args.should_save:
            self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)

    def _pad_tensors_to_target_len(
        self,
        src_tensor: torch.Tensor,
        tgt_tensor: torch.Tensor
    ) -> torch.Tensor:
        r"""
        Pads the tensor to the same length as the target tensor.
        """
        assert self.tokenizer.pad_token_id is not None, "Pad token is required."
        padded_tensor = self.tokenizer.pad_token_id * torch.ones_like(tgt_tensor)
        padded_tensor[:, -src_tensor.shape[-1]:] = src_tensor # adopt left-padding
        return padded_tensor.contiguous() # in contiguous memory

    def save_predictions(
        self,


Complete the function implementation replacing "raise NotImplementedError".