def _forward_one()

in custom/sequence_generator.py [0:0]

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

    def _forward_one(self, model, tokens, incremental_states=None, temperature=1., return_attn=False, return_logits=False, **decoder_kwargs):
        if incremental_states is not None:
            decoder_out = list(model.decoder(tokens, None, incremental_state=incremental_states, return_attn=return_attn, **decoder_kwargs))
        else:
            decoder_out = list(model.decoder(tokens, None, return_attn=return_attn, **decoder_kwargs))
        decoder_out[0] = decoder_out[0][:, -1:, :]
        if temperature != 1.:
            decoder_out[0].div_(temperature)
        attn = decoder_out[1]
        if type(attn) is dict:
            attn = attn['attn']
        if attn is not None:
            if type(attn) is dict:
                attn = attn['attn']
            attn = attn[:, :, -1, :]  # B x L x t
        if return_logits:
            logits_t = decoder_out[0][:, -1, :]
            return logits_t, attn
        log_probs = model.get_normalized_probs(decoder_out, log_probs=True)
        log_probs = log_probs[:, -1, :]
        return log_probs, attn