# generate the next token
def generate_token(decoder, encoder_hidden_states, sequence):
    outputs = decoder(
        sequence,
        encoder_hidden_states
    )
    logits = outputs[0][:, -1, :]
    return torch.argmax(logits, dim=-1, keepdim=True)

# simple auto-regressive sequence generator
def image_to_text_generator(encoder, decoder, image):
    # run encoder
    encoder_hidden_states = encoder(image)[0]

    # initialize sequence
    generated_ids = torch.ones(
        (image.shape[0], 1),
        dtype=torch.long,
        device=image.device
    ) * START_ID

    for _ in range(MAX_SEQ_LEN):
        # generate next token
        next_token = generate_token(
            decoder,
            encoder_hidden_states,
            generated_ids
        )
        generated_ids = torch.cat([generated_ids, next_token], dim=-1)
        if (next_token == END_ID).all():
            break

    return generated_ids