Struct rust_bert::pipelines::masked_language::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 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);
}