# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/hf.transformers.generation.ipynb.

# %% auto 0
__all__ = ['log', 'flat_pipeline', 'preprocess_generation_params', 'generate']

# %% ../../../nbs/hf.transformers.generation.ipynb 3
from copy import deepcopy

from tqdm import tqdm

from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/hf.transformers.generation.ipynb 4
def flat_pipeline(pipe):
    def func(inputs, **kwargs) -> list[str]:
        return [result[0]["generated_text"] for result in tqdm(pipe(inputs, **kwargs), total=len(inputs))]

    return func


def preprocess_generation_params(tokenizer, generation_params: dict) -> dict:
    generation_params = deepcopy(generation_params)

    terminators = generation_params.pop("terminators", [])
    eos_token_ids = {tokenizer.eos_token_id}
    for terminator in terminators:
        if isinstance(terminator, int):
            eos_token_ids.add(terminator)
        elif isinstance(terminator, str):
            eos_token_ids.add(tokenizer.convert_tokens_to_ids(terminator))
        else:
            raise ValueError(f"Invalid terminator token {terminator}.")
    generation_params["eos_token_id"] = sorted(eos_token_ids)

    generation_params["pad_token_id"] = tokenizer.eos_token_id

    return generation_params


def generate(
    pipe,
    inputs,
    **generation_kwargs,
) -> list[str]:
    # Setup generation parameters
    generation_kwargs = preprocess_generation_params(pipe.tokenizer, generation_kwargs)
    generation_kwargs["return_full_text"] = False

    # Generate text
    log.info(f"Running pipeline on {len(inputs)} samples...")
    return flat_pipeline(pipe)(inputs, **generation_kwargs)
