hugot

README

Hugot: Huggingface 🤗 pipelines for golang

What

The goal of this library is to provide an easy, scalable, and hassle-free way to run huggingface transformer pipelines in golang applications. It is built on the following principles:

1. Fidelity to the original Huggingface python implementations: the aim is to accurately replicate huggingface inference implementations for the implemented pipelines, so that models trained and tested in python can be seamlessly deployed in a golang application
2. Hassle-free and performant production use: we exclusively support onnx exports of huggingface models. Pytorch transformer models that don't have an onnx version can be easily exported to onnx via huggingface optimum, and used with the library
3. Run on your hardware: this library is for those who want to run transformer models tightly coupled with their go applications, without the performance drawbacks of having to hit a rest API, or the hassle of setting up and maintaining e.g. a python RPC service that talks to go.

We support inference on CPU and on all accelerators supported by ONNXRuntime. Note, however, that currently only CPU and GPU inference on nvidia GPU (with cuda) are tested (see below).

Why

Developing and fine-tuning transformer models with the huggingface python library is a great experience, but if your production stack is golang-based being able to reliably deploy and scale the resulting pytorch models can be challenging and require quite some setup. This library aims to allow you to just lift-and-shift your python model and use the same huggingface pipelines you use for development for inference in a go application.

For whom

For the golang developer or ML engineer who wants to run transformer piplines on their own hardware, tightly coupled with their own application.

Implemented pipelines

Currently, we have implementations for the following transfomer pipelines:

• featureExtraction
• textClassification
• tokenClassification

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.

Hardware acceleration 🚀

Hugot now also supports the following accelerator backends for your inference:

• CUDA (tested). See below for setup instructions.
• TensorRT (untested)
• DirectML (untested)
• CoreML (untested)
• OpenVINO (untested)

Please help us out by testing the untested options above and providing feedback, good or bad!

To use Hugot with nvidia gpu acceleration, you need to have the following:

• The cuda gpu version of onnxruntime on the machine/docker container. You can see how we get that by looking at the Dockerfile. You can also get the onnxruntime libraries that we use for testing from the release. Just download the gpu .so libraries and put them in /usr/lib64.
• the nvidia driver for your graphics card
• the required cuda libraries installed on your system that are compatible with the onnxruntime gpu version you use. See here. For instance, for onnxruntime-gpu 17.3, we need CUDA 12.x (any minor version should be compatible) and cuDNN 8.9.2.26.

On the last point above, you can install CUDA 12.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:

• cuda-cudart-12-4 libcublas-12-4 libcurand-12-4 libcufft-12-4 (from fedora repo)
• libcudnn8 (from RHEL repo, for cuDNN)

On different distros (e.g. Ubuntu), you should be able to install the equivalent packages and gpu inference should work.

Limitations

Apart from the fact that only the aforementioned pipelines are currently implemented, the current limitations are: - the library and cli are only built/tested on amd64-linux

Pipelines are also tested on specifically NLP use cases. In particular, we use the following models for testing:

• feature extraction: all-MiniLM-L6-v2
• text classification: distilbert-base-uncased-finetuned-sst-2-english
• token classification: distilbert-NER and Roberta-base-go_emotions

If you encounter any further issues or want further features, please open an issue.

Installation and usage

Hugot can be used in two ways: as a library in your go application, or as a command-line binary.

Use it as a library

To use Hugot as a library in your application, you will need the following dependencies on your system:

• the tokenizers.a file obtained from building the tokenizer go library (which is itself a fork of https://github.com/daulet/tokenizers). This file should be at /usr/lib/tokenizers.a so that hugot can load it.
• the onnxruntime.go file obtained from the onnxruntime project. This is dynamically linked by hugot and used by the onnxruntime inference libraryonnxruntime_go. This file should be at /usr/lib/onnxruntime.so or /usr/lib64/onnxruntime.so

You can get the libtokenizers.a in two ways. Assuming you have rust installed, you can compile the tokenizers library and get the required libtokenizers.a:

git clone https://github.com/Knights-Analytics/tokenizers -b main && \
    cd tokenizers && \
    cargo build --release
mv target/release/libtokenizers.a /usr/lib/libtokenizers.a

Alternatively, you can just download libtokenizers.a from the release section of the repo.

For onnxruntime, it suffices to download it, untar it, and place it in the right location:

curl -LO https://github.com/microsoft/onnxruntime/releases/download/v${ONNXRUNTIME_VERSION}/onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz && \
   tar -xzf onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}.tgz && \
   mv ./onnxruntime-linux-x64-${ONNXRUNTIME_VERSION}/lib/libonnxruntime.so.${ONNXRUNTIME_VERSION} /usr/lib/onnxruntime.so

