def _collate_fn()

in summarize_from_feedback/task_data.py [0:0]

• 13 lines of code
• 7 McCabe index (conditional complexity)

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