DOCSEARCHbyAlgolia
Submit Docs

Cybertron

repository·main·Indexed 43 days ago

https://github.com/nlpodyssey/cybertron

A pure Go package built on spaGO that provides an interface for running inference with pre-trained Transformer models from HuggingFace. It can be used as a library or as a standalone server supporting tasks such as text-generation, zero-shot-classification, question-answering, text-classification, token-classification, text-encoding, and language-modeling.

Tokens
3.6K
Snippets
9
Records
16
Agent score
54%

What's inside Cybertron

  2. Run Cybertron in Server mode

    main

    Cybertron can be run as a standalone server for inference. It automatically loads configuration from a .env file or accepts configuration via CLI flags.

    Configuration via .env

    Create a .env file in your current directory to define the model, directory, and task. For example, to set up a Machine Translation server:

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

    Then start the server:

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

    Testing the Server

    You can test the /v1/generate endpoint using curl:

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

  3. Generate gRPC and HTTP APIs

    main

    To extend or modify the Cybertron gRPC and HTTP APIs, you must use buf and the following Go tools to generate the necessary code:

    1. Install the required protoc plugins:
    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
    1. Generate the APIs:
    go generate ./...
    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

go generate ./...

  4. Initialize and start a Cybertron Server

    main

    To run a Cybertron server, you need a Config and a RequestHandler. The RequestHandler is a task-specific service (like text generation or classification) that implements the required registration interfaces. Use New() to create the server and Start(ctx) to begin serving gRPC and HTTP/2 requests. The server supports both TLS and insecure (h2c) modes.

    // 1. Define configuration
conf := &server.Config{
    Network:        "tcp4",
    Address:        ":8080",
    AllowedOrigins: []string{"*"},
    TLSEnabled:     false,
}

// 2. Resolve a handler for your model
// Assuming 'model' is an instance of a task interface (e.g., textgeneration.Interface)
handler, err := server.ResolveRequestHandler(model)
if err != nil {
    log.Fatal().Err(err).Msg("Failed to resolve handler")
}

// 3. Create and start the server
s := server.New(conf, handler)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

if err := s.Start(ctx); err != nil {
    log.Fatal().Err(err).Msg("Server failed")
}

  5. Run the Cybertron server

    main

    The Cybertron server can be started as a standalone executable. It initializes a server that hosts various NLP tasks (like Text Generation, Question Answering, etc.) by loading a specific model based on the task type.

    Configuration follows this priority order:

    1. Command line flags
    2. Environment variables (loaded from a .env file if present)
    3. Default values

    When running the server, it will automatically log system metrics (CPU and RAM usage) to stderr using zerolog.

  6. Configure Cybertron server via environment variables

    main

    The Cybertron server can be configured using environment variables prefixed with CYBERTRON_. These settings control logging, model loading, and server networking.

    Model Loading Configuration

    • CYBERTRON_MODELS_DIR: Base directory for models.
    • CYBERTRON_MODEL: Model name (and sub-path of MODELS_DIR).
    • CYBERTRON_HUB_ACCESS_TOKEN: Access token for private Hugging Face Hub models.
    • CYBERTRON_MODEL_DOWNLOAD: Policy for downloading models (always|missing|never).
    • CYBERTRON_MODEL_CONVERSION: Policy for model conversion (always|missing|never).
    • CYBERTRON_MODEL_CONVERSION_PRECISION: Floating-point precision for converted models (32|64).
    • CYBERTRON_MODEL_TASK: The specific task type the model performs (see Supported Task Types).

    Server Configuration

    • CYBERTRON_NETWORK: Network type for the server.
    • CYBERTRON_ADDRESS: Server listening address.
    • CYBERTRON_ALLOWED_ORIGINS: Comma-separated list of allowed origins.
    • CYBERTRON_TLS_ENABLED: Enable TLS (true|false).
    • CYBERTRON_TLS_CERT: TLS certificate filename.
    • CYBERTRON_TLS_KEY: TLS key filename.

    Logging

    • CYBERTRON_LOGLEVEL: Global zerolog level.

  7. Configure the Cybertron Server via Environment Variables

    main
    The server supports configuration via environment variables, which are loaded from a .env file if it exists in the working directory. Environment variables are applied after defaults but before command-line flags, allowing flags to override them.

  8. Run Cybertron NLP examples

    main

    You can run pre-configured examples to explore Cybertron's capabilities in library mode.

    Machine Translation (En -> It)

    GOARCH=amd64 go run ./examples/textgeneration

    Zero-Shot Text Classification

    Pass the target labels as a comma-separated list after the command.

    Note: If the model specified in your .env file is incompatible with zero-shot classification, an error will occur. In that case, remove the model from .env to use the default.

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

