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::pipelines::generation_utils::{
    GPT2Generator, GenerateConfig, LanguageGenerator,
};

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

let min_length = Some(32);
let max_length = Some(128);
let decoder_start_id = None;

let input_context = "The dog";
let second_input_context = "The cat was";
let output = gpt2_generator.generate(
    Some(vec![input_context, second_input_context]),
    None,
    min_length,
    max_length,
    decoder_start_id,
);

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

BartGenerator	Language generation model based on the Bart architecture
GPT2Generator	Language generation model based on the GPT2 architecture
GenerateConfig	Configuration for text generation
LMModelOutput	Container holding a language model output for generation tasks
MarianGenerator	Language generation model based on the Marian architecture for machine translation
OpenAIGenerator	Language generation model based on the GPT architecture
T5Generator
XLNetGenerator	Language generation model based on the XLNet architecture

Enums

Cache

Traits

LMHeadModel	Language Model trait
LanguageGenerator	Common trait for text generation models.