Struct rust_bert::pipelines::ner::NERModel

pub struct NERModel { /* fields omitted */ }

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 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(&self, input: &[&str]) -> Vec<Vec<Entity>>

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"),
    },
    Entity {
        word: String::from("Acme Corp"),
        score: 0.8847,
        label: String::from("I-LOC"),
    },
]]