Struct edge_transformers::modeling::classification::ClassificationModel

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

Onnx inference session wrapper for the conditional generation models.

Implementations

impl<'a> ClassificationModel<'a>

pub fn new_from_memory( env: Arc<Environment>, model_bytes: &'a [u8], device: Device, optimization_level: GraphOptimizationLevel ) -> Result<Self, Error>

pub fn new_from_file<'path>( env: Arc<Environment>, model_path: PathBuf, device: Device, optimization_level: GraphOptimizationLevel ) -> Result<Self, Error>

pub fn forward( &self, input_ids: Array2<u32>, attention_mask: Array2<u32>, token_type_ids: Option<Array2<u32>> ) -> Result<ArrayD<f32>, Error>

Does inference

Returns a vector of Embedding class that contains the embedding.

pub fn get_token_type_support(&self) -> bool

pub fn get_num_labels(&self) -> usize

pub fn is_token_classification(&self) -> bool