✨ Support text generation task for Nebius AI Studio provider (#1561)

In this PR, I'm adding text-generation task support for the `nebius`
provider (Nebius AI Studio). This is necessary to make completions for
model
[mistralai/Devstral-Small-2505](https://huggingface.co/mistralai/Devstral-Small-2505)
work.

Please let me know if my changes are inconsistent with how you think to
support such models. I will make the necessary changes. Thanks!

Co-authored-by: célina <[email protected]>

Author: kitcat-dev

packages/inference/src/providers/nebius.ts

@@ -14,7 +14,7 @@
  *
  * Thanks!
  */
-import type { FeatureExtractionOutput } from "@huggingface/tasks";
+import type { FeatureExtractionOutput, TextGenerationOutput } from "@huggingface/tasks";
 import type { BodyParams } from "../types.js";
 import { omit } from "../utils/omit.js";
 import {
@@ -40,6 +40,12 @@ interface NebiusEmbeddingsResponse {
 	}>;
 }
 
+interface NebiusTextGenerationOutput extends Omit<TextGenerationOutput, "choices"> {
+	choices: Array<{
+		text: string;
+	}>;
+}
+
 export class NebiusConversationalTask extends BaseConversationalTask {
 	constructor() {
 		super("nebius", NEBIUS_API_BASE_URL);
@@ -50,6 +56,29 @@ export class NebiusTextGenerationTask extends BaseTextGenerationTask {
 	constructor() {
 		super("nebius", NEBIUS_API_BASE_URL);
 	}
+
+	override preparePayload(params: BodyParams): Record<string, unknown> {
+		return {
+			...params.args,
+			model: params.model,
+			prompt: params.args.inputs,
+		};
+	}
+
+	override async getResponse(response: NebiusTextGenerationOutput): Promise<TextGenerationOutput> {
+		if (
+			typeof response === "object" &&
+			"choices" in response &&
+			Array.isArray(response?.choices) &&
+			response.choices.length > 0 &&
+			typeof response.choices[0]?.text === "string"
+		) {
+			return {
+				generated_text: response.choices[0].text,
+			};
+		}
+		throw new InferenceClientProviderOutputError("Received malformed response from Nebius text generation API");
+	}
 }
 
 export class NebiusTextToImageTask extends TaskProviderHelper implements TextToImageTaskHelper {

packages/inference/test/InferenceClient.spec.ts

@@ -1424,6 +1424,13 @@ describe.skip("InferenceClient", () => {
 					status: "live",
 					task: "feature-extraction",
 				},
+				"mistralai/Devstral-Small-2505": {
+					provider: "nebius",
+					providerId: "mistralai/Devstral-Small-2505",
+					hfModelId: "mistralai/Devstral-Small-2505",
+					status: "live",
+					task: "text2text-generation",
+				},
 			};
 
 			it("chatCompletion", async () => {
@@ -1471,6 +1478,19 @@ describe.skip("InferenceClient", () => {
 				expect(res).toBeInstanceOf(Array);
 				expect(res[0]).toEqual(expect.arrayContaining([expect.any(Number)]));
 			});
+
+			it("text2textGeneration", async () => {
+				const res = await client.textGeneration({
+					model: "mistralai/Devstral-Small-2505",
+					provider: "nebius",
+					inputs: "Once upon a time,",
+					temperature: 0,
+					max_tokens: 19,
+				});
+				expect(res).toMatchObject({
+					generated_text: " in a land far, far away, there lived a king who was very fond of flowers.",
+				});
+			});
 		},
 		TIMEOUT
 	);