✨ Add feature extraction task for Nebius (#1436)

Added support for [embedding
models](https://studio.nebius.com/?modality=embedding) for Nebius
provider.

---------

Co-authored-by: Nikita Barinov <[email protected]>
Co-authored-by: Celina Hanouti <[email protected]>

Author: 3 people

packages/inference/src/lib/getProviderHelper.ts

@@ -115,6 +115,7 @@ export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask,
 		"text-to-image": new Nebius.NebiusTextToImageTask(),
 		conversational: new Nebius.NebiusConversationalTask(),
 		"text-generation": new Nebius.NebiusTextGenerationTask(),
+		"feature-extraction": new Nebius.NebiusFeatureExtractionTask(),
 	},
 	novita: {
 		conversational: new Novita.NovitaConversationalTask(),

packages/inference/src/providers/nebius.ts

@@ -14,13 +14,15 @@
  *
  * Thanks!
  */
+import type { FeatureExtractionOutput } from "@huggingface/tasks";
 import { InferenceOutputError } from "../lib/InferenceOutputError";
 import type { BodyParams, UrlParams } from "../types";
 import { omit } from "../utils/omit";
 import {
 	BaseConversationalTask,
 	BaseTextGenerationTask,
 	TaskProviderHelper,
+	type FeatureExtractionTaskHelper,
 	type TextToImageTaskHelper,
 } from "./providerHelper";
 
@@ -32,6 +34,12 @@ interface NebiusBase64ImageGeneration {
 	}>;
 }
 
+interface NebiusEmbeddingsResponse {
+	data: Array<{
+		embedding: number[];
+	}>;
+}
+
 export class NebiusConversationalTask extends BaseConversationalTask {
 	constructor() {
 		super("nebius", NEBIUS_API_BASE_URL);
@@ -59,8 +67,7 @@ export class NebiusTextToImageTask extends TaskProviderHelper implements TextToI
 		};
 	}
 
-	makeRoute(params: UrlParams): string {
-		void params;
+	makeRoute(): string {
 		return "v1/images/generations";
 	}
 
@@ -88,3 +95,24 @@ export class NebiusTextToImageTask extends TaskProviderHelper implements TextToI
 		throw new InferenceOutputError("Expected Nebius text-to-image response format");
 	}
 }
+
+export class NebiusFeatureExtractionTask extends TaskProviderHelper implements FeatureExtractionTaskHelper {
+	constructor() {
+		super("nebius", NEBIUS_API_BASE_URL);
+	}
+
+	preparePayload(params: BodyParams): Record<string, unknown> {
+		return {
+			input: params.args.inputs,
+			model: params.model,
+		};
+	}
+
+	makeRoute(): string {
+		return "v1/embeddings";
+	}
+
+	async getResponse(response: NebiusEmbeddingsResponse): Promise<FeatureExtractionOutput> {
+		return response.data.map((item) => item.embedding);
+	}
+}

packages/inference/test/InferenceClient.spec.ts

@@ -1369,6 +1369,12 @@ describe.skip("InferenceClient", () => {
 					status: "live",
 					task: "text-to-image",
 				},
+				"BAAI/bge-multilingual-gemma2": {
+					providerId: "BAAI/bge-multilingual-gemma2",
+					hfModelId: "BAAI/bge-multilingual-gemma2",
+					status: "live",
+					task: "feature-extraction",
+				},
 			};
 
 			it("chatCompletion", async () => {
@@ -1406,6 +1412,16 @@ describe.skip("InferenceClient", () => {
 				});
 				expect(res).toBeInstanceOf(Blob);
 			});
+
+			it("featureExtraction", async () => {
+				const res = await client.featureExtraction({
+					model: "BAAI/bge-multilingual-gemma2",
+					inputs: "That is a happy person",
+				});
+
+				expect(res).toBeInstanceOf(Array);
+				expect(res[0]).toEqual(expect.arrayContaining([expect.any(Number)]));
+			});
 		},
 		TIMEOUT
 	);