Struct edge_transformers::modeling::conditional_generation::ConditionalGenerationModel

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

Onnx inference session wrapper for the conditional generation models.

Validates inputs and outputs of the model and provides a convenient interface to the model.

Implementations

impl<'a> ConditionalGenerationModel<'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( env: Arc<Environment>, model_path: PathBuf, device: Device, optimization_level: GraphOptimizationLevel ) -> Result<Self, Error>

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

Does inference. Returns logits and the past key values.

pub fn get_token_type_support(&self) -> bool