Enum TextGenerationOption

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

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

Variants

GPT2(GPT2Generator)

Text Generator based on GPT2 model

GPT(OpenAIGenerator)

Text Generator based on GPT model

GPTNeo(GptNeoGenerator)

Text Generator based on GPT-Neo model

GPTJ(GptJGenerator)

Text Generator based on GPT-J model

XLNet(XLNetGenerator)

Text Generator based on XLNet model

Reformer(ReformerGenerator)

Text Generator based on Reformer model

T5(T5Generator)

Text Generator based on T5 model

Implementations

impl TextGenerationOption

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

pub fn new_with_tokenizer( config: TextGenerationConfig, tokenizer: TokenizerOption, ) -> 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 get_tokenizer_mut(&mut self) -> &mut TokenizerOption

Interface method to access tokenizer

pub fn generate_indices<S>( &self, prompt_texts: Option<&[S]>, min_length: Option<i64>, max_length: Option<i64>, ) -> Result<Vec<Vec<i64>>, RustBertError>

where S: AsRef<str> + Send + Sync,

Interface method to generate() of the particular models.

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

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

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