Struct edge_transformers::pipelines::conditional_generation_with_pkvs::ConditionalGenerationPipelineWithPKVs

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

Wraps Huggingface Optimum pipeline exported to ONNX with causal-lm-with-past task.

!!! Note Does not add any special tokens to the input text. If you want to add special tokens to the input text, just provide them in the prompt.

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::{ConditionalGenerationPipelineWithPKVs, TopKSampler, Device};

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

let sampler = TopKSampler::new(50, 0.9);
let pipeline = ConditionalGenerationPipelineWithPKVs::from_pretrained(
    environment.into_arc(),
    "optimum/gpt2".to_string(),
    Device::CPU,
    GraphOptimizationLevel::Level3,
).unwrap();

let input = "This is a test";

println!("{}", pipeline.generate(input, 10, &sampler).unwrap());

Implementations

impl<'a> ConditionalGenerationPipelineWithPKVs<'a>

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

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

pub fn new_from_files( environment: Arc<Environment>, model: PathBuf, tokenizer_config: PathBuf, special_tokens_map: PathBuf, device: Device, optimization_level: GraphOptimizationLevel ) -> Result<Self, Error>

pub fn generate<'sampler>( &self, prompt: &str, max_length: i32, sampler: &'sampler dyn Sampler ) -> Result<String, Error>

pub fn generate_batch<'sampler>( &self, prompt: Vec<String>, max_length: i32, sampler: &'sampler dyn Sampler ) -> Result<Vec<String>, Error>

Generates text based on formatted prompt.