Struct rust_bert::pipelines::token_classification::TokenClassificationModel

pub struct TokenClassificationModel { /* fields omitted */ }

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

Implementations

impl TokenClassificationModel

pub fn new(
    config: TokenClassificationConfig
) -> Fallible<TokenClassificationModel>

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(
    &self,
    input: &[&str],
    consolidate_sub_tokens: bool,
    return_special: bool
) -> Vec<Token>

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<Token> containing Tokens with associated labels (for example POS tags)

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);