Struct edge_transformers::pipelines::token_classification::TokenClassificationPipeline

pub struct TokenClassificationPipeline<'a> { /* private fields */ }

Wraps Huggingface Optimum pipeline exported to ONNX with token-classification task.

Export docs https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model

Example

use std::fs;
use edge_transformers::Device;
use ort::{GraphOptimizationLevel, LoggingLevel};
use ort::environment::Environment;
use edge_transformers::TokenClassificationPipeline;

let environment = Environment::builder()
  .with_name("test")
 .with_log_level(LoggingLevel::Verbose)
.build()
.unwrap();

let pipeline = TokenClassificationPipeline::from_pretrained(
    environment.into_arc(),
    "optimum/bert-base-NER".to_string(),
    Device::CPU,
    GraphOptimizationLevel::Level3
).unwrap();

let input = "This is a test";

println!("{:?}", pipeline.tag(input).unwrap());

Implementations

impl<'a> TokenClassificationPipeline<'a>

pub fn from_pretrained( env: Arc<Environment>, model_id: String, device: Device, optimization_level: GraphOptimizationLevel ) -> Result<Self, Error>

pub fn new_from_files( environment: Arc<Environment>, model_path: PathBuf, tokenizer_config: PathBuf, special_tokens_map: PathBuf, device: Device, optimization_level: GraphOptimizationLevel, labels: Option<Vec<String>> ) -> Result<Self, Error>

Creates new pipeline from model and tokenizer configuration files.

Arguments

• environment - ONNX Runtime environment.
• model_path - Path to ONNX model file.
• tokenizer_config - Path to tokenizer configuration file.
• special_tokens_map - Path to special tokens map file. Maps token names to their string values.
• device - Device to run the model on.
• optimization_level - ONNX Runtime graph optimization level.

pub fn new_from_memory( environment: Arc<Environment>, model: &'a [u8], tokenizer_config: String, special_tokens_map: String, device: Device, optimization_level: GraphOptimizationLevel, labels: Option<Vec<String>> ) -> Result<Self, Error>

Creates new pipeline from model and tokenizer configuration files.

Arguments

• environment - ONNX Runtime environment.
• model - ONNX model file content.
• tokenizer_config - Path to tokenizer configuration file.
• special_tokens_map - Path to special tokens map file.
• device - Device to run the model on.
• optimization_level - ONNX Runtime graph optimization level.

pub fn tag(&self, input: &str) -> Result<TaggedString, Error>

Embeds input text.

Arguments

• input - Input text.

pub fn tag_batch(&self, inputs: Vec<String>) -> Result<Vec<TaggedString>, Error>

Embeds input texts.

Arguments

• inputs - Input texts.