Function generate

linear_moe/generation/api.py:114–211

Given prompts and input parameters, run inference and return the generated tokens, lengths, and output log probabilities. Args: model (torch.nn.Module): The model used for text generation. prompts (List[str], optional): A list of prompts to generate text from. t

(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)

def generate(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):
    """
    Given prompts and input parameters, run inference and return the generated tokens,
    lengths, and output log probabilities.

    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[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing the following elements:
            - tokens (torch.Tensor): The prompt plus generated tokens.
            - lengths (torch.Tensor): The lengths of the prompt plus the generated tokens.
            - output_log_probs (torch.Tensor): The output log probabilities for each generated token.
    """

    # Make sure input params are avaialble to all ranks.
    values = [
        tokens_to_generate, return_output_log_probs, top_k_sampling,
        top_p_sampling, top_p_decay, top_p_bound, temperature, add_BOS,
        use_eod_token_for_early_termination, stop_on_double_eol, stop_on_eol,
        prevent_newline_after_colon, random_seed
    ]
    values_float_tensor = broadcast_float_list(len(values), float_list=values)
    tokens_to_generate = int(values_float_tensor[0].item())
    return_output_log_probs = bool(values_float_tensor[1].item())
    top_k_sampling = int(values_float_tensor[2].item())
    top_p_sampling = values_float_tensor[3].item()
    top_p_decay = values_float_tensor[4].item()
    top_p_bound = values_float_tensor[5].item()
    temperature = values_float_tensor[6].item()