Struct edge_transformers::pipelines::embedding::EmbeddingPipeline

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

Wraps Huggingface Optimum pipeline exported to ONNX with default task.

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

Example

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

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

let pipeline = EmbeddingPipeline::from_pretrained(
 environment.into_arc(),
"optimum/all-MiniLM-L6-v2".to_string(),
PoolingStrategy::Mean,
Device::CPU,
GraphOptimizationLevel::Level3,
).unwrap();

let input = "This is a test";
let emb1 = pipeline.embed(input).unwrap();
let input = "This is a test2";
let emb2 = pipeline.embed(input).unwrap();
println!("Similarity: {:?}", emb1.similarity(&emb2));

Implementations

impl<'a> EmbeddingPipeline<'a>

pub fn from_pretrained( env: Arc<Environment>, model_id: String, pool_strategy: PoolingStrategy, 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, pooling: PoolingStrategy, device: Device, optimization_level: GraphOptimizationLevel ) -> 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.
• 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, pooling: PoolingStrategy, device: Device, optimization_level: GraphOptimizationLevel ) -> 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 embed(&self, input: &str) -> Result<Embedding, Error>

Embeds input text.

Arguments

• input - Input text.

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

Embeds input texts.

Arguments

• inputs - Input texts.