Struct rust_bert::pipelines::token_classification::TokenClassificationModel

pub struct TokenClassificationModel { /* private fields */ }

TokenClassificationModel for Named Entity Recognition or Part-of-Speech tagging

Implementations

impl TokenClassificationModel

pub fn new(
    config: TokenClassificationConfig
) -> Result<TokenClassificationModel, RustBertError>

Build a new TokenClassificationModel

Arguments

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

Example

use rust_bert::pipelines::token_classification::TokenClassificationModel;

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

pub fn predict<S>(
    &self,
    input: &[S],
    consolidate_sub_tokens: bool,
    return_special: bool
) -> Vec<Vec<Token>> where
    S: AsRef<str>,

Classify tokens in a text sequence

Arguments

• input - &[&str] Array of texts to extract entities from.
• consolidate_subtokens - bool flag indicating if subtokens should be consolidated at the token level
• return_special - bool flag indicating if labels for special tokens should be returned

Returns

• Vec<Vec<Token>> containing Tokens with associated labels (for example POS tags) for each input provided

Example

let ner_model = TokenClassificationModel::new(Default::default())?;
let input = [
    "My name is Amy. I live in Paris.",
    "Paris is a city in France.",
];
let output = ner_model.predict(&input, true, true);