Project: 046_Learning_Dynamics_of_LLM_Finetuning | Function: get_batch_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:** Unified SFT/DPO analysis framework orchestrating comprehensive learning dynamics analysis for LLM finetuning, supporting both supervised fine-tuning and preference optimization modes

**KEY TERMS:** {
  "train_test": "training vs evaluation phase",
  "policy_chosen_logps": "πθt (y+u | χ+u)",
  "policy_rejected_logps": "πθt (y−u | χ−u)",
  "losses": "learning dynamics loss",
  "prob_dicts": "different response types for analysis",
  "argmax_token": "argmax predictions y*"
}

**CONTEXT CODE:**             FSDP.summon_full_params(
                self.policy, writeback=False, recurse=False
            ) if 'FSDP' in self.config.trainer else contextlib.nullcontext()
        )
        with ctx():
            policy_output = self.policy.generate(
                batch['prompt_input_ids'],
                attention_mask=batch['prompt_attention_mask'],
                max_length=self.config.max_length,
                do_sample=sample_flag,
                pad_token_id=self.tokenizer.pad_token_id
            )

        if self.config.loss.name in {'dpo', 'ipo'}:
            ctx = lambda: (FSDP.summon_full_params(self.reference_model, writeback=False, recurse=False) if 'FSDP' in self.config.trainer else contextlib.nullcontext())
            with ctx():
                reference_output = self.reference_model.generate(
                    batch['prompt_input_ids'], attention_mask=batch['prompt_attention_mask'], max_length=self.config.max_length, do_sample=True, pad_token_id=self.tokenizer.pad_token_id)
        
        policy_output = pad_to_length(
            policy_output, self.config.max_length, self.tokenizer.pad_token_id
        )
        policy_output = all_gather_if_needed(
            policy_output, self.rank, self.world_size
        )
        policy_output_decoded = self.tokenizer.batch_decode(
            policy_output, skip_special_tokens=True
        )

        if self.config.loss.name in {'dpo', 'ipo'}:
            reference_output = pad_to_length(reference_output, self.config.max_length, self.tokenizer.pad_token_id)
            reference_output = all_gather_if_needed(reference_output, self.rank, self.world_size)
            reference_output_decoded = self.tokenizer.batch_decode(reference_output, skip_special_tokens=True)
        else:
            reference_output_decoded = []

        return policy_output_decoded, reference_output_decoded

    def concatenated_forward(self, model: nn.Module, batch: Dict[str, Union[List, torch.LongTensor]]) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
        """Run the given model on the given batch of inputs, concatenating the chosen and rejected inputs together.
        
           We do this to avoid doing two forward passes, because it's faster for FSDP.
        """
        concatenated_batch = concatenated_inputs(batch)
        all_logits = model(concatenated_batch['concatenated_input_ids'], attention_mask=concatenated_batch['concatenated_attention_mask']).logits.to(torch.float32)
        all_logps, _ = _get_batch_logps(all_logits, concatenated_batch['concatenated_labels'], average_log_prob=False)
        chosen_logps = all_logps[:batch['chosen_input_ids'].shape[0]]
        rejected_logps = all_logps[batch['chosen_input_ids'].shape[0]:]
        return chosen_logps, rejected_logps

    def get_batch_metrics(
        self,
        batch: Dict[str, Union[List, torch.LongTensor]],
        loss_config: DictConfig,
        train=True,
        prob_set=None,
        force_sft=False
    ) -> Tuple[torch.FloatTensor, Dict[str, List]]:
        """Compute the SFT loss and other metrics for the given batch of inputs.
        """
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def evaluation(self, prob_set='prob_train'):
        if prob_set.lower()=='prob_train':
            data_batches = self.probtrain_batches
        elif prob_set.lower()=='prob_test':
            data_batches = self.probtest_batches
        else:
            raise ('only have prob_set naming prob_train or prob_test')

        self.policy.eval()
        all_eval_metrics = defaultdict(list)
        all_argmax_token = []
        for eval_batch in (tqdm.tqdm(data_batches, desc='Computing eval metrics') if self.rank == 0 else data_batches):
            local_eval_batch = slice_and_move_batch_for_device(eval_batch, self.rank, self.world_size, self.rank)     
            with torch.no_grad():
                # ----- detail_eval_matrics is token-wise, for k, v... then each v[i] contains logp of M tokens
                _, eval_metrics, argmax_token = self.get_batch_metrics(local_eval_batch, self.config.loss, train=False, prob_set=prob_set)                          
            all_argmax_token.append(argmax_token)  # argmax_token is [B, M], all_argmax_token is a list storing those argtokens with different length
            for k, v in eval_metrics.items():
                all_eval_metrics[k].extend(v)
                 
        # -------- Save the corresponding results
        logp_npy = np.zeros((1,1))
        mean_eval_metrics = {k: sum(v) / len(v) for k, v in all_eval_metrics.items()}
        if self.config.wandb.enabled and self.rank == 0:
            wandb.log(mean_eval_metrics, step=self.example_counter)
        if self.rank==0:
            output_dir = os.path.join(self.config.save_path, f'{prob_set}_metrics.json')
            self.save_metrics(output_dir, mean_eval_metrics)
            print(mean_eval_metrics)
            if self.config.fine_evaluation:
                # --------- Save logp; first N are for different y, last four are argmax, expect_argmax, |A_o|, |p-e| for each example
                tmp = []       
                for k, v in all_eval_metrics.items():
                    tmp.append(v)
                logp_npy = np.array(tmp)
        return logp_npy, all_argmax_token
                
    def evaluation_get_response(self, prob_set='prob_train_gen'):
        data_iterator_kwargs = dict(
                    names=self.config.datasets,
                    tokenizer=self.tokenizer,
                    max_length=self.config.max_length,
                    max_prompt_length=self.config.max_prompt_length,
                )
        probtest_iterator = get_batch_iterator(
            **data_iterator_kwargs,
            split=prob_set,
            n_examples=500,
            shuffle=False,
            batch_size=self.config.eval_batch_size,
            silent=True,
        )
        data_batches = list(probtest_iterator)
        rank0_print(f'========Loaded {len(data_batches)} prob_test batches ')

        self.policy.eval()
        all_policy_samples = []
        for eval_batch in (tqdm.tqdm(data_batches, desc='Computing eval metrics') if self.rank == 0 else data_batches):
            local_eval_batch = slice_and_move_batch_for_device(eval_batch, self.rank, self.world_size, self.rank)
            with torch.no_grad():
                policy_samples, _ = self.get_batch_samples(local_eval_batch, sample_flag=True)  # For greedy decoding, convert this to False
                for prompt, sample in zip(eval_batch['prompt'], policy_samples):
                    all_policy_samples.append({'prompt':prompt, 'response':sample})
        output_dir = os.path.join(self.config.save_path, f'{prob_set}_response.jsonl')
        with open(output_dir, 'a',newline='\n') as f:
            for i in range(len(all_policy_samples)):
                f.write(json.dumps(all_policy_samples[i]))
                f.write('\n')

    def train(self):
        """Begin either SFT or DPO training, with periodic evaluation."""

        rank0_print(f'Using {self.config.optimizer} optimizer')
        self.optimizer = getattr(torch.optim, self.config.optimizer)(
            self.policy.parameters(), lr=self.config.lr
        )
        self.scheduler = torch.optim.lr_scheduler.LambdaLR(
            self.optimizer,
            lr_lambda=lambda step: min(
                1.0, (step + 1) / (self.config.warmup_steps + 1)
            )
        )

        torch.manual_seed(self.seed)
        np.random.seed(self.seed)
        random.seed(self.seed)

        if self.config.loss.name in {'dpo', 'ipo'}:


Complete the function implementation replacing "raise NotImplementedError".