Struct edge_transformers::pipelines::seq2seq_generation::Seq2SeqGenerationPipeline

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

Wraps Huggingface AutoModelForCausalLM exported to ONNX with seq2seq-lm feature and pretrained tokenizer into simple text to text generative interface.

!!! Note Uses bos_token if it exists to start decoder generation and eos_token to stop generation. If bos_token is not in special_tokens_map, then eos_token is used. If you want to use a different first decoder token, provide decoder_prompt parameter. No additional special tokens are added to decoder_prompt.

Export docs https://huggingface.co/docs/transformers/serialization#export-to-onnx

Example

 use std::fs;
 use ort::{GraphOptimizationLevel, LoggingLevel};
 use ort::environment::Environment;
 use edge_transformers::{Seq2SeqGenerationPipeline, 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 = Seq2SeqGenerationPipeline::from_pretrained(
     environment.into_arc(),
    "optimum/t5-small".to_string(),
     Device::CPU,
     GraphOptimizationLevel::Level3,
 ).unwrap();

 let input = "This is a test";

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

Implementations

impl<'a> Seq2SeqGenerationPipeline<'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, decoder_prompt: Option<&str>, max_length: i32, sampler: &'sampler dyn Sampler ) -> Result<String, Error>

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

Generates text from input batch text.