npc-engine / edge-transformers

Rust implementation of Huggingface transformers pipelines using onnxruntime backend with bindings to C# and C.

Edge transformers is a Rust implementation of Huggingface's pipelines based on ONNX Runtime backend.

Features

• C# and C wrappers (proper C++ wrapper is planned)
• Text to output interface abstracting over tokenizers.
• Multiple ORT providers support:
  • CPU
  • CUDA (Requires building onnxruntime for with CUDA provider)
  • DirectML
  • More planned

Tasks implemented

Model export feature/task	Class name
causal-lm	ConditionalGenerationPipeline
causal-lm-with-past	ConditionalGenerationPipelineWithPKVs
default	EmbeddingPipeline
seq2seq-lm	Seq2SeqGenerationPipeline or OptimumSeq2SeqPipeline
seq2seq-lm-with-past	OptimumSeq2SeqPipelineWithPKVs
sequence-classification	SequenceClassificationPipeline
token-classification	TokenClassificationPipeline

Usage

Your linker must be able to find onnxruntime.dll and edge-transformers.dll (or *.so on Linux). You can find C and C# wrappers in c and csharp folders respectively.

C#

Documentation is WIP, refer to Rust documentation for now.

using EdgeTransformers;

...
        var env = EnvContainer.New();

        var conditionalGen = ConditionalGenerationPipelineFFI.FromPretrained(
            env.Context, "optimum/gpt2", DeviceFFI.CPU, GraphOptimizationLevelFFI.Level3
        );
        var outp = conditionalGen.GenerateTopkSampling("Hello", 2, 50, 0.9f);
        Assert.IsNotNull(outp);
...

Batch processing is supported, but is a bit unintuitive and requires StringBatch class.

using EdgeTransformers;
  ...
        var env = EnvContainer.New();
        var condPipelinePkv = ConditionalGenerationPipelineFFI.FromPretrained(
            env.Context, "optimum/gpt2", DeviceFFI.DML, GraphOptimizationLevelFFI.All);
        var string_batch = StringBatch.New();
        string_batch.Add("Hello world");
        string_batch.Add("Hello world");

        var o_batch_pkv = condPipelinePkv.GenerateRandomSamplingBatch(string_batch.Context, 10, 0.5f);

        Debug.LogFormat("Cond generation output 0: {0} 1: {1}", o_batch_pkv[0].ascii_string, o_batch_pkv[1].ascii_string);
  ...

C

TODO

Rust

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());

Roadmap

• C# wrapper
• C wrapper
• Migrate to ort crate
• Proper CI/CD to test and build for more execution providers
• C++ wrapper
• More pipelines (e.g. extractive QA, ASR, etc.)
• Better huggingface config.json parsing

Building

Please refer to ONNX Runtime bindings docs on detailed how to.

Testing

Tests require to be running on a single thread at least for the first time. The reason is that they use *::from_pretrained function that downloads data from Huggingface Hub and some tests rely on the same files being downloaded. Second time they can run in parallel because they use cached files.

e.g. First time command:

cargo test -- --test-threads=1

e.g. Second time:

cargo test