Configure a Text Classification Model

A container administrator can configure a text classification model in the container configuration file. The exact configuration depends on the model type and runtime.

The following example shows a vLLM model configured for text classification:

{
  "models": [
    {
      "name": "twitter-roberta-base-sentiment-latest",
      "path": "cardiffnlp/twitter-roberta-base-sentiment-latest",
      "runtime": "vllm",
      "capabilities": ["TEXT_CLASSIFICATION"]
    }
  ]
}

For ONNX models, the service derives the classification capability from the validated model function and metadata.

The following examples provide sample configurations for reranking models using the vLLM and ONNX runtimes, respectively:

• In this vLLM example, the classification model is downloaded from Hugging Face. Replace the path value with the location of a zip file or a PAR link for other variants.

  {
    "models": [
      {
        "name":"text-classify",
        "path":"nie3e/sentiment-polish-gpt2-small",
        "runtime": "vllm",
        "capabilities":"TEXT_CLASSIFICATION"
      }
    ]
  }

• In this example, a local ONNX file is specified using the path property. The container will look for this file in /privateai/models in the container file system.

  ONNX configurations do not contain the capabilities property. However, for classification models, the labels need to be specified in the config.json file. These labels will be used in the response of the classification API.

  {
    "models": [
      {
        "name":"text-classify",
        "path":"distilbert-base-multilingual-cased-sentiment.zip",
        "function":"classification",
        "labels":["negative","positive","neutral"]
      }
    ]
  }

The following example sends a request to the /classify API endpoint, asking the sentiment model to classify the sentiment of the two given inputs. Ensure any model specified in a request to /classify is configured for TEXT_CLASSIFICATION.

curl -X POST \
  --header 'Content-Type: application/json' \
  --header 'Accept: application/json' \
  --header "Authorization: Bearer $API_KEY" \
  --cacert "$SECRETS_DIR/cert.pem" \
  -d '{
    "model": "sentiment",
    "input": [
      "I loved the service.",
      "The experience was frustrating."
    ]
}' https://localhost:9091/classify

For more information about configuring the Private AI Services Container, see Configure the Private AI Services Container.