def reward()

in summarize_from_feedback/reward_model.py [0:0]

• 17 lines of code
• 3 McCabe index (conditional complexity)

    def reward(self, query_tokens, response_tokens, eval_fn=None, eval_inputs=None, **kwargs):
        """
        :return: A dict with structure:
            reward: [batch, num_responses]
            eval_stats: dict of stats returned by eval_fn
        """
        last_response_indices = _response_indices(response_tokens).to(self.device)
        if self.task_hparams is not None:
            assert_eq(query_tokens.size(1), self.task_hparams.query.length)
            assert_eq(response_tokens.size(2), self.task_hparams.response.length)
        assert query_tokens.size(0) == response_tokens.size(0)
        if eval_fn is not None:
            eval_fn = _wrap_reward_model_fn(eval_fn)
            eval_inputs["last_response_index"] = last_response_indices
        result = self._eval(
            query_tokens, response_tokens, eval_fn=eval_fn, eval_inputs=eval_inputs, **kwargs
        )
        result["reward"] = gather_one(
            result["reward"]["response"][:, :, 1:], last_response_indices, dim=2
        )

        assert_shape_eq(result["reward"], (response_tokens.size(0), response_tokens.size(1)))
        return result