Skip to main content

Batch Prediction

Batch prediction allows you to send a large number of prompts to an LLM at once. In Predibase, this is implemented as a row-wise operation over a Dataset. This means that you can use any Predibase Dataset as input to a batch prediction job.

To better understand how Batch Prediction is possible, let's use an example.

Consider a dataset of wine reviews, which may look like the following:

We can do batch prediction on this dataset in a similar manner to how we make vanilla queries to the model (see Prompting). The only difference is that we must specify a dataset in the Dataset for Batch Inference dropdown menu in the left panel.

Once we have done this, we can write queries in the text box that will be applied as a prompt for each row of the dataset. Because datasets can be quite large, the default behavior is to only run the first 10 rows of the dataset. Here is an example of a batch query:

Columns returned to the user:

  1. sample_index: The index of the row relative to the number of rows in your query. This is helpful for keeping track of which row of the dataset each response corresponds to during analysis.
  2. prompt: The prompt that was sent to the LLM for each row of the dataset.
    • Contains JSON for data from the row that was inserted into the prompt
  3. response: The response from the LLM for each row of the dataset.

Note that the prompt column of the dataset is populated with a stringified JSON object that represents the data in the row. This is because our prompt makes no explicit reference to any particular columns of the dataset. To further customize the prompt with specific column data for a particular row, we can use curly brace notation to dynamically inject values into the prompt. For example, we can isolate the "description" column of the dataset by writing {description} directly in our query. Here is an example of this: