Compatible third party NLP models

The Elastic Stack machine learning features support transformer models that conform to the standard BERT model interface and use the WordPiece tokenization algorithm.

The current list of supported architectures is:

In general, any trained model that has a supported architecture is deployable in Elasticsearch by using eland. However, it is not possible to test every third party model. The following lists are therefore provided for informational purposes only and may not be current. Elastic makes no warranty or assurance that the machine learning features will continue to interoperate with these third party models in the way described, or at all.

These models are listed by NLP task; for more information about those tasks, refer to Overview.

Using SentenceTransformerWrapper:

Using DPREncoderWrapper:

Models used for each NLP task type must output tensors of a specific format to be used in the Elasticsearch NLP pipelines.

Here are the expected outputs for each task type.

Fill mask is a specific kind of token classification; it is the base training task of many transformer models.

For the Elastic stack’s fill mask NLP task to understand the model output, it must have a specific format. It needs to be a float tensor with shape(<number of sequences>, <number of tokens>, <vocab size>).

Here is an example with a single sequence "The capital of [MASK] is Paris" and with vocabulary ["The", "capital", "of", "is", "Paris", "France", "[MASK]"].

Should output:

 [
   [
     [ 0, 0, 0, 0, 0, 0, 0 ], // The
     [ 0, 0, 0, 0, 0, 0, 0 ], // capital
     [ 0, 0, 0, 0, 0, 0, 0 ], // of
     [ 0.01, 0.01, 0.3, 0.01, 0.2, 1.2, 0.1 ], // [MASK]
     [ 0, 0, 0, 0, 0, 0, 0 ], // is
     [ 0, 0, 0, 0, 0, 0, 0 ] // Paris
   ]
]

The predicted value here for [MASK] is "France" with a score of 1.2.

Named entity recognition is a specific token classification task. Each token in the sequence is scored related to a specific set of classification labels. For the Elastic Stack, we use Inside-Outside-Beginning (IOB) tagging. Additionally, only the following classification labels are supported: "O", "B_MISC", "I_MISC", "B_PER", "I_PER", "B_ORG", "I_ORG", "B_LOC", "I_LOC".

The "O" entity label indicates that the current token is outside any entity. "I" indicates that the token is inside an entity. "B" indicates the beginning of an entity. "MISC" is a miscellaneous entity. "LOC" is a location. "PER" is a person. "ORG" is an organization.

The response format must be a float tensor with shape(<number of sequences>, <number of tokens>, <number of classification labels>).

Here is an example with a single sequence "Waldo is in Paris":

 [
   [
//    "O", "B_MISC", "I_MISC", "B_PER", "I_PER", "B_ORG", "I_ORG", "B_LOC", "I_LOC"
     [ 0,  0,         0,       0.4,     0.5,     0,       0.1,     0,       0 ], // Waldo
     [ 1,  0,         0,       0,       0,       0,       0,       0,       0 ], // is
     [ 1,  0,         0,       0,       0,       0,       0,       0,       0 ], // in
     [ 0,  0,         0,       0,       0,       0,       0,       0,       1.0 ] // Paris
   ]
]

Text embedding allows for semantic embedding of text for dense information retrieval. The output of the model must be the specific embedding directly without any additional pooling.

Eland does this wrapping for the aforementioned models. But if supplying your own, the model must output the embedding for each inferred sequence.

With text classification (for example, in tasks like sentiment analysis), the entire sequence is classified. The output of the model must be a float tensor with shape(<number of sequences>, <number of classification labels>).

Here is an example with two sequences for a binary classification model of "happy" and "sad":

 [
   [
//     happy, sad
     [ 0,     1], // first sequence
     [ 1,     0] // second sequence
   ]
]

Zero-shot text classification allows text to be classified for arbitrary labels not necessarily part of the original training. Each sequence is combined with the label given some hypothesis template. The model then scores each of these combinations according to [entailment, neutral, contradiction]. The output of the model must be a float tensor with shape(<number of sequences>, <number of labels>, 3).

Here is an example with a single sequence classified against 4 labels:

 [
   [
//     entailment, neutral, contradiction
     [ 0.5,        0.1,     0.4], // first label
     [ 0,          0,       1], // second label
     [ 1,          0,       0], // third label
     [ 0.7,        0.2,     0.1] // fourth label
   ]
]