Enum rust_bert::pipelines::text_generation::TextGenerationOption

pub enum TextGenerationOption {
    GPT2(GPT2Generator),
    GPT(OpenAIGenerator),
    GPTNeo(GptNeoGenerator),
    XLNet(XLNetGenerator),
    Reformer(ReformerGenerator),
}

Abstraction that holds one particular text generation model, for any of the supported models

Variants

GPT2(GPT2Generator)

Tuple Fields

0: GPT2Generator

Text Generator based on GPT2 model

GPT(OpenAIGenerator)

Tuple Fields

0: OpenAIGenerator

Text Generator based on GPT model

GPTNeo(GptNeoGenerator)

Tuple Fields

0: GptNeoGenerator

Text Generator based on GPT-Neo model

XLNet(XLNetGenerator)

Tuple Fields

0: XLNetGenerator

Text Generator based on XLNet model

Reformer(ReformerGenerator)

Tuple Fields

0: ReformerGenerator

Text Generator based on Reformer model

Implementations

impl TextGenerationOption

pub fn new(config: TextGenerationConfig) -> Result<Self, RustBertError>

pub fn model_type(&self) -> ModelType

Returns the ModelType for this TextGenerationOption

pub fn get_tokenizer(&self) -> &TokenizerOption

Interface method to access tokenizer

pub fn generate_indices<S>(
    &self,
    prompt_texts: Option<&[S]>,
    min_length: Option<i64>,
    max_length: Option<i64>
) -> Vec<Vec<i64>> where
    S: AsRef<str> + Sync, 

Interface method to generate() of the particular models.

pub fn half(&mut self)

pub fn float(&mut self)

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