Struct rust_bert::pipelines::text_generation::TextGenerationModel

pub struct TextGenerationModel { /* private fields */ }

TextGenerationModel to generate texts from a prompt

Implementations

impl TextGenerationModel

pub fn new(
    generation_config: TextGenerationConfig
) -> Result<TextGenerationModel, RustBertError>

Build a new TextGenerationModel

Arguments

• generation_config - GenerateConfig object containing the resource references (model, vocabulary, configuration), generation options and device placement (CPU/GPU)
• model_type - ModelType enum variant indicating the type of model to use for generation

Example

use rust_bert::pipelines::common::ModelType;
use rust_bert::pipelines::text_generation::TextGenerationModel;

let generation_model = TextGenerationModel::new(Default::default())?;

pub fn half(&mut self)

pub fn float(&mut self)

pub fn set_device(&mut self, device: Device)

pub fn generate<'a, S>(
    &self,
    texts: &[S],
    prefix: impl Into<Option<&'a str>>
) -> Vec<String> where
    S: AsRef<str> + Sync,

Generate texts from provided prompts

Arguments

• input - &[&str] Array of texts to summarize.
• prefix - impl Into<Option<&'a str>>: Optional string to pass as a prefix for generation. Will be excluded from generated sequences.

Returns

• Vec<String> Generated texts

Example

use rust_bert::pipelines::common::ModelType;
use rust_bert::pipelines::text_generation::TextGenerationModel;

let model = TextGenerationModel::new(Default::default())?;

let input = ["The dog", "The cat was"];
let prefix = None;

let output = model.generate(&input, prefix);