code: def _collate_fn(raw_data, all_fields=False, device="cpu"):
      context_input = torch.as_tensor(
          [x["context"]["tokens"] for x in raw_data],
          dtype=torch.long,
          device=device,
      )
      reference_input = torch.as_tensor(
          [x["reference"]["tokens"] for x in raw_data],
          dtype=torch.long,
          device=device,
      )
      input_dict = dict(
          context=dict(tokens=context_input),
          reference=dict(tokens=reference_input),
      )
      if "text" in raw_data[0]["reference"]:
          input_dict["reference"]["text"] = [
              x["reference"]["text"] for x in raw_data
          ]
      if all_fields:
          input_dict["extra_fields"] = [
              x["extra_fields"] for x in raw_data
          ]
      return input_dict
  
file_path: summarize-from-feedback/summarize_from_feedback/task_data.py
imports: import json
  import os
  from functools import partial
  from glob import glob
  from typing import Optional
  import blobfile
  import torch
  import summarize_from_feedback
  from summarize_from_feedback import tasks
  from summarize_from_feedback.datasets import jsonl_encoding, get_dataset
  from summarize_from_feedback.model_layout import ModelLayout
  from summarize_from_feedback.utils import even_more_itertools, blobs
label_desc: Use image node_id 1 for calling a global function with example usage: _collate_fn(raw_data, all_fields, device) and returns: input_dict