[Enterprise Search] update supported ml models filter (#141307)

Updated the list of keys for the inference_config that we support for models in enterprise search.
elastic · Sep 21, 2022 · 77f7329 · 77f7329
1 parent b77de15
commit 77f7329
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/...ns/enterprise_search_content/components/search_index/pipelines/ml_inference/utils.test.ts b/...ns/enterprise_search_content/components/search_index/pipelines/ml_inference/utils.test.ts
@@ -31,11 +31,6 @@ describe('ml inference utils', () => {
             ner: {},
           },
         }),
-        makeFakeModel({
-          inference_config: {
-            classification: {},
-          },
-        }),
         makeFakeModel({
           inference_config: {
             text_classification: {},
@@ -53,6 +48,16 @@ describe('ml inference utils', () => {
             },
           },
         }),
+        makeFakeModel({
+          inference_config: {
+            question_answering: {},
+          },
+        }),
+        makeFakeModel({
+          inference_config: {
+            fill_mask: {},
+          },
+        }),
       ];
 
       for (const model of models) {
@@ -61,7 +66,14 @@ describe('ml inference utils', () => {
     });
 
     it('returns false for expected models', () => {
-      const models: TrainedModelConfigResponse[] = [makeFakeModel({})];
+      const models: TrainedModelConfigResponse[] = [
+        makeFakeModel({}),
+        makeFakeModel({
+          inference_config: {
+            classification: {},
+          },
+        }),
+      ];
 
       for (const model of models) {
         expect(isSupportedMLModel(model)).toBe(false);

diff --git a/...cations/enterprise_search_content/components/search_index/pipelines/ml_inference/utils.ts b/...cations/enterprise_search_content/components/search_index/pipelines/ml_inference/utils.ts
@@ -11,10 +11,11 @@ import { TrainedModelConfigResponse } from '@kbn/ml-plugin/common/types/trained_
 import { AddInferencePipelineFormErrors, InferencePipelineConfiguration } from './types';
 
 const NLP_CONFIG_KEYS = [
+  'fill_mask',
   'ner',
-  'classification',
   'text_classification',
   'text_embedding',
+  'question_answering',
   'zero_shot_classification',
 ];
 export const isSupportedMLModel = (model: TrainedModelConfigResponse): boolean => {