Struct rust_bert::pipelines::ner::NERModel

pub struct NERModel { /* fields omitted */ }

NERModel to extract named entities

Methods

impl NERModel

pub fn new(
    vocab_path: &Path,
    config_path: &Path,
    weights_path: &Path,
    device: Device
) -> Fallible<NERModel>

Build a new NERModel

Arguments

• vocab_path - Path to the model vocabulary, expected to have a structure following the Transformers library convention
• config_path - Path to the model configuration, expected to have a structure following the Transformers library convention
• weights_path - Path to the model weight files. These need to be converted form the .bin to .ot format using the utility script provided.
• device - Device to run the model on, e.g. Device::Cpu or Device::Cuda(0)

Example

use tch::Device;
use std::path::{Path, PathBuf};
use rust_bert::pipelines::ner::NERModel;

let mut home: PathBuf = dirs::home_dir().unwrap();
let config_path = &home.as_path().join("config.json");
let vocab_path = &home.as_path().join("vocab.txt");
let weights_path = &home.as_path().join("model.ot");
let device = Device::Cpu;
let ner_model =  NERModel::new(vocab_path,
                               config_path,
                               weights_path,
                               device)?;

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

Extract entities from a text

Arguments

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

Returns

• Vec<Entity> containing extracted entities

Example

let ner_model =  NERModel::new(vocab_path,
                               config_path,
                               weights_path,
                               device)?;
let input = [
    "My name is Amy. I live in Paris.",
    "Paris is a city in France."
];
let output = ner_model.predict(&input);