Allow single input for generation endpoint (huggingface#803)

28Smiles · web-flow · commit f7c29737e4a9 · 2024-07-17T18:49:26.000+02:00
According to spec inputs can be either a single `string` or a `string`
array, with the result beeing either a single `TextGenerationOutput` or
and array of them.
diff --git a/packages/inference/src/tasks/nlp/textGeneration.ts b/packages/inference/src/tasks/nlp/textGeneration.ts
@@ -1,6 +1,7 @@
 import type { TextGenerationInput, TextGenerationOutput } from "@huggingface/tasks";
 import { InferenceOutputError } from "../../lib/InferenceOutputError";
 import type { BaseArgs, Options } from "../../types";
+import { toArray } from "../../utils/toArray";
 import { request } from "../custom/request";
 
 export type { TextGenerationInput, TextGenerationOutput };
@@ -12,10 +13,12 @@ export async function textGeneration(
 	args: BaseArgs & TextGenerationInput,
 	options?: Options
 ): Promise<TextGenerationOutput> {
-	const res = await request<TextGenerationOutput[]>(args, {
-		...options,
-		taskHint: "text-generation",
-	});
+	const res = toArray(
+		await request<TextGenerationOutput | TextGenerationOutput[]>(args, {
+			...options,
+			taskHint: "text-generation",
+		})
+	);
 	const isValidOutput = Array.isArray(res) && res.every((x) => typeof x?.generated_text === "string");
 	if (!isValidOutput) {
 		throw new InferenceOutputError("Expected Array<{generated_text: string}>");
