[ML] Trained Models: Add a tech preview lable for rerank model (elast…

…ic#203587) ## Summary Adds a tech preview lable for the`.rerank-v1` model <img width="1365" alt="image" src="https://github.com/user-attachments/assets/dd179f4b-f482-4b1d-beac-74ac3d374446">
Ikuni17 · Dec 10, 2024 · e3ec477 · e3ec477
1 parent 35aeac1
commit e3ec477
Show file tree

Hide file tree

Showing 7 changed files with 77 additions and 4 deletions.
diff --git a/x-pack/platform/packages/shared/ml/trained_models_utils/src/constants/trained_models.ts b/x-pack/platform/packages/shared/ml/trained_models_utils/src/constants/trained_models.ts
@@ -12,6 +12,7 @@ export const ELSER_MODEL_ID = '.elser_model_2';
 export const ELSER_LINUX_OPTIMIZED_MODEL_ID = '.elser_model_2_linux-x86_64';
 export const E5_MODEL_ID = '.multilingual-e5-small';
 export const E5_LINUX_OPTIMIZED_MODEL_ID = '.multilingual-e5-small_linux-x86_64';
+export const RERANK_MODEL_ID = '.rerank-v1';
 export const LANG_IDENT_MODEL_ID = 'lang_ident_model_1';
 export const ELSER_ID_V1 = '.elser_model_1' as const;
 export const LATEST_ELSER_VERSION: ElserVersion = 2;
@@ -148,9 +149,25 @@ export const ELASTIC_MODEL_DEFINITIONS: Record<
         'This E5 model, as defined, hosted, integrated and used in conjunction with our other Elastic Software is covered by our standard warranty.',
     }),
   },
+  [RERANK_MODEL_ID]: {
+    techPreview: true,
+    default: true,
+    hidden: true,
+    modelName: 'rerank',
+    version: 1,
+    config: {
+      input: {
+        field_names: ['input', 'query'],
+      },
+    },
+    description: i18n.translate('xpack.ml.trainedModels.modelsList.rerankDescription', {
+      defaultMessage: 'Elastic Rerank v1',
+    }),
+    type: ['pytorch', 'text_similarity'],
+  },
 } as const);
 
-export type ElasticCuratedModelName = 'elser' | 'e5';
+export type ElasticCuratedModelName = 'elser' | 'e5' | 'rerank';
 
 export interface ModelDefinition {
   /**
@@ -177,6 +194,8 @@ export interface ModelDefinition {
   licenseUrl?: string;
   type?: readonly string[];
   disclaimer?: string;
+  /** Indicates if model is in tech preview */
+  techPreview?: boolean;
 }
 
 export type ModelDefinitionResponse = ModelDefinition & {

diff --git a/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts b/x-pack/platform/plugins/shared/ml/common/types/trained_models.ts
@@ -333,6 +333,7 @@ interface BaseNLPModelItem extends BaseModelItem {
   supported?: boolean;
   state: ModelState | undefined;
   downloadState?: ModelDownloadState;
+  techPreview?: boolean;
 }
 
 /** Model available for download */

diff --git a/...ared/ml/public/application/components/technical_preview_badge/technical_preview_badge.tsx b/...ared/ml/public/application/components/technical_preview_badge/technical_preview_badge.tsx
@@ -25,6 +25,7 @@ export const TechnicalPreviewBadge: FC<{ compressed?: boolean }> = ({ compressed
           'This functionality is in technical preview and may be changed or removed completely in a future release. Elastic will work to fix any issues, but features in technical preview are not subject to the support SLA of official GA features.',
       })}
       tooltipPosition={'right'}
+      css={{ verticalAlign: 'middle' }}
     />
   );
 };