See also the dockerfile used for building & testing.

Once these pieces are in place, the library can be used as follows:

import (
	"github.com/knights-analytics/hugot"
	"github.com/knights-analytics/hugot/pipelines"
)

func check(err error) {
    if err != nil {
        panic(err.Error())
    }
}
// start a new session. This looks for the onnxruntime.so library in its default path, e.g. /usr/lib/onnxruntime.so
session, err := hugot.NewSession()
// if your onnxruntime.so is somewhere else, you can explicitly set it by using WithOnnxLibraryPath
// session, err := hugot.NewSession(WithOnnxLibraryPath("/path/to/onnxruntime.so"))
check(err)
// A successfully created hugot session needs to be destroyed when you're done
defer func(session *hugot.Session) {
    err := session.Destroy()
    check(err)
}(session)

// Let's download an onnx sentiment test classification model in the current directory
// note: if you compile your library with build flag NODOWNLOAD, this will exclude the downloader.
// Useful in case you just want the core engine (because you already have the models) and want to
// drop the dependency on huggingfaceModelDownloader.
modelPath, err := session.DownloadModel("KnightsAnalytics/distilbert-base-uncased-finetuned-sst-2-english", "./", hugot.NewDownloadOptions())
check(err)

// we now create the configuration for the text classification pipeline we want to create.
// Options to the pipeline can be set here using the Options field
config := TextClassificationConfig{
    ModelPath: modelPath,
    Name:      "testPipeline",
}
// then we create out pipeline.
// Note: the pipeline will also be added to the session object so all pipelines can be destroyed at once
sentimentPipeline, err := NewPipeline(session, config)
check(err)

// we can now use the pipeline for prediction on a batch of strings
batch := []string{"This movie is disgustingly good !", "The director tried too much"}
batchResult, err := sentimentPipeline.RunPipeline(batch)
check(err)

// and do whatever we want with it :)
s, err := json.Marshal(batchResult)
check(err)
fmt.Println(string(s))
// {"ClassificationOutputs":[[{"Label":"POSITIVE","Score":0.9998536}],[{"Label":"NEGATIVE","Score":0.99752176}]]}

See also hugot_test.go for further examples.

Use it as a cli: Huggingface 🤗 pipelines from the command line

With hugot you don't need python, pytorch, or even go to run huggingface transformers. Simply install the hugot cli (alpha):

curl https://raw.githubusercontent.com/knights-analytics/hugot/main/scripts/install-hugot-cli.sh | bash

This will install the hugot binary at $HOME/.local/bin/hugot, and the corresponding onnxruntime.so library at $HOME/lib/hugot/onnxruntime.so. The if $HOME/.local/bin is on your $PATH, you can do:

hugot run --model=KnightsAnalytics/distilbert-base-uncased-finetuned-sst-2-english --input=/path/to/input.jsonl --output=/path/to/folder/output --type=textClassification

Hugot will load the model, process the input, and write the results in the output folder. Note that the hugot cli currently expects the input in a specific format: json lines with an "input" key containing the string to process. Example:

{"input": "The director tried too much"}
{"input": "The film was excellent"}

Will produce a file called result_0.jsonl in the output folder with contents:

{"input":"The director tried too much","output":[{"Label":"NEGATIVE","Score":0.99752176}]}
{"input":"The film was excellent","output":[{"Label":"POSITIVE","Score":0.99986285}]}

Note that if --input is not provided, hugot will read from stdin, and if --output is not provided, it will write to stdout. This allows to chain things like:

echo '{"input":"The director tried too much"}' | hugot run --model=/path/to/model --type=textClassification | jq

To be able to run transformers fully from the command line.

Note that the --model parameter can be: 1. the full path to a model to load 2. the name of a huggingface model. Hugot will first try to look for the model at $HOME/hugot, or will try to download the model from huggingface.

Performance Tuning

Firstly, the throughput of onnxruntime depends largely on the size of the input requests. The best batch size is affected by the number of tokens per input, but we find batches of roughly 32 inputs per call to be optimal.

The library defaults to onnxruntime's default tuning settings. These are optimised for latency over throughput, and will attempt to parallelize single threaded calls to onnxruntime over multiple cores.

For maximum throughput, it is best to call a single shared hugot pipeline from multiple goroutines (1 per core), using a channel to pass the input data. In this scenario, the following settings will greatly increase inference throughput.

