Struct 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 new_with_tokenizer( config: TokenClassificationConfig, tokenizer: TokenizerOption, ) -> Result<TokenClassificationModel, RustBertError>

Build a new TokenClassificationModel with a provided tokenizer.

Arguments

• config - TokenClassificationConfig object containing the resource references (model, vocabulary, configuration) and device placement (CPU/GPU)
• tokenizer - TokenizerOption tokenizer to use for token classification

Example

use rust_bert::pipelines::common::{ModelType, TokenizerOption};
use rust_bert::pipelines::token_classification::TokenClassificationModel;
let tokenizer = TokenizerOption::from_file(
    ModelType::Bert,
    "path/to/vocab.txt",
    None,
    false,
    None,
    None,
)?;
let model = TokenClassificationModel::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<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);