diff --git a/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/shared/ml_inference/utils.ts b/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/shared/ml_inference/utils.ts
index 0b2955cb7f30e..57b9c91b52e57 100644
--- a/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/shared/ml_inference/utils.ts
+++ b/x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/shared/ml_inference/utils.ts
@@ -8,14 +8,9 @@
 import { i18n } from '@kbn/i18n';
 import { TrainedModelConfigResponse } from '@kbn/ml-plugin/common/types/trained_models';
 
-export const NLP_CONFIG_KEYS = [
-  'fill_mask',
-  'ner',
-  'text_classification',
-  'text_embedding',
-  'question_answering',
-  'zero_shot_classification',
-];
+import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/ml_inference_pipeline';
+
+export const NLP_CONFIG_KEYS = Object.values(SUPPORTED_PYTORCH_TASKS);
 export const RECOMMENDED_FIELDS = ['body', 'body_content', 'title'];
 
 export const NLP_DISPLAY_TITLES: Record<string, string | undefined> = {