session, err := hugot.NewSession(
	hugot.WithInterOpNumThreads(1),
	hugot.WithIntraOpNumThreads(1),
	hugot.WithCpuMemArena(false),
	hugot.WithMemPattern(false),
)

InterOpNumThreads and IntraOpNumThreads constricts each goroutine's call to a single core, greatly reducing locking and cache penalties. Disabling CpuMemArena and MemPattern skips pre-allocation of some memory structures, increasing latency, but also throughput efficiency.

For GPU the config above also applies. We are still testing the optimum GPU configuration, whether it is better to run in parallel or with a single thread, and what size of input batch is fastest.

Contributing

Development environment

The easiest way to contribute to hugot is by developing inside a docker container that has the tokenizer and onnxruntime libraries. From the source folder, it should be as easy as:

make start-dev-container

which will download the test models, build the test container, and launch it (see compose-dev), mounting the source code at /home/testuser/repositories/hugot. Then you can attach to the container with e.g. vscode remote extension as testuser. The vscode attached container configuration file can be set to:

{
    "remoteUser": "testuser",
    "workspaceFolder": "/home/testuser/repositories/hugot",
    "extensions": [
		"bierner.markdown-preview-github-styles",
		"golang.go",
		"ms-azuretools.vscode-docker"
	],
    "remoteEnv": {"GOPATH": "/home/testuser/go"}
}

Once you're done, you can tear the container down with:

make stop-dev-container

Alternatively, you can use your IDE devcontainer support, and point it to the Dockerfile.

If you prefer to develop on bare metal, you will need to download the tokenizers.a to /usr/lib/tokenizers.a and onnxruntime.so to /usr/lib/onnxruntime.so.

Run the tests

The full test suite can be run as follows. From the source folder:

make clean run-tests

This will build a test image and run all tests in a container. A testTarget folder will appear in the source directory with the test results.

Contribution process

1. create or find an issue for your contribution
2. fork and develop
3. add tests and make sure the full test suite passes and test coverage does not dip below 80%
4. create a MR linking to the relevant issue

Thank you for contributing to hugot!

Documentation

Index

• func GetPipeline[T pipelines.Pipeline](s *Session, name string) (T, error)
• func NewPipeline[T pipelines.Pipeline](s *Session, pipelineConfig pipelines.PipelineConfig[T]) (T, error)
• type DownloadOptions
  • func NewDownloadOptions() DownloadOptions
• type FeatureExtractionConfig
• type Session
  • func NewSession(options ...WithOption) (*Session, error)
  • func (s *Session) Destroy() error
  • func (s *Session) DownloadModel(modelName string, destination string, options DownloadOptions) (string, error)
  • func (s *Session) GetFeatureExtractionPipeline(name string) (*pipelines.FeatureExtractionPipeline, error)
  • func (s *Session) GetStats() []string
  • func (s *Session) GetTextClassificationPipeline(name string) (*pipelines.TextClassificationPipeline, error)
  • func (s *Session) GetTokenClassificationPipeline(name string) (*pipelines.TokenClassificationPipeline, error)
  • func (s *Session) NewFeatureExtractionPipeline(modelPath string, name string) (*pipelines.FeatureExtractionPipeline, error)
  • func (s *Session) NewTextClassificationPipeline(modelPath string, name string, opts ...TextClassificationOption) (*pipelines.TextClassificationPipeline, error)
  • func (s *Session) NewTokenClassificationPipeline(modelPath string, name string, opts ...TokenClassificationOption) (*pipelines.TokenClassificationPipeline, error)
• type TextClassificationConfig
• type TextClassificationOption
• type TokenClassificationConfig
• type TokenClassificationOption
• type WithOption
  • func WithCoreML(flags uint32) WithOption
  • func WithCpuMemArena(enable bool) WithOption
  • func WithCuda(options map[string]string) WithOption
  • func WithDirectML(deviceID int) WithOption
  • func WithInterOpNumThreads(numThreads int) WithOption
  • func WithIntraOpNumThreads(numThreads int) WithOption
  • func WithMemPattern(enable bool) WithOption
  • func WithOnnxLibraryPath(ortLibraryPath string) WithOption
  • func WithOpenVINO(options map[string]string) WithOption
  • func WithTelemetry() WithOption
  • func WithTensorRT(options map[string]string) WithOption

Functions

func GetPipeline

func GetPipeline[T pipelines.Pipeline](s *Session, name string) (T, error)

GetPipeline can be used to retrieve a pipeline of type T with the given name from the session

