Struct MaskedLanguageModel

pub struct MaskedLanguageModel { /* private fields */ }

MaskedLanguageModel for Masked Language (e.g. Fill Mask)

Implementations

impl MaskedLanguageModel

pub fn new( config: MaskedLanguageConfig, ) -> Result<MaskedLanguageModel, RustBertError>

Build a new MaskedLanguageModel

Arguments

• config - MaskedLanguageConfig object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)

Example

use rust_bert::pipelines::masked_language::MaskedLanguageModel;

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

pub fn new_with_tokenizer( config: MaskedLanguageConfig, tokenizer: TokenizerOption, ) -> Result<MaskedLanguageModel, RustBertError>

Build a new MaskedLanguageModel with a provided tokenizer.

Arguments

• config - MaskedLanguageConfig object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)
• tokenizer - TokenizerOption tokenizer to use for masked language modeling

Example

use rust_bert::pipelines::common::{ModelType, TokenizerOption};
use rust_bert::pipelines::masked_language::MaskedLanguageModel;
let tokenizer = TokenizerOption::from_file(
    ModelType::Bert,
    "path/to/vocab.txt",
    None,
    false,
    None,
    None,
)?;
let model = MaskedLanguageModel::new_with_tokenizer(Default::default(), tokenizer)?;

pub fn get_tokenizer(&self) -> &TokenizerOption

Get a reference to the model tokenizer.

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

Get a mutable reference to the model tokenizer.

pub fn predict<'a, S>( &self, input: S, ) -> Result<Vec<Vec<MaskedToken>>, RustBertError>

where S: AsRef<[&'a str]>,

Mask texts

Arguments

• input - &[&str] Array of texts to mask.

Returns

• Vec<String> containing masked words for input texts

Example

use rust_bert::pipelines::masked_language::MaskedLanguageModel;
//    Set-up model
let mask_language_model = MaskedLanguageModel::new(Default::default())?;

//    Define input
let input = [
    "Looks like one [MASK] is missing",
    "It was a very nice and [MASK] day",
];

//    Run model
let output = mask_language_model.predict(&input)?;
for word in output {
    println!("{:?}", word);
}