def text_generation_post_process(output_data, **kwargs):
    """Post-process data for text generation task.

    Args:
        output_data (object): Model output to be post-processed.
        **kwargs: Additional arguments.

    Returns:
        object: Post-processed data.

    """
    if isinstance(output_data, transformers.modeling_outputs.ModelOutput):
        preds = output_data.logits
    elif isinstance(output_data, dict):
        if "logits" in output_data:
            preds = output_data["logits"]
        elif "last_hidden_state" in output_data:
            preds = output_data["last_hidden_state"]
        else:
            raise ValueError("`logits` or `last_hidden_state` not found in model output.")
    else:
        preds = output_data
    return preds
Post-process data for text generation task.

    Args:
        output_data (object): Model output to be post-processed.
        **kwargs: Additional arguments.

    Returns:
        object: Post-processed data.

    
olive/data/component/post_process_data.py