Module rust_bert::pipelines::generation

Natural Language Generation pipeline

Generate language based on a prompt. GPT2 and GPT available as base models. 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. All resources for this model can be downloaded using the Python utility script included in this repository.

1. Set-up a Python virtual environment and install dependencies (in ./requirements.txt)
2. Run the conversion script python /utils/download-dependencies_gpt2.py (or /utils/download-dependencies_openaigpt.py) The dependencies will be downloaded to the user's home directory, under ~/rustbert/gpt2 (~/rustbert/openai-gpt respectively)

use rust_bert::pipelines::generation::{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 input_context = "The dog";
let second_input_context = "The cat was";
let output = gpt2_generator.generate(Some(vec![input_context, second_input_context]), None);

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
MarianGenerator	Language generation model based on the Marian architecture for machine translation
OpenAIGenerator	Language generation model based on the GPT architecture
T5Generator

Enums

Cache

Traits

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