# Zero-Shot Classification
GOARCH=amd64 go run ./examples/zeroshotclassification politics,business,science,technology,health,culture,sports

  9. Configure the Server

    main

    The Config struct defines the network and security settings for the server.

    FieldTypeDescription
    NetworkstringThe network type (e.g., tcp4). Defaults to tcp4 if empty.
    AddressstringThe listening address (e.g., :8080). Defaults to :8080 if empty.
    AllowedOrigins[]stringList of origins allowed for CORS.
    TLSEnabledboolWhether to enable TLS encryption.
    TLSCertstringPath to the TLS certificate file.
    TLSKeystringPath to the TLS private key file.

  10. Cybertron Server CLI flags

    main

    When running the Cybertron server, you can use the following flags to override or specify configurations:

    FlagDescription
    -addressserver listening address
    -allowed-originsallowed origins (comma separated)
    -loglevelzerolog global level
    -modelmodel name (and sub-path of models-dir)
    -model-conversionmodel conversion policy ("always"|"missing"|"never")
    -model-conversion-precisionfloating-point bits of precision to use if the model is converted ("32"|"64")
    -model-downloadmodel downloading policy ("always"|"missing"|"never")
    -models-dirmodels's base directory
    -networknetwork type for server listening
    -tasktype of inference/computation ("textgeneration"|"zero-shot-classification"|"question-answering"|"text-classification"|"token-classification"|"text-encoding")
    -tlswhether to enable TLS ("true"|"false")
    -tls-certTLS cert filename
    -tls-keyTLS key filename
    Usage of server:
  -address value
        server listening address
  -allowed-origins value
        allowed origins (comma separated)
  -loglevel value
        zerolog global level
  -model value
        model name (and sub-path of models-dir)
  -model-conversion value
        model conversion policy ("always"|"missing"|"never")
  -model-conversion-precision value
        floating-point bits of precision to use if the model is converted ("32"|"64")
  -model-download value
        model downloading policy ("always"|"missing"|"never")
  -models-dir value
        models's base directory
  -network value
        network type for server listening
  -task value
        type of inference/computation that the model can fulfill ("textgeneration"|"zero-shot-classification"|"question-answering"|"text-classification"|"token-classification"|"text-encoding")
  -tls value
        whether to enable TLS ("true"|"false")
  -tls-cert value
        TLS cert filename
  -tls-key value
        TLS key filename

  11. Check if the server is ready for connections

    main

    The ReadyForConnections(dur time.Duration) bool method allows you to poll the server to see if it is successfully listening on its configured address. It returns true if a connection can be established within the specified duration, and false otherwise.

    // Wait up to 5 seconds for the server to be ready
if s.ReadyForConnections(5 * time.Second) {
    fmt.Println("Server is up!")
}

  12. Resolve a RequestHandler from a model

    main

    The ResolveRequestHandler(model any) (RequestHandler, error) function is a factory that automatically selects the correct task-specific server implementation based on the type of the provided model.

    Supported model interfaces include:

    • textgeneration.Interface
    • zeroshotclassifier.Interface
    • questionanswering.Interface
    • textclassification.Interface
    • textencoding.Interface
    • tokenclassification.Interface
    • languagemodeling.Interface
    handler, err := server.ResolveRequestHandler(myModel)