def process_responses()

in summarize_from_feedback/tasks.py [0:0]

• 12 lines of code
• 2 McCabe index (conditional complexity)

    def process_responses(self, unprocessed_tokens: SampledTokens) -> ProcessedTokens:
        assert unprocessed_tokens.size(-1) == self.H.length
        if self.H.truncate_token is not None:
            assert self.padding_token is not None
            trunc_idxs = first_true_indices(unprocessed_tokens == self.H.truncate_token).unsqueeze(
                -1
            )
            new_size = [1] * (len(unprocessed_tokens.size()) - 1) + [self.H.length]
            idxs = torch.arange(self.H.length, device=unprocessed_tokens.device).view(*new_size)
            return torch.masked_fill(unprocessed_tokens, idxs > trunc_idxs, self.padding_token)
        else:
            return unprocessed_tokens