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

Example

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

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

pub fn new_with_tokenizer( generation_config: TextGenerationConfig, tokenizer: TokenizerOption, ) -> Result<TextGenerationModel, RustBertError>

Build a new TextGenerationModel with a given tokenizer

Arguments

• generation_config - GenerateConfig object containing the resource references (model, vocabulary, configuration), generation options and device placement (CPU/GPU)
• tokenizer - TokenizerOption tokenizer to use for text generation

Example

use rust_bert::pipelines::common::{ModelType, TokenizerOption};
use rust_bert::pipelines::text_generation::TextGenerationModel;

let tokenizer = TokenizerOption::from_file(
    ModelType::GPT2,
    "path/to/vocab.json",
    Some("path/to/merges.txt"),
    false,
    None,
    None,
)?;
let generation_model = TextGenerationModel::new_with_tokenizer(Default::default(), tokenizer)?;

pub fn get_tokenizer(&self) -> &TokenizerOption

pub fn get_tokenizer_mut(&mut self) -> &mut TokenizerOption

pub fn half(&mut self) -> Result<(), RustBertError>

pub fn float(&mut self) -> Result<(), RustBertError>

pub fn set_device(&mut self, device: Device) -> Result<(), RustBertError>

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

where S: AsRef<str> + Send + 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);