Cybertron

Cybertron 是一个纯 Go 语言编写的包，基于 spaGO 构建，为 Go 开发者提供了一个简单的接口来使用 NLP 技术，无需其他编程语言或复杂的框架。它专为使用 HuggingFace 模型仓库中可用的预训练 Transformer 模型而设计。

该包主要用于进行推理，未来可能会添加微调功能。

团队欢迎贡献者参与，以加速其发展。

支持的模型

BERT
ELECTRA
BART
PEGASUS
MarianMT

支持的任务

掩码语言建模
监督和零样本文本分类（情感分析、主题分类、意图检测等）
令牌分类（命名实体识别、词性标注等）
抽取式和生成式问答
文本编码（文本嵌入、语义搜索等）
文本生成（翻译、释义、摘要等）
关系抽取

使用方法

要求：

Go 1.21

克隆此仓库或获取库：

go get -u github.com/nlpodyssey/cybertron

Cybertron 支持两种主要使用场景，下面将详细解释。

服务器模式

设置在 .env 文件中配置，Cybertron 会自动加载。或者，也可以通过标志接受配置。

要查看完整列表，请运行：

GOARCH=amd64 go run ./cmd/server -h

输出：

Usage of server:
  -address value
        服务器监听地址
  -allowed-origins value
        允许的源（逗号分隔）
  -loglevel value
        zerolog 全局级别
  -model value
        模型名称（和 models-dir 的子路径）
  -model-conversion value
        模型转换策略（"always"|"missing"|"never"）
  -model-conversion-precision value
        如果模型被转换，使用的浮点精度位数（"32"|"64"）
  -model-download value
        模型下载策略（"always"|"missing"|"never"）
  -models-dir value
        模型的基础目录
  -network value
        服务器监听的网络类型
  -task value
        模型可以执行的推理/计算类型（"textgeneration"|"zero-shot-classification"|"question-answering"|"text-classification"|"token-classification"|"text-encoding"）
  -tls value
        是否启用 TLS（"true"|"false"）
  -tls-cert value
        TLS 证书文件名
  -tls-key value
        TLS 密钥文件名

例如，要使用默认设置在服务器模式下运行 Cybertron 进行机器翻译（如 en 到 it），只需在当前目录创建一个 .env 文件：

echo "CYBERTRON_MODEL=Helsinki-NLP/opus-mt-en-it" > .env
echo "CYBERTRON_MODELS_DIR=models" >> .env
echo "CYBERTRON_MODEL_TASK=text-generation" >> .env

然后执行以下命令：

GOARCH=amd64 go run ./cmd/server -address 0.0.0.0:8080

要测试服务器，运行：

curl -X 'POST' \
  '0.0.0.0:8080/v1/generate' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "input": "You must be the change you wish to see in the world.",
  "parameters": {}
}'

库模式

可以利用多个示例来了解 Cybertron 当前的 NLP 功能。以下是示例列表。

机器翻译（英 -> 意）

GOARCH=amd64 go run ./examples/textgeneration

零样本文本分类

⚠️ 如果 .env 文件中指定的模型不兼容，将返回错误。在这种情况下，从配置文件中删除指定的模型，以使用默认模型。

GOARCH=amd64 go run ./examples/zeroshotclassification politics,business,science,technology,health,culture,sports

依赖项

Cybertron 的主要依赖项包括：

Spago - 一个轻量级的纯 Go 自包含机器学习框架
GoPickle - 一个用于加载 Python 使用 pickle 序列化的数据和 PyTorch 模块文件的 Go 模块
GoTokenizers - 当今最常用分词器的 Go 实现

其余主要用于 gRPC 和 HTTP API 开发。

开发工具

本节适用于想要更改或丰富 Cybertron gRPC 和 HTTP API 的开发者。

首先，你需要在机器上安装 buf。

然后安装以下工具：

go install github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-grpc-gateway \
  github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2 \
  google.golang.org/protobuf/cmd/protoc-gen-go \
  google.golang.org/grpc/cmd/protoc-gen-go-grpc

然后运行以下命令生成 gRPC 和 HTTP API：

go generate ./...