def _split_query_response_output_parts()

in summarize_from_feedback/query_response_model.py [0:0]

• 31 lines of code
• 5 McCabe index (conditional complexity)

def _split_query_response_output_parts(x, query_length, response_padding_mask):
    """
    Given an output x with shape [batch, num_responses, query_length + response_length, *rest],
    returns a dictionary with it split into query/response parts with shapes
    [batch, query_length + 1, *rest] and [batch, num_responses, response_length + 1, *rest]
    """
    assert x.ndim >= 3
    rest_shape = x.size()[3:]
    d = dict()
    # Add this back if it's ever actually useful
    # d["query"] = torch.cat(
    #     [nans([x.size(0), 1, *rest_shape], dtype=x.dtype, device=x.device), x[:, 0, :query_length]],
    #     dim=1,
    # )
    if query_length > 0:
        d["response"] = x[:, :, query_length - 1 :]
    else:
        d["response"] = torch.cat(
            [
                nans([x.size(0), x.size(1), 1, *rest_shape], dtype=x.dtype, device=x.device),
                x[:, :, :query_length],
            ],
            dim=2,
        )
    for _ in range(len(rest_shape)):
        response_padding_mask = response_padding_mask.unsqueeze(-1)
    # fill with NaNs in places where response had padding
    d["response"].masked_fill_(
        torch.cat(
            [
                torch.zeros(
                    [x.size(0), x.size(1), 1] + [1 for _ in range(len(rest_shape))],
                    dtype=torch.bool,
                    device=x.device,
                ),
                response_padding_mask,
            ],
            dim=2,
        ),
        np.nan,
    )
    return d