func NewPipeline

func NewPipeline[T pipelines.Pipeline](s *Session, pipelineConfig pipelines.PipelineConfig[T]) (T, error)

NewPipeline can be used to create a new pipeline of type T. The initialised pipeline will be returned and it will also be stored in the session object so that all created pipelines can be destroyed with session.Destroy() at once.

Types

type DownloadOptions

type DownloadOptions struct {
	AuthToken             string
	SkipSha               bool
	Branch                string
	MaxRetries            int
	RetryInterval         int
	ConcurrentConnections int
	Verbose               bool
}

DownloadOptions is a struct of options that can be passed to DownloadModel

func NewDownloadOptions

func NewDownloadOptions() DownloadOptions

NewDownloadOptions creates new DownloadOptions struct with default values. Override the values to specify different download options.

type FeatureExtractionConfig

type FeatureExtractionConfig = pipelines.PipelineConfig[*pipelines.FeatureExtractionPipeline]

FeatureExtractionConfig is the configuration for a feature extraction pipeline

type Session

type Session struct {
	// contains filtered or unexported fields
}

Session allows for the creation of new pipelines and holds the pipeline already created.

func NewSession

func NewSession(options ...WithOption) (*Session, error)

NewSession is the main entrypoint to hugot and is used to create a new hugot session object. ortLibraryPath should be the path to onnxruntime.so. If it's the empty string, hugot will try to load the library from the default location (/usr/lib/onnxruntime.so). A new session must be destroyed when it's not needed anymore to avoid memory leaks. See the Destroy method. Note moreover that there can be at most one hugot session active (i.e., the Session object is a singleton), otherwise NewSession will return an error.

func (*Session) Destroy

func (s *Session) Destroy() error

Destroy deletes the hugot session and onnxruntime environment and all initialized pipelines, freeing memory. A hugot session should be destroyed when not neeeded anymore, preferably with a defer() call.

func (*Session) DownloadModel

func (s *Session) DownloadModel(modelName string, destination string, options DownloadOptions) (string, error)

DownloadModel can be used to download a model directly from huggingface. Before the model is downloaded, validation occurs to ensure there is an .onnx and tokenizers.json file. Hugot only works with onnx models.

func (*Session) GetFeatureExtractionPipeline

func (s *Session) GetFeatureExtractionPipeline(name string) (*pipelines.FeatureExtractionPipeline, error)

GetFeatureExtractionPipeline returns a feature extraction pipeline by name. If the name does not exist, it will return an error. Deprecated: use GetPipeline.

func (*Session) GetStats

func (s *Session) GetStats() []string

GetStats returns runtime statistics for all initialized pipelines for profiling purposes. We currently record for each pipeline: the total runtime of the tokenization step the number of batch calls to the tokenization step the average time per tokenization batch call the total runtime of the inference (i.e. onnxruntime) step the number of batch calls to the onnxruntime inference the average time per onnxruntime inference batch call

func (*Session) GetTextClassificationPipeline

func (s *Session) GetTextClassificationPipeline(name string) (*pipelines.TextClassificationPipeline, error)

GetTextClassificationPipeline returns a text classification pipeline by name. If the name does not exist, it will return an error. Deprecated: use GetPipeline.

func (*Session) GetTokenClassificationPipeline

func (s *Session) GetTokenClassificationPipeline(name string) (*pipelines.TokenClassificationPipeline, error)

GetTokenClassificationPipeline returns a token classification pipeline by name. If the name does not exist, it will return an error. Deprecated: use GetPipeline.

func (*Session) NewFeatureExtractionPipeline

func (s *Session) NewFeatureExtractionPipeline(modelPath string, name string) (*pipelines.FeatureExtractionPipeline, error)

NewFeatureExtractionPipeline creates and returns a new feature extraction pipeline object. modelPath should be the path to a folder with the onnx exported transformer model. Name is an identifier for the pipeline (see GetFeatureExtractionPipeline). Deprecated: use NewPipeline

func (*Session) NewTextClassificationPipeline

func (s *Session) NewTextClassificationPipeline(modelPath string, name string, opts ...TextClassificationOption) (*pipelines.TextClassificationPipeline, error)

NewTextClassificationPipeline creates and returns a new text classification pipeline object. modelPath should be the path to a folder with the onnx exported transformer model. Name is an identifier for the pipeline (see GetTextClassificationPipeline). Deprecated: use NewPipeline

func (*Session) NewTokenClassificationPipeline

