Module rust_bert::pipelines::generation_utils

Natural Language Generation utilities

Set of text generation utilities, serving as a basis for TextGenerationModel, SummarizationModels and TranslationModels. Include techniques such as beam search, top-k and nucleus sampling, temperature setting and repetition penalty. Supports batch generation of sentences from several prompts. Sequences will be left-padded with the model’s padding token if present, the unknown token otherwise. This may impact the results and it is recommended to submit prompts of similar length for best results.

use rust_bert::gpt2::GPT2Generator;
use rust_bert::pipelines::generation_utils::{
    GenerateConfig, GenerateOptions, LanguageGenerator,
};

let generate_config = GenerateConfig {
    do_sample: true,
    num_beams: 5,
    temperature: 1.1,
    num_return_sequences: 3,
    ..Default::default()
};
let mut gpt2_generator = GPT2Generator::new(generate_config)?;

let input_context = "The dog";
let second_input_context = "The cat was";

let generate_options = GenerateOptions {
    min_length: Some(32),
    max_length: Some(128),
    output_scores: true,
    ..Default::default()
};

let output = gpt2_generator.generate(
    Some(&[input_context, second_input_context]),
    Some(generate_options),
);

Example output:

[
    "The dog's owners, however, did not want to be named. According to the lawsuit, the animal's owner, a 29-year",
    "The dog has always been part of the family. \"He was always going to be my dog and he was always looking out for me",
    "The dog has been able to stay in the home for more than three months now. \"It's a very good dog. She's",
    "The cat was discovered earlier this month in the home of a relative of the deceased. The cat\'s owner, who wished to remain anonymous,",
    "The cat was pulled from the street by two-year-old Jazmine.\"I didn't know what to do,\" she said",
    "The cat was attacked by two stray dogs and was taken to a hospital. Two other cats were also injured in the attack and are being treated."
]

Structs

GenerateConfig

Configuration for text generation

GenerateOptions

Generation options for text generation.

GeneratedIndicesOutput

Generated indices output

GeneratedTextOutput

Generated text output

LMModelOutput

Container holding a language model output for generation tasks

Enums

Cache

Traits

LMHeadModel

Language Model trait

LanguageGenerator

Common trait for text generation models.