[ML] Adding trained models utils package (#154046)

Creates a package for sharing general trained model util code.
Initially only contains the constants which had been copied in the
enterprise search plugin.

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

.github/CODEOWNERS

@@ -462,6 +462,7 @@ x-pack/plugins/ml @elastic/ml-ui
 x-pack/packages/ml/query_utils @elastic/ml-ui
 x-pack/packages/ml/route_utils @elastic/ml-ui
 x-pack/packages/ml/string_hash @elastic/ml-ui
+x-pack/packages/ml/trained_models_utils @elastic/ml-ui
 x-pack/packages/ml/url_state @elastic/ml-ui
 packages/kbn-monaco @elastic/appex-sharedux
 x-pack/plugins/monitoring_collection @elastic/infra-monitoring-ui

package.json

@@ -477,6 +477,7 @@
     "@kbn/ml-query-utils": "link:x-pack/packages/ml/query_utils",
     "@kbn/ml-route-utils": "link:x-pack/packages/ml/route_utils",
     "@kbn/ml-string-hash": "link:x-pack/packages/ml/string_hash",
+    "@kbn/ml-trained-models-utils": "link:x-pack/packages/ml/trained_models_utils",
     "@kbn/ml-url-state": "link:x-pack/packages/ml/url_state",
     "@kbn/monaco": "link:packages/kbn-monaco",
     "@kbn/monitoring-collection-plugin": "link:x-pack/plugins/monitoring_collection",

tsconfig.base.json

@@ -918,6 +918,8 @@
       "@kbn/ml-route-utils/*": ["x-pack/packages/ml/route_utils/*"],
       "@kbn/ml-string-hash": ["x-pack/packages/ml/string_hash"],
       "@kbn/ml-string-hash/*": ["x-pack/packages/ml/string_hash/*"],
+      "@kbn/ml-trained-models-utils": ["x-pack/packages/ml/trained_models_utils"],
+      "@kbn/ml-trained-models-utils/*": ["x-pack/packages/ml/trained_models_utils/*"],
       "@kbn/ml-url-state": ["x-pack/packages/ml/url_state"],
       "@kbn/ml-url-state/*": ["x-pack/packages/ml/url_state/*"],
       "@kbn/monaco": ["packages/kbn-monaco"],

x-pack/.i18nrc.json

@@ -45,7 +45,7 @@
     "xpack.main": "legacy/plugins/xpack_main",
     "xpack.maps": ["plugins/maps"],
     "xpack.aiops": ["packages/ml/aiops_components", "plugins/aiops"],
-    "xpack.ml": ["packages/ml/date_picker", "plugins/ml"],
+    "xpack.ml": ["packages/ml/date_picker", "packages/ml/trained_models_utils", "plugins/ml"],
     "xpack.monitoring": ["plugins/monitoring"],
     "xpack.osquery": ["plugins/osquery"],
     "xpack.painlessLab": "plugins/painless_lab",

x-pack/packages/ml/trained_models_utils/README.md

@@ -0,0 +1,3 @@
+# @kbn/ml-trained-models-utils
+
+Static utils for interacting with ML trained models

x-pack/packages/ml/trained_models_utils/index.ts

@@ -0,0 +1,17 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export {
+  BUILT_IN_MODEL_TAG,
+  BUILT_IN_MODEL_TYPE,
+  DEPLOYMENT_STATE,
+  SUPPORTED_PYTORCH_TASKS,
+  TRAINED_MODEL_TYPE,
+  type DeploymentState,
+  type SupportedPytorchTasksType,
+  type TrainedModelType,
+} from './src/constants/trained_models';

x-pack/packages/ml/trained_models_utils/jest.config.js

@@ -0,0 +1,12 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+module.exports = {
+  preset: '@kbn/test',
+  rootDir: '../../../..',
+  roots: ['<rootDir>/x-pack/packages/ml/trained_models_utils'],
+};

x-pack/packages/ml/trained_models_utils/kibana.jsonc

@@ -0,0 +1,5 @@
+{
+  "type": "shared-common",
+  "id": "@kbn/ml-trained-models-utils",
+  "owner": "@elastic/ml-ui"
+}

x-pack/packages/ml/trained_models_utils/package.json

@@ -0,0 +1,6 @@
+{
+  "name": "@kbn/ml-trained-models-utils",
+  "private": true,
+  "version": "1.0.0",
+  "license": "Elastic License 2.0"
+}

x-pack/packages/ml/trained_models_utils/src/constants/trained_models.ts

@@ -29,6 +29,7 @@ export const SUPPORTED_PYTORCH_TASKS = {
   TEXT_CLASSIFICATION: 'text_classification',
   TEXT_EMBEDDING: 'text_embedding',
   FILL_MASK: 'fill_mask',
+  TEXT_EXPANSION: 'text_expansion',
 } as const;
 export type SupportedPytorchTasksType =
   typeof SUPPORTED_PYTORCH_TASKS[keyof typeof SUPPORTED_PYTORCH_TASKS];
@@ -37,3 +38,5 @@ export const BUILT_IN_MODEL_TYPE = i18n.translate(
   'xpack.ml.trainedModels.modelsList.builtInModelLabel',
   { defaultMessage: 'built-in' }
 );
+
+export const BUILT_IN_MODEL_TAG = 'prepackaged';

x-pack/packages/ml/trained_models_utils/tsconfig.json

@@ -0,0 +1,21 @@
+{
+  "extends": "../../../../tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "target/types",
+    "types": [
+      "jest",
+      "node",
+      "react"
+    ]
+  },
+  "include": [
+    "**/*.ts",
+    "**/*.tsx",
+  ],
+  "exclude": [
+    "target/**/*"
+  ],
+  "kbn_references": [
+    "@kbn/i18n",
+  ]
+}