func (s *Session) NewTokenClassificationPipeline(modelPath string, name string, opts ...TokenClassificationOption) (*pipelines.TokenClassificationPipeline, error)

NewTokenClassificationPipeline creates and returns a new token classification pipeline object. modelPath should be the path to a folder with the onnx exported transformer model. Name is an identifier for the pipeline (see GetTokenClassificationPipeline). Deprecated: use NewPipeline

type TextClassificationConfig

type TextClassificationConfig = pipelines.PipelineConfig[*pipelines.TextClassificationPipeline]

TextClassificationConfig is the configuration for a text classification pipeline

type TextClassificationOption

type TextClassificationOption = pipelines.PipelineOption[*pipelines.TextClassificationPipeline]

TextClassificationOption is an option for a text classification pipeline

type TokenClassificationConfig

type TokenClassificationConfig = pipelines.PipelineConfig[*pipelines.TokenClassificationPipeline]

TokenClassificationConfig is the configuration for a token classification pipeline

type TokenClassificationOption

type TokenClassificationOption = pipelines.PipelineOption[*pipelines.TokenClassificationPipeline]

TokenClassificationOption is an option for a token classification pipeline

type WithOption

type WithOption func(o *ortOptions)

WithOption is the interface for all option functions

func WithCoreML

func WithCoreML(flags uint32) WithOption

WithCoreML Use this function to set the CoreML options flags for the ONNX Runtime configuration. The `flags` parameter represents the CoreML options flags. The `o.coreMLOptions` field in `ortOptions` struct will be set to the provided flags parameter. The `o.coreMLOptionsSet` field in `ortOptions` struct will be set to true.

func WithCpuMemArena

func WithCpuMemArena(enable bool) WithOption

WithCpuMemArena Enable/Disable the usage of the memory arena on CPU. Arena may pre-allocate memory for future usage. Default is true.

func WithCuda

func WithCuda(options map[string]string) WithOption

WithCuda Use this function to set the options for CUDA provider. It takes a pointer to an instance of CUDAProviderOptions struct as input. The options will be applied to the ortOptions struct and the cudaOptionsSet flag will be set to true.

func WithDirectML

func WithDirectML(deviceID int) WithOption

WithDirectML Use this function to set the DirectML device ID for the onnxruntime. By default, this option is not set.

func WithInterOpNumThreads

func WithInterOpNumThreads(numThreads int) WithOption

WithInterOpNumThreads Sets the number of threads used to parallelize execution across separate onnxruntime graph nodes. If unspecified, onnxruntime uses the number of physical CPU cores.

func WithIntraOpNumThreads

func WithIntraOpNumThreads(numThreads int) WithOption

WithIntraOpNumThreads Sets the number of threads used to parallelize execution within onnxruntime graph nodes. If unspecified, onnxruntime uses the number of physical CPU cores.

func WithMemPattern

func WithMemPattern(enable bool) WithOption

WithMemPattern Enable/Disable the memory pattern optimization. If this is enabled memory is preallocated if all shapes are known. Default is true.

func WithOnnxLibraryPath

func WithOnnxLibraryPath(ortLibraryPath string) WithOption

WithOnnxLibraryPath Use this function to set the path to the "onnxruntime.so" or "onnxruntime.dll" function. By default, it will be set to "onnxruntime.so" on non-Windows systems, and "onnxruntime.dll" on Windows.

func WithOpenVINO

func WithOpenVINO(options map[string]string) WithOption

WithOpenVINO Use this function to set the OpenVINO options for the OpenVINO execution provider. The options parameter should be a map of string keys and string values, representing the configuration options. For each key-value pair in the map, the specified option will be set in the OpenVINO execution provider. Once the options are set, the openVINOOptionsSet flag in the ortOptions struct will be set to true. Example usage: WithOpenVINO(map[string]string{"device_type": "CPU", "num_threads": "4"}) This will configure the OpenVINO execution provider to use CPU as the device type and set the number of threads to 4.

func WithTelemetry

func WithTelemetry() WithOption

WithTelemetry Enables telemetry events for the onnxruntime environment. Default is off.

func WithTensorRT

func WithTensorRT(options map[string]string) WithOption

WithTensorRT Use this function to set the options for the TensorRT provider. The options parameter should be a pointer to an instance of TensorRTProviderOptions. By default, the options will be nil and the TensorRT provider will not be used. Example usage:

options := &onnxruntime_go.TensorRTProviderOptions{
    DeviceID: 0,
}
WithTensorRT(options)

Note: For the TensorRT provider to work, the onnxruntime library must be built with TensorRT support.