diff --git a/x-pack/platform/plugins/shared/ml/public/application/model_management/models_list.tsx b/x-pack/platform/plugins/shared/ml/public/application/model_management/models_list.tsx
@@ -427,7 +427,8 @@ export const ModelsList: FC<Props> = ({
       render: (item: TrainedModelUIItem) => {
         const { description, model_id: modelId, type } = item;
 
-        const isTechPreview = description?.includes('(Tech Preview)');
+        const isTechPreview =
+          description?.includes('(Tech Preview)') || (isNLPModelItem(item) && item.techPreview);
 
         let descriptionText = description?.replace('(Tech Preview)', '');
 

diff --git a/x-pack/platform/plugins/shared/ml/server/models/model_management/model_provider.test.ts b/x-pack/platform/plugins/shared/ml/server/models/model_management/model_provider.test.ts
@@ -126,6 +126,22 @@ describe('modelsProvider', () => {
           licenseUrl: 'https://huggingface.co/elastic/multilingual-e5-small_linux-x86_64',
           type: ['pytorch', 'text_embedding'],
         },
+        {
+          model_id: '.rerank-v1',
+          techPreview: true,
+          recommended: true,
+          supported: true,
+          hidden: true,
+          modelName: 'rerank',
+          version: 1,
+          config: {
+            input: {
+              field_names: ['input', 'query'],
+            },
+          },
+          description: 'Elastic Rerank v1',
+          type: ['pytorch', 'text_similarity'],
+        },
       ]);
     });
 
@@ -215,6 +231,22 @@ describe('modelsProvider', () => {
           license: 'MIT',
           licenseUrl: 'https://huggingface.co/elastic/multilingual-e5-small_linux-x86_64',
         },
+        {
+          model_id: '.rerank-v1',
+          techPreview: true,
+          recommended: true,
+          supported: true,
+          hidden: true,
+          modelName: 'rerank',
+          version: 1,
+          config: {
+            input: {
+              field_names: ['input', 'query'],
+            },
+          },
+          description: 'Elastic Rerank v1',
+          type: ['pytorch', 'text_similarity'],
+        },
       ]);
     });
   });

diff --git a/x-pack/platform/plugins/shared/ml/server/models/model_management/models_provider.ts b/x-pack/platform/plugins/shared/ml/server/models/model_management/models_provider.ts
@@ -237,12 +237,15 @@ export class ModelsProvider {
     const forDownload = await this.getModelDownloads();
 
     const notDownloaded: TrainedModelUIItem[] = forDownload
-      .filter(({ model_id: modelId, hidden, recommended, supported, disclaimer }) => {
+      .filter(({ model_id: modelId, hidden, recommended, supported, disclaimer, techPreview }) => {
         if (idMap.has(modelId)) {
           const model = idMap.get(modelId)! as NLPModelItem;
           if (recommended) {
             model.recommended = true;
           }
+          if (techPreview) {
+            model.techPreview = true;
+          }
           model.supported = supported;
           model.disclaimer = disclaimer;
         }

diff --git a/x-pack/test/api_integration/apis/ml/trained_models/model_downloads.ts b/x-pack/test/api_integration/apis/ml/trained_models/model_downloads.ts
@@ -45,7 +45,7 @@ export default ({ getService }: FtrProviderContext) => {
         .set(getCommonRequestHeader('1'));
       ml.api.assertResponseStatusCode(200, status, body);
 
-      expect(body.length).to.eql(5);
+      expect(body.length).to.eql(6);
 
       expect(body).to.eql([
         {
@@ -129,6 +129,22 @@ export default ({ getService }: FtrProviderContext) => {
           model_id: '.multilingual-e5-small_linux-x86_64',
           ...(isIntelBased ? { recommended: true, supported: true } : { supported: false }),
         },
+        {
+          model_id: '.rerank-v1',
+          techPreview: true,
+          recommended: true,
+          supported: true,
+          hidden: true,
+          modelName: 'rerank',
+          version: 1,
+          config: {
+            input: {
+              field_names: ['input', 'query'],
+            },
+          },
+          description: 'Elastic Rerank v1',
+          type: ['pytorch', 'text_similarity'],
+        },
       ]);
     });