Function decoder_inference

demo/HuggingFace/BART/measurements.py:50–68

(
    BART_decoder, input_ids, encoder_last_hidden_state, timing_profile, use_cuda=True, use_cache=False, past_key_values=None
)

@use_cuda
def decoder_inference(
    BART_decoder, input_ids, encoder_last_hidden_state, timing_profile, use_cuda=True, use_cache=False, past_key_values=None
):
    # This implementation is a bit ugly. Moving implementation of the model to check HFRunner would be cleaner.
    if isinstance(BART_decoder, TRTNativeRunner):
        # Function is technically in BARTTRTDecoder however due to circular import, TRTNativeRunner in this module scope
        # implies the existence of this function.
        BART_decoder.set_encoder_hidden_states_for_inference_cycle(encoder_last_hidden_state)
        BART_decoder.set_return_device("cuda" if use_cuda else "cpu")

    def decoder_stmt():
        BART_decoder(
            input_ids=input_ids, encoder_hidden_states=encoder_last_hidden_state, use_cache=use_cache,
            past_key_values=past_key_values
        )

    decoder_e2e_time = measure_python_inference_code(decoder_stmt, timing_profile)

    return (decoder_stmt(), decoder_e2e_time)


@use_cuda