def evaluate(self):
        """Samples model on `eval_prompts`, logs stats with `reward_fn` or `metric_fn` if provided"""

        stats = {}
        all_samples = []
        all_aes_score = []
        generate_time = time()
        for prompts in self.eval_dataloader:
            if hasattr(prompts, "aes_score"):
                all_aes_score.append(prompts.aes_score.cpu())

            if isinstance(prompts, torch.Tensor):
                samples = self.generate(prompts)
            else:
                samples = self.generate(**prompts)

            num_return_sequences = self.generate_kwargs["num_return_sequences"]
            samples = samples.reshape(-1, num_return_sequences, samples.shape[-1])
            samples = samples[:,0,:]

            if isinstance(samples, tuple):
                samples, *_ = samples

            pad_token = self.tokenizer.eos_token_id if self.tokenizer else 0
            all_samples.append(
                F.pad(
                    samples,
                    (0, self.max_length - samples.shape[1]),
                    value=pad_token,
                )
            )
        stats["generate_time"] = time() - generate_time

        # samples = self.accelerator.gather(torch.vstack(all_samples))
        # NOTE walkaround for parallel eval
        samples = torch.vstack(all_samples)
        plain_aes_score = torch.cat(all_aes_score) if len(all_aes_score) > 0 else None
        if self.tokenizer:
            samples = self.tokenizer.batch_decode(samples, skip_special_tokens=True)
        if isinstance(samples[0], str):
            columns_data = [samples]
        else:
            columns_data = [samples.tolist()]
        columns = ["samples"]

        if self.reward_fn:
            rewards = torch.as_tensor(self.reward_fn(samples, plain_aes_score), dtype=torch.float, device=self.accelerator.device)
            all_rewards = self.accelerator.gather(rewards)
            if self.accelerator.is_main_process:
                print("[I] all_rewards.shape=")
                print(all_rewards.shape)
                mean_reward = all_rewards.mean()
                columns.append("reward")
                columns_data.append(all_rewards)
                stats["mean_reward"] = mean_reward
                print(f"mean_reward={mean_reward}")

        if self.metric_fn:
            raise NotImplementedError
        
        if self.accelerator.is_main_process and not self.eval_only:
            rows = list(zip(*columns_data))
            stats["samples"] = wandb.Table(columns=columns, rows=rows)

            print(rows[0])

        return stats
Samples model on `eval_prompts`, logs stats with `reward_fn` or `metric_fn` if provided
promptist/trlx/trlx/model/accelerate_base_model.py