Struct edge_transformers::pipelines::conditional_generation::ConditionalGenerationPipeline

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

Wraps Huggingface Optimum pipeline exported to ONNX with causal-lm 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::{GraphOptimizationLevel, LoggingLevel};
use ort::environment::Environment;
use edge_transformers::{ConditionalGenerationPipeline, 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 = ConditionalGenerationPipeline::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> ConditionalGenerationPipeline<'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>

Creates new pipeline from ONNX model bytes and tokenizer configuration.

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>

Creates new pipeline from ONNX model file and tokenizer configuration.

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 from input batch text.