[ML] add rerank to _inference (#2479)

elastic · Apr 3, 2024 · f95a873 · f95a873
1 parent 141c83e
commit f95a873
Show file tree

Hide file tree

Showing 6 changed files with 139 additions and 6 deletions.
diff --git a/output/openapi/elasticsearch-serverless-openapi.json b/output/openapi/elasticsearch-serverless-openapi.json
diff --git a/output/schema/schema.json b/output/schema/schema.json
diff --git a/output/typescript/types.ts b/output/typescript/types.ts
diff --git a/specification/inference/_types/Results.ts b/specification/inference/_types/Results.ts
@@ -63,6 +63,17 @@ export class CompletionResult {
   result: string
 }
 
+/**
+ * The rerank result object representing a single ranked document
+ * id: the original index of the document in the request
+ * score: the score of the document relative to the query
+ * text: Optional, the text of the document, if requested
+ */
+export class RankedDocument {
+  id: string
+  score: string
+  text?: string
+}
 /**
  * InferenceResult is an aggregation of mutually exclusive variants
  * @variants container
@@ -72,4 +83,5 @@ export class InferenceResult {
   text_embedding?: Array<TextEmbeddingResult>
   sparse_embedding?: Array<SparseEmbeddingResult>
   completion?: Array<CompletionResult>
+  rerank?: Array<RankedDocument>
 }
diff --git a/specification/inference/_types/TaskType.ts b/specification/inference/_types/TaskType.ts
@@ -22,5 +22,7 @@
  */
 export enum TaskType {
   sparse_embedding,
-  text_embedding
+  text_embedding,
+  rerank,
+  completion
 }
diff --git a/specification/inference/inference/InferenceRequest.ts b/specification/inference/inference/InferenceRequest.ts
@@ -40,6 +40,11 @@ export interface Request extends RequestBase {
     inference_id: Id
   }
   body: {
+    /**
+     * Query input, required for rerank task.
+     * Not required for other tasks.
+     */
+    query?: string
     /**
      * Text input to the model.
      * Either a string or an array of strings.