code: def _wrap_policy_fn(fn, heads=()):
      @functools.wraps(fn)
      def wrapped(outputs_mb, inputs_mb):
          for key in heads:
              outputs_mb[key] = outputs_mb[key]["response"]
          return fn(outputs_mb, inputs_mb)
  
      return wrapped
  
file_path: summarize-from-feedback/summarize_from_feedback/policy.py
imports: import functools
  from summarize_from_feedback import tasks
  from summarize_from_feedback.query_response_model import QueryResponseModel
label_desc: Use image node_id 1 for calling a global function with example usage: _wrap_policy_fn(fn, heads) and returns: wrapped, fn