Function generate_and_post_process

linear_moe/generation/api.py:25–111

Run inference and post-process outputs, i.e., detokenize, move to cpu and convert to list. Args: model (torch.nn.Module): The model used for text generation. prompts (List[str], optional): A list of prompts to generate text from. tokens_to_generate (int): The ma

(model,
                              prompts=None,
                              tokens_to_generate=0,
                              return_output_log_probs=False,
                              top_k_sampling=0,
                              top_p_sampling=0.0,
                              top_p_decay=0.0,
                              top_p_bound=0.0,
                              temperature=1.0,
                              add_BOS=False,
                              use_eod_token_for_early_termination=True,
                              stop_on_double_eol=False,
                              stop_on_eol=False,
                              prevent_newline_after_colon=False,
                              random_seed=-1,
                              return_logits=False)

def generate_and_post_process(model,
                              prompts=None,
                              tokens_to_generate=0,
                              return_output_log_probs=False,
                              top_k_sampling=0,
                              top_p_sampling=0.0,
                              top_p_decay=0.0,
                              top_p_bound=0.0,
                              temperature=1.0,
                              add_BOS=False,
                              use_eod_token_for_early_termination=True,
                              stop_on_double_eol=False,
                              stop_on_eol=False,
                              prevent_newline_after_colon=False,
                              random_seed=-1,
                              return_logits=False):
    """
    Run inference and post-process outputs, i.e., detokenize,
    move to cpu and convert to list.

    Args:
        model (torch.nn.Module): The model used for text generation.
        prompts (List[str], optional): A list of prompts to generate text from.
        tokens_to_generate (int): The maximum number of tokens to generate.
        return_output_log_probs (bool): A flag indicating whether to return the output log probabilities for each generated token.
        top_k_sampling (int): The value of k for top-k sampling.
        top_p_sampling (float): The value of p for top-p sampling.
        top_p_decay (float): The amount by which to decay the value of p for each token generated.
        top_p_bound (float): The minimum value of p for top-p sampling.
        temperature (float): The temperature value to apply during sampling.
        add_BOS (bool): A flag indicating whether to add a beginning-of-sentence token to the generated output.
        use_eod_token_for_early_termination (bool): A flag indicating whether to use the end-of-document token for early termination.
        stop_on_double_eol (bool): A flag indicating whether to stop generating text when a double end-of-line token is generated.
        stop_on_eol (bool): A flag indicating whether to stop generating text when an end-of-line token is generated.
        prevent_newline_after_colon (bool): A flag indicating whether to prevent newline characters after a colon.
        random_seed (int): The random seed to use for text generation.

    Returns:
        Tuple[List[str], List[str], List[List[float]], List[int]]: A tuple containing the following elements:
            - prompts_plus_generations (List[str]): A list of prompts followed by the generated text.
            - prompts_plus_generations_segments (List[str]): A list of segments corresponding to each prompt and generated text.
            - output_log_probs (List[List[float]]): The output log probabilities for each generated token (if return_output_log_probs is True).
            - tokens (List[int]): The generated tokens.

    """

    # Main inference.
    tokens, lengths, output_log_probs, logits = generate(
        model,
        prompts=prompts,
        tokens_to_generate=tokens_to_generate,
        return_output_log_probs=return_output_log_probs,
        top_k_sampling=top_k_sampling,
        top_p_sampling=top_p_sampling,
        top_p_decay=top_p_decay,
        top_p_bound=top_p_bound,
        temperature=temperature,
        add_BOS=add_BOS,