def METHOD_NAME(model, dataset, vocab_path):
    _, tokens_to_ids, _ = data.vocab.create_lookup_tables(vocab_path)
    elem = next(iter(dataset))
    elem_ids = tf.cast(tf.squeeze(elem["ids"]), dtype=tf.dtypes.int32)
    _, pred = model(elem)
    pred_ids = tf.squeeze(tokens_to_ids.lookup(pred["tokens"]))
    tflite_concrete_fn = model.tflite_function().get_concrete_function()
    tflite_pred_ids = tflite_concrete_fn(elem_ids)
    # Model can sometimes output a tensor of shape (), if it does, converts to shape (1,)
    if tf.size(tf.shape(pred_ids)) == 0:
        pred_ids = tf.convert_to_tensor([pred_ids])
    # Modify tflite ids tensor to be the same size as normal model output
    if tf.size(pred_ids) < tf.size(tflite_pred_ids):
        tflite_pred_ids = tf.slice(tflite_pred_ids, [0], [tf.size(pred_ids)])
    elif tf.size(pred_ids) > tf.size(tflite_pred_ids):
        tflite_pred_ids = tf.pad(
            tflite_pred_ids, [[0, tf.size(pred_ids) - tf.size(tflite_pred_ids)]]
        )
    return pred_ids, tflite_pred_ids