Golang实现Huggingface Transformer Pipelines的ONNX运行时插件Hugot的使用

Hugot是一个用于在Go应用程序中运行Transformer Pipelines推理和训练的库，兼容Hugging Face的Transformer模型，并支持ONNX格式。

什么是Hugot

Hugot提供了一种简单、可扩展且无麻烦的方式在Go应用程序中运行Transformer Pipelines，如Hugging Face的Transformers Pipelines。它支持以下功能：

• 嵌入向量
• 文本分类
• 命名实体识别
• 微调等AI用例

主要特点

1. Hugging Face兼容性：使用Python Hugging Face Transformer库训练和测试的模型可以导出为ONNX格式，并通过Hugot Pipelines获得与Python版本相同的预测结果
2. 高性能生产使用：仅支持ONNX模型
3. 在自有硬件上运行：无需调用REST API或设置Python RPC服务
4. 简单性：Hugot API允许轻松部署Pipelines而无需编写自己的推理或训练代码

安装和使用

作为库使用

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

func check(err error) {
    if err != nil {
        panic(err.Error())
    }
}

func main() {
    // 启动一个新会话
    session, err := hugot.NewGoSession()
	// 使用XLA后端（需要构建标签"XLA"或"ALL"）：
	// session, err := NewXLASession()
	// 使用ORT后端（需要构建标签"ORT"或"ALL"）：
	// session, err := NewORTSession()
	// 默认在/usr/lib/onnxruntime.so查找onnxruntime.so库
    // 如果onnxruntime.so在其他位置，可以使用WithOnnxLibraryPath明确指定
    // session, err := hugot.NewORTSession(WithOnnxLibraryPath("/path/to/onnxruntime.so"))
	check(err)
	
    // 成功创建的hugot会话在使用完毕后需要销毁
    defer func (session *hugot.Session) {
    err := session.Destroy()
    check(err)
    }(session)

    // 下载一个ONNX情感测试分类模型到当前目录
    // 注意：如果使用NODOWNLOAD构建标签编译库，将排除下载器
    modelPath, err := hugot.DownloadModel("KnightsAnalytics/distilbert-base-uncased-finetuned-sst-2-english", "./models/", hugot.NewDownloadOptions())
    check(err)

    // 创建文本分类Pipeline的配置
    config := TextClassificationConfig{
    ModelPath: modelPath,
    Name:      "testPipeline",
    }
    // 创建Pipeline
    // 注意：Pipeline也会添加到会话对象中，以便一次性销毁所有Pipeline
    sentimentPipeline, err := NewPipeline(session, config)
    check(err)

    // 现在可以在字符串批次上使用Pipeline进行预测
    batch := []string{"This movie is disgustingly good !", "The director tried too much"}
    batchResult, err := sentimentPipeline.RunPipeline(batch)
    check(err)

    // 处理结果
    s, err := json.Marshal(batchResult)
    check(err)
    fmt.Println(string(s))
}
// 输出: {"ClassificationOutputs":[[{"Label":"POSITIVE","Score":0.9998536}],[{"Label":"NEGATIVE","Score":0.99752176}]]}

作为CLI使用

安装Hugot CLI：

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

运行模型：

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

输入文件格式示例：

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

硬件加速

Hugot支持以下加速器后端：

• CUDA（在ONNX Runtime和OpenXLA上测试）
• TensorRT（未测试，仅在ONNX Runtime中可用）
• DirectML（未测试，仅在ONNX Runtime中可用）
• CoreML（未测试，仅在ONNX Runtime中可用）
• OpenVINO（未测试，仅在ONNX Runtime中可用）

CUDA配置示例

opts := []options.WithOption{
  options.WithOnnxLibraryPath("/usr/lib64/onnxruntime-gpu/libonnxruntime.so"),
  options.WithCuda(map[string]string{
    "device_id": "0",
  }),
}
session, err := NewORTSession(opts...)

训练和微调Pipeline

Hugot支持Transformer Pipeline的训练和微调（beta版），目前仅支持FeatureExtractionPipeline。

训练数据集格式示例：

{"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}

性能调优

对于最大吞吐量，建议：

1. 从多个goroutine（每个核心1个）调用单个共享的Hugot Pipeline
2. 使用通道传递输入数据
3. 使用以下设置：

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

文件系统

Hugot使用抽象文件系统，支持各种操作系统文件系统。要使用S3等对象存储，请导入适当的afsc插件：

import _ "github.com/viant/afsc/s3"

当前支持的Pipeline

• featureExtraction
• textClassification
• tokenClassification
• zeroShotClassification

局限性

• 目前仅在amd64-linux上构建/测试
• 仅支持上述Pipeline

Hugot是Knights Analytics团队开发的产品，用于自动化AI驱动的数据管理。