Struct NERModel

pub struct NERModel { /* private fields */ }

NERModel to extract named entities

Implementations

impl NERModel

pub fn new( ner_config: TokenClassificationConfig, ) -> Result<NERModel, RustBertError>

Build a new NERModel

Arguments

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

Example

use rust_bert::pipelines::ner::NERModel;

let ner_model = NERModel::new(Default::default())?;

pub fn new_with_tokenizer( ner_config: TokenClassificationConfig, tokenizer: TokenizerOption, ) -> Result<NERModel, RustBertError>

Build a new NERModel with a provided tokenizer.

Arguments

• ner_config - NERConfig 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::ner::NERModel;
let tokenizer = TokenizerOption::from_file(
    ModelType::Bert,
    "path/to/vocab.txt",
    None,
    false,
    None,
    None,
)?;
let ner_model = NERModel::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]) -> Vec<Vec<Entity>>

where S: AsRef<str>,

Extract entities from a text

Arguments

• input - &[&str] Array of texts to extract entities from.

Returns

• Vec<Vec<Entity>> containing extracted entities

Example

let ner_model = NERModel::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);

pub fn predict_full_entities<S>(&self, input: &[S]) -> Vec<Vec<Entity>>

where S: AsRef<str>,

Extract full entities from a text performing entity chunking. Follows the algorithm for entities chunking described in Erik F. Tjong Kim Sang, Jorn Veenstra, Representing Text Chunks The proposed implementation is inspired by the Python seqeval library (shared under MIT license).

Arguments

• input - &[&str] Array of texts to extract entities from.

Returns

• Vec<Entity> containing consolidated extracted entities

Example

let ner_model = NERModel::new(Default::default())?;
let input = ["Asked John Smith about Acme Corp"];
let output = ner_model.predict_full_entities(&input);

Outputs:

Output: \

[[
    Entity {
        word: String::from("John Smith"),
        score: 0.9747,
        label: String::from("PER"),
        offset: Offset { begin: 6, end: 16 },
    },
    Entity {
        word: String::from("Acme Corp"),
        score: 0.8847,
        label: String::from("I-LOC"),
        offset: Offset { begin: 23, end: 32 },
    },
]]