Batch predict

This is an experimental feature and is not recommended for use in a production environment. For updates on the progress of the feature or if you want to leave feedback, see the associated GitHub issue.

ML Commons can perform inference on large datasets in an offline asynchronous mode using a model deployed on external model servers. To use the Batch Predict API, you must provide the model_id for an externally hosted model. Amazon SageMaker, Cohere, and OpenAI are currently the only verified external servers that support this API.

For information about user access for this API, see Model access control considerations.

For information about externally hosted models, see Connecting to externally hosted models.

For instructions on how set up batch inference and connector blueprints, see the following:

Path and HTTP methods

  1. POST /_plugins/_ml/models/<model_id>/_batch_predict

Prerequisites

Before using the Batch Predict API, you need to create a connector to the externally hosted model. For example, to create a connector to an OpenAI text-embedding-ada-002 model, send the following request:

  1. POST /_plugins/_ml/connectors/_create
  2. {
  3.   "name": "OpenAI Embedding model",
  4.   "description": "OpenAI embedding model for testing offline batch",
  5.   "version": "1",
  6.   "protocol": "http",
  7.   "parameters": {
  8.     "model": "text-embedding-ada-002",
  9.     "input_file_id": "<your input file id in OpenAI>",
  10.     "endpoint": "/v1/embeddings"
  11.   },
  12.   "credential": {
  13.     "openAI_key": "<your openAI key>"
  14.   },
  15.   "actions": [
  16.     {
  17.       "action_type": "predict",
  18.       "method": "POST",
  19.       "url": "https://api.openai.com/v1/embeddings",
  20.       "headers": {
  21.         "Authorization": "Bearer ${credential.openAI_key}"
  22.       },
  23.       "request_body": "{ \"input\": ${parameters.input}, \"model\": \"${parameters.model}\" }",
  24.       "pre_process_function": "connector.pre_process.openai.embedding",
  25.       "post_process_function": "connector.post_process.openai.embedding"
  26.     },
  27.     {
  28.       "action_type": "batch_predict",
  29.       "method": "POST",
  30.       "url": "https://api.openai.com/v1/batches",
  31.       "headers": {
  32.         "Authorization": "Bearer ${credential.openAI_key}"
  33.       },
  34.       "request_body": "{ \"input_file_id\": \"${parameters.input_file_id}\", \"endpoint\": \"${parameters.endpoint}\", \"completion_window\": \"24h\" }"
  35.     }
  36.   ]
  37. }

copy

The response contains a connector ID that you’ll use in the next steps:

  1. {
  2.   "connector_id": "XU5UiokBpXT9icfOM0vt"
  3. }

Next, register an externally hosted model and provide the connector ID of the created connector:

  1. POST /_plugins/_ml/models/_register?deploy=true
  2. {
  3.     "name": "OpenAI model for realtime embedding and offline batch inference",
  4.     "function_name": "remote",
  5.     "description": "OpenAI text embedding model",
  6.     "connector_id": "XU5UiokBpXT9icfOM0vt"
  7. }

copy

The response contains the task ID for the register operation:

  1. {
  2.   "task_id": "rMormY8B8aiZvtEZIO_j",
  3.   "status": "CREATED",
  4.   "model_id": "lyjxwZABNrAVdFa9zrcZ"
  5. }

To check the status of the operation, provide the task ID to the Tasks API. Once the registration is complete, the task state changes to COMPLETED.

Example request

Once you have completed the prerequisite steps, you can call the Batch Predict API. The parameters in the batch predict request override those defined in the connector:

  1. POST /_plugins/_ml/models/lyjxwZABNrAVdFa9zrcZ/_batch_predict
  2. {
  3.   "parameters": {
  4.     "model": "text-embedding-3-large"
  5.   }
  6. }

copy

Example response

  1. {
  2.   "inference_results": [
  3.     {
  4.       "output": [
  5.         {
  6.           "name": "response",
  7.           "dataAsMap": {
  8.             "id": "batch_<your file id>",
  9.             "object": "batch",
  10.             "endpoint": "/v1/embeddings",
  11.             "errors": null,
  12.             "input_file_id": "file-<your input file id>",
  13.             "completion_window": "24h",
  14.             "status": "validating",
  15.             "output_file_id": null,
  16.             "error_file_id": null,
  17.             "created_at": 1722037257,
  18.             "in_progress_at": null,
  19.             "expires_at": 1722123657,
  20.             "finalizing_at": null,
  21.             "completed_at": null,
  22.             "failed_at": null,
  23.             "expired_at": null,
  24.             "cancelling_at": null,
  25.             "cancelled_at": null,
  26.             "request_counts": {
  27.               "total": 0,
  28.               "completed": 0,
  29.               "failed": 0
  30.             },
  31.             "metadata": null
  32.           }
  33.         }
  34.       ],
  35.       "status_code": 200
  36.     }
  37.   ]
  38. }

For the definition of each field in the result, see OpenAI Batch API. Once the batch inference is complete, you can download the output by calling the OpenAI Files API and providing the file name specified in the id field of the response.