Project: 007_Spread_Preference_Annotation_Direct_Preference_Jud | Function: get_batch_loss_metrics (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:** Computes batch-wise preference losses and metrics including direct preference judgment from model logits as described in SPA framework

**KEY TERMS:** {
  "policy_chosen_logps": "log πi−1(y1|x) chosen response logits",
  "policy_rejected_logps": "log πi−1(y2|x) rejected response logits",
  "reference_chosen_logps": "log πinit(y1|x) reference model logits",
  "chosen_rewards": "preference rewards for chosen responses",
  "rejected_rewards": "preference rewards for rejected responses"
}

**CONTEXT CODE:** 
        We do this to avoid doing two forward passes, because it's faster for FSDP.
        """
        concatenated_batch = self.concatenated_inputs(
            batch,
            is_encoder_decoder=self.is_encoder_decoder,
            label_pad_token_id=self.label_pad_token_id,
            padding_value=self.padding_value,
            device=self.accelerator.device,
        )
        len_chosen = batch["chosen_labels"].shape[0]

        model_kwargs = (
            {
                "labels": concatenated_batch["concatenated_labels"],
                "decoder_input_ids": concatenated_batch.pop("concatenated_decoder_input_ids", None),
            }
            if self.is_encoder_decoder
            else {}
        )
        all_logits = model(
            concatenated_batch["concatenated_input_ids"],
            attention_mask=concatenated_batch["concatenated_attention_mask"],
            **model_kwargs,
        ).logits
        
        if self.loss_type == "simpo" :
            all_logps = self.get_batch_logps(
                all_logits,
                concatenated_batch["concatenated_labels"],
                average_log_prob=True,
                is_encoder_decoder=self.is_encoder_decoder,
                label_pad_token_id=self.label_pad_token_id,
            )
        else :
            all_logps = self.get_batch_logps(
                all_logits,
                concatenated_batch["concatenated_labels"],
                average_log_prob=False,
                is_encoder_decoder=self.is_encoder_decoder,
                label_pad_token_id=self.label_pad_token_id,
            )
        chosen_logps = all_logps[:len_chosen]
        rejected_logps = all_logps[len_chosen:]

        chosen_logits = all_logits[:len_chosen]
        rejected_logits = all_logits[len_chosen:]
        
        return (chosen_logps, rejected_logps, chosen_logits, rejected_logits)

    def get_batch_loss_metrics(
        self,
        model,
        batch: Dict[str, Union[List, torch.LongTensor]],
        train_eval: Literal["train", "eval"] = "train",
    ):
        """Compute the DPO loss and other metrics for the given batch of inputs for train or test."""
        raise NotImplementedError("Function implementation has been masked for testing")



    def get_batch_loss_metrics_eval(
        self,
        model,
        batch: Dict[str, Union[List, torch.LongTensor]],
        train_eval: Literal["train", "eval"] = "train",
    ):
        """Compute individual rewards for chosen and rejected responses."""
        metrics = {}

        (
            policy_chosen_logps,
            policy_rejected_logps,
            policy_chosen_logits,
            policy_rejected_logits,
        ) = self.concatenated_forward(model, batch)

        if "reference_chosen_logps" in batch and "reference_rejected_logps" in batch:
            reference_chosen_logps = batch["reference_chosen_logps"]
            reference_rejected_logps = batch["reference_rejected_logps"]
        else:
            with torch.no_grad():
                if self.ref_model is None:
                    with self.null_ref_context():
                        (
                            reference_chosen_logps,
                            reference_rejected_logps,
                            reference_chosen_logits,
                            reference_reject_logits,
                        ) = self.concatenated_forward(self.model, batch)
                else:
                    (
                        reference_chosen_logps,
                        reference_rejected_logps,
                        reference_chosen_logits,
                        reference_reject_logits,
                    ) = self.concatenated_forward(self.ref_model, batch)

        # Calculate individual rewards
        chosen_rewards = (policy_chosen_logps - reference_chosen_logps).detach()
        rejected_rewards = (policy_rejected_logps - reference_rejected_logps).detach()

        # Save individual rewards
        output_dir = self.args.confidence_dir
        name = self.save_confidence_name
        file_path = output_dir + '/' + name

        if not os.path.exists(file_path):
            with open(file_path, 'w') as file:
                json.dump({}, file)

        if os.path.exists(file_path):
            try:
                import fcntl
                with open(file_path, 'r+', encoding='utf-8') as file:
                    fcntl.flock(file, fcntl.LOCK_EX)
                    data = json.load(file)
                    
                    for index, c_r, r_r in zip(batch["index"], 
                                            chosen_rewards.cpu().tolist(),
                                            rejected_rewards.cpu().tolist()):
                        self.save_dataset[index] = [c_r, r_r]
                        data[index] = [c_r, r_r]

                    file.seek(0)
                    json.dump(data, file, ensure_ascii=False, indent=4)
                    file.truncate()
                    fcntl.flock(file, fcntl.LOCK_UN)

            except json.JSONDecodeError:
                print("JSON 파일을 읽는 중 오류가 발생했습니다.")
            except ValueError as e:
                print(e)
        else:
            data = {}

        # Calculate metrics
        prefix = "eval_" if train_eval == "eval" else ""
        metrics[f"{prefix}rewards/chosen"] = chosen_rewards.mean().cpu()
        metrics[f"{prefix}rewards/rejected"] = rejected_rewards.mean().cpu()
        metrics[f"{prefix}logps/rejected"] = policy_rejected_logps.detach().mean().cpu()
        metrics[f"{prefix}logps/chosen"] = policy_chosen_logps.detach().mean().cpu()
        metrics[f"{prefix}logits/rejected"] = policy_rejected_logits.detach().mean().cpu()
        metrics[f"{prefix}logits/chosen"] = policy_chosen_logits.detach().mean().cpu()

        # Calculate loss (needed for trainer)
        losses = -(chosen_rewards - rejected_rewards)

        return losses.mean(), metrics



Complete the function implementation replacing "raise NotImplementedError".