Onnx transformer pipelines in Golang
TL;DR: AI use-cases such as embeddings, text generation (generative/LLMs), image classification, entity recognition, fine-tuning, and more, natively running in Go!
The goal of this library is to provide an easy, scalable, and hassle-free way to run transformer pipelines inference and training in golang applications, such as Hugging Face transformers pipelines. It is built on the following principles:
We support inference on CPU and on all accelerators supported by ONNX Runtime/OpenXLA. Note, however, that currently only CPU, TPU, and GPU inference on Nvidia GPUs via CUDA, are tested (see below).
IMPORTANT: The Go backend is designed for simpler workloads, environments that disallow cgo, and for smaller models such as all-MiniLM-L6-v2. It works best with small batches of roughly 32 inputs per call. If you have performance requirements, please move to a C backend such as XLA or ORT (detailed below).
Hugot loads and saves models in the ONNX format.
Developing and fine-tuning transformer models with the Hugging Face python library is great, but if your production stack is golang-based being able to reliably deploy and scale the resulting pytorch models can be challenging. This library aims to allow you to just lift-and-shift your python model and use the same Hugging Face pipelines you use for development for inference in a go application.
For the golang developer or ML engineer who wants to run or fine-tune transformer pipelines on their own hardware and tightly coupled with their own application, without having to deal with writing their own inference or training code.
Hugot is brought to you by the friendly folks at Knights Analytics, who use Hugot in production to automate ai-powered data curation.
Currently, we have implementations for the following transformer pipelines:
Implementations for additional pipelines will follow. We also very gladly accept PRs to expand the set of pipelines! See here for the missing pipelines that can be implemented, and the contributing section below if you want to lend a hand.
Hugot can be used both as a library and as a command-line application. See below for usage instructions.
Hugot supports pluggable backends to perform the tokenization and run the ONNX models. Currently, we support the following backends:
The native Go backend is always available, as it does not require any cgo dependencies. It is provided by the GoMLX project.
Onnx Runtime can be included at compile time via the build tag "-tags ORT". It is currently the fastest backend for inference and supports all pipelines, including generative pipelines such as text generation.
OpenXLA can be included at compile time via the build tag "-tags XLA". This is the only backend that supports TPUs. Note that it does not yet support generative pipelines or dynamic shapes.
CUDA requires a C backend, either OpenXLA or Onnx Runtime.
Once compiled, Hugot can be instantiated with your backend of choice via calling NewGoSession(), NewXLASession() or NewORTSession() respectively.
You may combine build tags "-tags XLA,ORT" or use "-tags ALL" to be able to use all available backends interchangeably.
To use Hugot as a library in your application, you can directly import it and follow the example below.
linux/amd64 you will have to download it from the ONNX Runtime releases page, see the dockerfile as an example. Hugot looks for this file at /usr/lib/libonnxruntime.so by default. A different location can be specified by passing the WithOnnxLibraryPath() option to NewORTSession(), e.g:session, err := NewORTSession(
ctx,
options.WithOnnxLibraryPath("/path/to/my/lib/directory"),
)if using XLA, the easiest way is to run "GOPROXY=direct go run github.com/gomlx/go-xla/cmd/pjrt_installer@latest -plugin=linux -version=v${GOPJRT_VERSION} -path=/usr/local/lib/go-xla", which will install the XLA backend provided by the goMLX project.
if using XLA or ORT, you will also need to use the rust-based tokenizer. The tokenizers.a file can be obtained from the releases section of this page (if you want to use alternative architecture from linux/amd64 you will have to build the tokenizers.a yourself, see here). This file should be at /usr/lib/tokenizers.a so that Hugot can load it. Alternatively, you can explicitly specify the path to the folder with the libtokenizers.a file using the CGO_LDFLAGS env variable, see the dockerfile. The tokenizer is statically linked at build time.
Alternatively, you can also use the docker image which has all the above dependencies already baked in.
The library can be used as follows:
…See also hugot_test.go for further examples for all pipelines.
Hugot uses the Onnx Runtime Generative AI backend to run generative models.
We currently support generative models only within the text generation pipeline. Please look at the ORT tests for an example of its usage.
To use the experimental Engine support for concurrent requests and inference batching, use the WithGenerativeEngine() option when creating a session.
Hugot now also supports the following accelerator backends for your inference:
Please provide feedback if encountering any issues with the accelerators above!
To use Hugot with Nvidia gpu acceleration, you need to have the following:
ctx := context.Background()
opts := []options.WithOption{
options.WithCuda(map[string]string{
"device_id": "0",
}),
}
session, err := NewORTSession(ctx, opts...)GOPROXY=direct go run github.com/gomlx/go-xla/cmd/pjrt_installer@latest -plugin=cuda13 -version=${JAX_CUDA_VERSION} -path=/usr/local/lib/go-xlactx := context.Background()
opts := []options.WithOption{
options.WithCuda(map[string]string{
"device_id": "0",
}),
}
session, err := NewXLASession(ctx, opts...)For the ONNX Runtime Cuda libraries, you can install CUDA 13.x by installing the full cuda toolkit, but that's quite a big package. In our testing on awslinux/fedora, we have been able to limit the libraries needed to run Hugot with Nvidia gpu acceleration to just these:
libcufft and libcudnn9 are lazy loaded when needed, so may be skippable depending on the models you load.
On different distros (e.g. Ubuntu), you should be able to install the equivalent packages.
Hugot now also supports the training and fine-tuning of transformer pipelines! This functionality requires that you build with XLA enabled as we use gomlx behind the scenes for training/fine-tuning: the onnx model will be loaded, converted to xla and trained using goMLX, and serialized back to onnx format.
This is currently supported only for the FeatureExtractionPipeline. This can be used to fine-tune the vector embeddings for e.g. semantic textual similarity (for applications like RAG and semantic search). In order to fine-tune the feature extraction pipeline for semantic search you will need to collect a training dataset in the following format:
{"sentence1": "The quick brown fox jumps over the lazy dog", "sentence2": "A quick brown fox jumps over a lazy dog", "score": 1}
{"sentence1": "The quick brown fox jumps over the lazy dog", "sentence2": "A quick brown cow jumps over a lazy caterpillar", "score": 0.5}See the example for a sample dataset.
The s
No open issues yet, or sync has not completed.