Supported Tasks

Each task has an associated a Dataset Format, a Query Format and a Prediction Format.

A task’s Dataset Format specifies the format of the dataset files. Datasets are prepared by Application Developers when they create Train Jobs and received by Model Developers when they define rafiki.model.BaseModel.train() and rafiki.model.BaseModel.evaluate().

A task’s Query Format specifies the format of queries when they are passed to models. Queries are generated by Application Users when they send queries to Inference Jobs and received by Model Developers when they define rafiki.model.BaseModel.predict().

A task’s Prediction Format specifies the format of predictions made by models. Predictions are generated by Model Developers when they define rafiki.model.BaseModel.predict() and received by Application Users as predictions to their queries sent to Inference Jobs.

IMAGE_CLASSIFICATION

Dataset Format

IMAGE_FILES

  • There is only 1 tag column of class, corresponding to the class of the image as an integer from 0 to k - 1, where k is the total no. of classes.

  • The train & validation dataset’s images should be have the same dimensions W x H and same total no. of classes.

An example:

path,class
image-0-of-class-0.png,0
image-1-of-class-0.png,0
...
image-0-of-class-1.png,1
...
image-99-of-class-9.png,9

Note

You can refer to and run ./examples/datasets/image_files/load_folder_format.py for converting directories of images to Rafiki’s IMAGE_CLASSIFICATION format.

Query Format

A W x H x 3 3D array representing a RGB version of the query image. The query image can be of any dimensions.

Prediction Format

A size-k array of floats, representing the probabilities of each class, by index, from 0 to k-1. For example, the float at index 0 corresponds to the probability of class 0.

POS_TAGGING

Dataset Format

CORPUS

  • Sentences are delimited by \n tokens.

  • There is only 1 tag column of tag corresponding to the POS tag of the token as an integer from 0 to k-1.

An example:

token       tag
Two         3
leading     2
...
line-item   1
veto        5
.           4
\n          0
Professors  6
Philip      6
...
previous    1
presidents  8
.           4
\n          0

Query Format

An array of strings representing a sentence as a list of tokens in that sentence.

Prediction Format

A array of integers representing the list of predicted tag for each token, in sequence, for the sentence.

TABULAR_CLASSIFICATION

Dataset Type

TABULAR

The following optional train arguments are supported:

Train Argument

Description

features

List of feature columns’ names as a list of strings (defaults to first N-1 columns in the CSV file)

target

Target column name as a string (defaults to the last column in the CSV file)

The train & validation datasets should have the same columns.

Query Format

An size-N-1 dictionary representing feature-value pairs.

Prediction Format

A size-k list of floats, representing the probabilities of each class from 0 to k-1 for the target column.

TABULAR_REGRESSION

Dataset Type

TABULAR

The following optional train arguments are supported:

Train Argument

Description

features

List of feature columns’ names as a list of strings (defaults to first N-1 columns in the CSV file)

target

Target column name as a string (defaults to the last column in the CSV file)

The train & validation datasets should have the same columns.

An example of the dataset follows:

density,bodyfat,age,weight,height,neck,chest,abdomen,hip,thigh,knee,ankle,biceps,forearm,wrist
1.0708,12.3,23,154.25,67.75,36.2,93.1,85.2,94.5,59,37.3,21.9,32,27.4,17.1
1.0853,6.1,22,173.25,72.25,38.5,93.6,83,98.7,58.7,37.3,23.4,30.5,28.9,18.2
1.0414,25.3,22,154,66.25,34,95.8,87.9,99.2,59.6,38.9,24,28.8,25.2,16.6
...

Query Format

An size-N-1 dictionary representing feature-value pairs.

Prediction Format

A float, representing the value of the target column.

SPEECH_RECOGNITION

Speech recognition for the English language.

Dataset Type

AUDIO_FILES

The audios.csv should be of a .CSV format with 3 columns of wav_filename, wav_filesize and transcript.

For each row,

wav_filename should be a file path to a .wav audio file within the archive, relative to the root of the directory. Each audio file’s sample rate must equal to 16kHz.

wav_filesize should be an integer representing the size of the .wav audio file, in number of bytes.

transcript should be a string of the true transcript for the audio file. Transcripts should only contain the following alphabets:

a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z


'

An example of audios.csv follows:

wav_filename,wav_filesize,transcript
6930-81414-0000.wav,412684,audio transcript one
6930-81414-0001.wav,559564,audio transcript two
...
672-122797-0005.wav,104364,audio transcript one thousand
...
1995-1837-0001.wav,279404,audio transcript three thousand

Query Format

A Base64-encoded string of the bytes of the audio as a 16kHz .wav file

Prediction Format

A string, representing the predicted transcript for the audio.