x-pack/plugins/enterprise_search/common/ml_inference_pipeline/index.test.ts

@@ -11,13 +11,11 @@ import {
   MlTrainedModelConfig,
   MlTrainedModelStats,
 } from '@elastic/elasticsearch/lib/api/types';
-import { BUILT_IN_MODEL_TAG } from '@kbn/ml-plugin/common/constants/data_frame_analytics';
-import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-plugin/common/constants/trained_models';
+import { SUPPORTED_PYTORCH_TASKS, BUILT_IN_MODEL_TAG } from '@kbn/ml-trained-models-utils';
 
 import { MlInferencePipeline, TrainedModelState } from '../types/pipelines';
 
 import {
-  BUILT_IN_MODEL_TAG as LOCAL_BUILT_IN_MODEL_TAG,
   generateMlInferencePipelineBody,
   getMlModelTypesForModelConfig,
   getRemoveProcessorForInferenceType,
@@ -25,7 +23,6 @@ import {
   parseMlInferenceParametersFromPipeline,
   parseModelStateFromStats,
   parseModelStateReasonFromStats,
-  SUPPORTED_PYTORCH_TASKS as LOCAL_SUPPORTED_PYTORCH_TASKS,
 } from '.';
 
 const mockModel: MlTrainedModelConfig = {
@@ -65,10 +62,6 @@ describe('getMlModelTypesForModelConfig lib function', () => {
     const response = getMlModelTypesForModelConfig(builtInMockModel);
     expect(response.sort()).toEqual(expected.sort());
   });
-
-  it('local BUILT_IN_MODEL_TAG matches ml plugin', () => {
-    expect(LOCAL_BUILT_IN_MODEL_TAG).toEqual(BUILT_IN_MODEL_TAG);
-  });
 });
 
 describe('getRemoveProcessorForInferenceType lib function', () => {
@@ -105,11 +98,6 @@ describe('getRemoveProcessorForInferenceType lib function', () => {
 
 describe('getSetProcessorForInferenceType lib function', () => {
   const destinationField = 'dest';
-
-  it('local LOCAL_SUPPORTED_PYTORCH_TASKS matches ml plugin', () => {
-    expect(SUPPORTED_PYTORCH_TASKS).toEqual(LOCAL_SUPPORTED_PYTORCH_TASKS);
-  });
-
   it('should return expected value for TEXT_CLASSIFICATION', () => {
     const inferenceType = SUPPORTED_PYTORCH_TASKS.TEXT_CLASSIFICATION;
 

x-pack/plugins/enterprise_search/common/ml_inference_pipeline/index.ts

@@ -13,6 +13,12 @@ import {
   MlTrainedModelStats,
 } from '@elastic/elasticsearch/lib/api/types';
 
+import {
+  SUPPORTED_PYTORCH_TASKS,
+  TRAINED_MODEL_TYPE,
+  BUILT_IN_MODEL_TAG,
+} from '@kbn/ml-trained-models-utils';
+
 import {
   MlInferencePipeline,
   CreateMlInferencePipelineParameters,
@@ -20,23 +26,7 @@ import {
   InferencePipelineInferenceConfig,
 } from '../types/pipelines';
 
-// Getting an error importing this from @kbn/ml-plugin/common/constants/data_frame_analytics'
-// So defining it locally for now with a test to make sure it matches.
-export const BUILT_IN_MODEL_TAG = 'prepackaged';
-
-// Getting an error importing this from @kbn/ml-plugin/common/constants/trained_models'
-// So defining it locally for now with a test to make sure it matches.
-export const SUPPORTED_PYTORCH_TASKS = {
-  FILL_MASK: 'fill_mask',
-  NER: 'ner',
-  QUESTION_ANSWERING: 'question_answering',
-  TEXT_CLASSIFICATION: 'text_classification',
-  TEXT_EMBEDDING: 'text_embedding',
-  ZERO_SHOT_CLASSIFICATION: 'zero_shot_classification',
-} as const;
-
 export const ELSER_TASK_TYPE = 'text_expansion';
-export const LANG_IDENT_MODEL_TYPE = 'lang_ident';
 
 export interface MlInferencePipelineParams {
   description?: string;
@@ -213,7 +203,10 @@ export const parseModelStateFromStats = (
   model?: Partial<MlTrainedModelStats> & Partial<MlTrainedModelConfig>,
   modelTypes?: string[]
 ) => {
-  if (model?.model_type === LANG_IDENT_MODEL_TYPE || modelTypes?.includes(LANG_IDENT_MODEL_TYPE))
+  if (
+    model?.model_type === TRAINED_MODEL_TYPE.LANG_IDENT ||
+    modelTypes?.includes(TRAINED_MODEL_TYPE.LANG_IDENT)
+  )
     return TrainedModelState.Started;
   switch (model?.deployment_stats?.state) {
     case 'started':

x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/search_index/pipelines/ml_inference/inference_config.tsx

@@ -12,10 +12,9 @@ import { useValues } from 'kea';
 import { EuiSpacer, EuiText } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 
-import {
-  getMlModelTypesForModelConfig,
-  SUPPORTED_PYTORCH_TASKS,
-} from '../../../../../../../common/ml_inference_pipeline';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
+
+import { getMlModelTypesForModelConfig } from '../../../../../../../common/ml_inference_pipeline';
 import { getMLType } from '../../../shared/ml_inference/utils';
 
 import { MLInferenceLogic } from './ml_inference_logic';

x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/search_index/pipelines/ml_inference/target_field_help_text.tsx

@@ -12,10 +12,9 @@ import { EuiText } from '@elastic/eui';
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage, FormattedNumber } from '@kbn/i18n-react';
 
-import {
-  getMlModelTypesForModelConfig,
-  SUPPORTED_PYTORCH_TASKS,
-} from '../../../../../../../common/ml_inference_pipeline';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
+
+import { getMlModelTypesForModelConfig } from '../../../../../../../common/ml_inference_pipeline';
 import { TrainedModel } from '../../../../api/ml_models/ml_trained_models_logic';
 import { getMLType } from '../../../shared/ml_inference/utils';
 

x-pack/plugins/enterprise_search/public/applications/enterprise_search_content/components/shared/ml_inference/utils.ts

@@ -8,10 +8,7 @@
 import { i18n } from '@kbn/i18n';
 import { TrainedModelConfigResponse } from '@kbn/ml-plugin/common/types/trained_models';
 
-import {
-  LANG_IDENT_MODEL_TYPE,
-  SUPPORTED_PYTORCH_TASKS,
-} from '../../../../../../common/ml_inference_pipeline';
+import { TRAINED_MODEL_TYPE, SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 
 export const NLP_CONFIG_KEYS: string[] = Object.values(SUPPORTED_PYTORCH_TASKS);
 export const RECOMMENDED_FIELDS = ['body', 'body_content', 'title'];
@@ -52,7 +49,7 @@ export const NLP_DISPLAY_TITLES: Record<string, string | undefined> = {
 export const isSupportedMLModel = (model: TrainedModelConfigResponse): boolean => {
   return (
     Object.keys(model.inference_config).some((key) => NLP_CONFIG_KEYS.includes(key)) ||
-    model.model_type === LANG_IDENT_MODEL_TYPE
+    model.model_type === TRAINED_MODEL_TYPE.LANG_IDENT
   );
 };
 
@@ -75,7 +72,7 @@ export const getMLType = (modelTypes: string[]): string => {
       return type;
     }
   }
-  if (modelTypes?.includes(LANG_IDENT_MODEL_TYPE)) return LANG_IDENT_MODEL_TYPE;
+  if (modelTypes?.includes(TRAINED_MODEL_TYPE.LANG_IDENT)) return TRAINED_MODEL_TYPE.LANG_IDENT;
   return modelTypes?.[0] ?? '';
 };
 

x-pack/plugins/enterprise_search/tsconfig.json

@@ -30,6 +30,7 @@
     "@kbn/spaces-plugin",
     "@kbn/i18n",
     "@kbn/ml-plugin",
+    "@kbn/ml-trained-models-utils",
     "@kbn/guided-onboarding-plugin",
     "@kbn/i18n-react",
     "@kbn/test-jest-helpers",

x-pack/plugins/ml/common/constants/data_frame_analytics.ts

@@ -34,6 +34,4 @@ export const INDEX_CREATED_BY = {
   DATA_FRAME_ANALYTICS: 'data-frame-analytics',
 } as const;
 
-export const BUILT_IN_MODEL_TAG = 'prepackaged';
-
 export type JobMapNodeTypes = typeof JOB_MAP_NODE_TYPES[keyof typeof JOB_MAP_NODE_TYPES];

x-pack/plugins/ml/common/types/trained_models.ts

@@ -5,10 +5,10 @@
  * 2.0.
  */
 import * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import type { DeploymentState, TrainedModelType } from '@kbn/ml-trained-models-utils';
 import type { DataFrameAnalyticsConfig } from './data_frame_analytics';
 import type { FeatureImportanceBaseline, TotalFeatureImportance } from './feature_importance';
 import type { XOR } from './common';
-import type { DeploymentState, TrainedModelType } from '../constants/trained_models';
 import type { MlSavedObjectType } from './saved_objects';
 
 export interface IngestStats {

x-pack/plugins/ml/public/application/data_frame_analytics/pages/components/analytics_selector/analytics_id_selector.tsx

@@ -23,7 +23,7 @@ import {
 import { i18n } from '@kbn/i18n';
 import { FormattedMessage } from '@kbn/i18n-react';
 
-import { BUILT_IN_MODEL_TAG } from '../../../../../../common/constants/data_frame_analytics';
+import { BUILT_IN_MODEL_TAG } from '@kbn/ml-trained-models-utils';
 import { useTrainedModelsApiService } from '../../../../services/ml_api_service/trained_models';
 import { GetDataFrameAnalyticsResponse } from '../../../../services/ml_api_service/data_frame_analytics';
 import { useToastNotificationService } from '../../../../services/toast_notification_service';

x-pack/plugins/ml/public/application/model_management/model_actions.tsx

@@ -10,7 +10,11 @@ import { i18n } from '@kbn/i18n';
 import { isPopulatedObject } from '@kbn/ml-is-populated-object';
 import { EuiToolTip } from '@elastic/eui';
 import React, { useCallback, useMemo } from 'react';
-import { BUILT_IN_MODEL_TAG } from '../../../common/constants/data_frame_analytics';
+import {
+  DEPLOYMENT_STATE,
+  TRAINED_MODEL_TYPE,
+  BUILT_IN_MODEL_TAG,
+} from '@kbn/ml-trained-models-utils';
 import { useTrainedModelsApiService } from '../services/ml_api_service/trained_models';
 import { getUserConfirmationProvider } from './force_stop_dialog';
 import { useToastNotificationService } from '../services/toast_notification_service';
@@ -19,7 +23,6 @@ import { useMlKibana, useMlLocator, useNavigateToPath } from '../contexts/kibana
 import { getAnalysisType } from '../../../common/util/analytics_utils';
 import { DataFrameAnalysisConfigType } from '../../../common/types/data_frame_analytics';
 import { ML_PAGES } from '../../../common/constants/locator';
-import { DEPLOYMENT_STATE, TRAINED_MODEL_TYPE } from '../../../common/constants/trained_models';
 import { isTestable } from './test_models';
 import { ModelItem } from './models_list';
 

x-pack/plugins/ml/public/application/model_management/models_list.tsx

@@ -27,6 +27,7 @@ import { FIELD_FORMAT_IDS } from '@kbn/field-formats-plugin/common';
 import { isPopulatedObject } from '@kbn/ml-is-populated-object';
 import { usePageUrlState } from '@kbn/ml-url-state';
 import { useTimefilter } from '@kbn/ml-date-picker';
+import { BUILT_IN_MODEL_TYPE, BUILT_IN_MODEL_TAG } from '@kbn/ml-trained-models-utils';
 import { useModelActions } from './model_actions';
 import { ModelsTableToConfigMapping } from '.';
 import { ModelsBarStats, StatsBar } from '../components/stats_bar';
@@ -37,7 +38,6 @@ import {
   TrainedModelConfigResponse,
   TrainedModelStat,
 } from '../../../common/types/trained_models';
-import { BUILT_IN_MODEL_TAG } from '../../../common/constants/data_frame_analytics';
 import { DeleteModelsModal } from './delete_models_modal';
 import { ML_PAGES } from '../../../common/constants/locator';
 import { ListingPageUrlState } from '../../../common/types/common';
@@ -46,7 +46,6 @@ import { useTableSettings } from '../data_frame_analytics/pages/analytics_manage
 import { useToastNotificationService } from '../services/toast_notification_service';
 import { useFieldFormatter } from '../contexts/kibana/use_field_formatter';
 import { useRefresh } from '../routing/use_refresh';
-import { BUILT_IN_MODEL_TYPE } from '../../../common/constants/trained_models';
 import { SavedObjectsWarning } from '../components/saved_objects_warning';
 import { TestTrainedModelFlyout } from './test_models';
 

x-pack/plugins/ml/public/application/model_management/test_models/models/inference_base.ts

@@ -10,8 +10,8 @@ import * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { i18n } from '@kbn/i18n';
 
 import { map } from 'rxjs/operators';
+import { SupportedPytorchTasksType } from '@kbn/ml-trained-models-utils';
 import { MLHttpFetchError } from '../../../../../common/util/errors';
-import { SupportedPytorchTasksType } from '../../../../../common/constants/trained_models';
 import { trainedModelsApiProvider } from '../../../services/ml_api_service/trained_models';
 import { getInferenceInfoComponent } from './inference_info';
 

x-pack/plugins/ml/public/application/model_management/test_models/models/ner/ner_inference.ts

@@ -7,12 +7,12 @@
 
 import * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { i18n } from '@kbn/i18n';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import type { InferResponse } from '../inference_base';
 import { getGeneralInputComponent } from '../text_input';
 import { getNerOutputComponent } from './ner_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 
 export type FormattedNerResponse = Array<{
   value: string;

x-pack/plugins/ml/public/application/model_management/test_models/models/question_answering/question_answering_inference.ts

@@ -9,11 +9,11 @@ import { BehaviorSubject } from 'rxjs';
 import { map } from 'rxjs/operators';
 import { i18n } from '@kbn/i18n';
 import { estypes } from '@elastic/elasticsearch';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import type { InferResponse } from '../inference_base';
 import { getQuestionAnsweringInput } from './question_answering_input';
 import { getQuestionAnsweringOutputComponent } from './question_answering_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 
 export interface RawQuestionAnsweringResponse {

x-pack/plugins/ml/public/application/model_management/test_models/models/text_classification/fill_mask_inference.ts

@@ -8,12 +8,12 @@
 import { i18n } from '@kbn/i18n';
 import { estypes } from '@elastic/elasticsearch';
 import { map } from 'rxjs/operators';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import type { TextClassificationResponse, RawTextClassificationResponse } from './common';
 import { processResponse, processInferenceResult } from './common';
 import { getGeneralInputComponent } from '../text_input';
 import { getFillMaskOutputComponent } from './fill_mask_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 
 const MASK = '[MASK]';

x-pack/plugins/ml/public/application/model_management/test_models/models/text_classification/text_classification_inference.ts

@@ -7,12 +7,12 @@
 
 import { i18n } from '@kbn/i18n';
 import { estypes } from '@elastic/elasticsearch';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import { processInferenceResult, processResponse } from './common';
 import type { TextClassificationResponse, RawTextClassificationResponse } from './common';
 import { getGeneralInputComponent } from '../text_input';
 import { getTextClassificationOutputComponent } from './text_classification_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 
 export class TextClassificationInference extends InferenceBase<TextClassificationResponse> {

x-pack/plugins/ml/public/application/model_management/test_models/models/text_classification/zero_shot_classification_inference.ts

@@ -9,6 +9,7 @@ import { i18n } from '@kbn/i18n';
 import { BehaviorSubject } from 'rxjs';
 import { map } from 'rxjs/operators';
 import { estypes } from '@elastic/elasticsearch';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import { processInferenceResult, processResponse } from './common';
@@ -16,7 +17,6 @@ import type { TextClassificationResponse, RawTextClassificationResponse } from '
 
 import { getZeroShotClassificationInput } from './zero_shot_classification_input';
 import { getTextClassificationOutputComponent } from './text_classification_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 
 export class ZeroShotClassificationInference extends InferenceBase<TextClassificationResponse> {
   protected inferenceType = SUPPORTED_PYTORCH_TASKS.ZERO_SHOT_CLASSIFICATION;

x-pack/plugins/ml/public/application/model_management/test_models/models/text_embedding/text_embedding_inference.ts

@@ -7,11 +7,11 @@
 
 import { i18n } from '@kbn/i18n';
 import { estypes } from '@elastic/elasticsearch';
+import { SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { InferenceBase, INPUT_TYPE } from '../inference_base';
 import type { InferResponse } from '../inference_base';
 import { getGeneralInputComponent } from '../text_input';
 import { getTextEmbeddingOutputComponent } from './text_embedding_output';
-import { SUPPORTED_PYTORCH_TASKS } from '../../../../../../common/constants/trained_models';
 import { trainedModelsApiProvider } from '../../../../services/ml_api_service/trained_models';
 
 export interface RawTextEmbeddingResponse {

x-pack/plugins/ml/public/application/model_management/test_models/selected_model.tsx

@@ -8,6 +8,7 @@
 import * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import React, { FC, useMemo, useEffect } from 'react';
 
+import { TRAINED_MODEL_TYPE, SUPPORTED_PYTORCH_TASKS } from '@kbn/ml-trained-models-utils';
 import { NerInference } from './models/ner';
 import { QuestionAnsweringInference } from './models/question_answering';
 
@@ -20,10 +21,6 @@ import {
 
 import { TextEmbeddingInference } from './models/text_embedding';
 
-import {
-  TRAINED_MODEL_TYPE,
-  SUPPORTED_PYTORCH_TASKS,
-} from '../../../../common/constants/trained_models';
 import { useMlApiContext } from '../../contexts/kibana';
 import { InferenceInputForm } from './models/inference_input_form';
 import { InferrerType } from './models';

x-pack/plugins/ml/public/application/model_management/test_models/utils.ts

@@ -9,8 +9,8 @@ import {
   TRAINED_MODEL_TYPE,
   DEPLOYMENT_STATE,
   SUPPORTED_PYTORCH_TASKS,
-} from '../../../../common/constants/trained_models';
-import type { SupportedPytorchTasksType } from '../../../../common/constants/trained_models';
+  type SupportedPytorchTasksType,
+} from '@kbn/ml-trained-models-utils';
 import type { ModelItem } from '../models_list';
 
 const PYTORCH_TYPES = Object.values(SUPPORTED_PYTORCH_TASKS);

x-pack/plugins/ml/server/routes/management.ts

@@ -6,14 +6,13 @@
  */
 
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import { BUILT_IN_MODEL_TYPE, BUILT_IN_MODEL_TAG } from '@kbn/ml-trained-models-utils';
 import { wrapError } from '../client/error_wrapper';
 import type { RouteInitialization } from '../types';
 import { listTypeSchema } from './schemas/management_schema';
 
 import { jobServiceProvider } from '../models/job_service';
 import { checksFactory } from '../saved_objects';
-import { BUILT_IN_MODEL_TAG } from '../../common/constants/data_frame_analytics';
-import { BUILT_IN_MODEL_TYPE } from '../../common/constants/trained_models';
 import type {
   AnomalyDetectionManagementItems,
   AnalyticsManagementItems,

x-pack/plugins/ml/tsconfig.json

@@ -46,6 +46,7 @@
     "@kbn/es-query",
     "@kbn/ui-theme",
     "@kbn/ml-string-hash",
+    "@kbn/ml-trained-models-utils",
     "@kbn/core-http-browser",
     "@kbn/home-plugin",
     "@kbn/management-plugin",
@@ -86,6 +87,7 @@
     "@kbn/saved-objects-finder-plugin",
     "@kbn/monaco",
     "@kbn/repo-info",
+    "@kbn/ml-trained-models-utils",
     "@kbn/ml-route-utils",
   ],
 }

yarn.lock

@@ -4573,6 +4573,10 @@
   version "0.0.0"
   uid ""
 
+"@kbn/ml-trained-models-utils@link:x-pack/packages/ml/trained_models_utils":
+  version "0.0.0"
+  uid ""
+
 "@kbn/ml-url-state@link:x-pack/packages/ml/url_state":
   version "0.0.0"
   uid ""