Struct edge_transformers::modeling::conditional_generation_with_pkvs::ConditionalGenerationModelWithPKVs

pub struct ConditionalGenerationModelWithPKVs<'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> ConditionalGenerationModelWithPKVs<'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>>, past_key_values: Option<HashMap<String, ArrayD<f32>>> ) -> Result<(Array3<f32>, HashMap<String, ArrayD<f32>>), Error>

Does inference. Returns logits and the past key values.

pub fn get_token_type_support(&self) -> bool