@huggingface/inference 3.7.0 → 3.8.0

package/src/tasks/cv/textToImage.ts

@@ -1,48 +1,15 @@
-import type { TextToImageInput, TextToImageOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
-import type { BaseArgs, InferenceProvider, Options } from "../../types";
-import { delay } from "../../utils/delay";
-import { omit } from "../../utils/omit";
+import type { TextToImageInput } from "@huggingface/tasks";
+import { getProviderHelper } from "../../lib/getProviderHelper";
+import { makeRequestOptions } from "../../lib/makeRequestOptions";
+import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
 export type TextToImageArgs = BaseArgs & TextToImageInput;
 
-interface Base64ImageGeneration {
-	data: Array<{
-		b64_json: string;
-	}>;
-}
-interface OutputUrlImageGeneration {
-	output: string[];
-}
-interface HyperbolicTextToImageOutput {
-	images: Array<{ image: string }>;
-}
-
-interface BlackForestLabsResponse {
-	id: string;
-	polling_url: string;
-}
-
 interface TextToImageOptions extends Options {
 	outputType?: "url" | "blob";
 }
 
-function getResponseFormatArg(provider: InferenceProvider) {
-	switch (provider) {
-		case "fal-ai":
-			return { sync_mode: true };
-		case "nebius":
-			return { response_format: "b64_json" };
-		case "replicate":
-			return undefined;
-		case "together":
-			return { response_format: "base64" };
-		default:
-			return undefined;
-	}
-}
-
 /**
  * This task reads some text input and outputs an image.
  * Recommended model: stabilityai/stable-diffusion-2
@@ -56,108 +23,13 @@ export async function textToImage(
 	options?: TextToImageOptions & { outputType?: undefined | "blob" }
 ): Promise<Blob>;
 export async function textToImage(args: TextToImageArgs, options?: TextToImageOptions): Promise<Blob | string> {
-	const payload =
-		!args.provider || args.provider === "hf-inference" || args.provider === "sambanova"
-			? args
-			: {
-					...omit(args, ["inputs", "parameters"]),
-					...args.parameters,
-					...getResponseFormatArg(args.provider),
-					prompt: args.inputs,
-			  };
-	const { data: res } = await innerRequest<
-		| TextToImageOutput
-		| Base64ImageGeneration
-		| OutputUrlImageGeneration
-		| BlackForestLabsResponse
-		| HyperbolicTextToImageOutput
-	>(payload, {
+	const provider = args.provider ?? "hf-inference";
+	const providerHelper = getProviderHelper(provider, "text-to-image");
+	const { data: res } = await innerRequest<Record<string, unknown>>(args, providerHelper, {
 		...options,
 		task: "text-to-image",
 	});
 
-	if (res && typeof res === "object") {
-		if (args.provider === "black-forest-labs" && "polling_url" in res && typeof res.polling_url === "string") {
-			return await pollBflResponse(res.polling_url, options?.outputType);
-		}
-		if (args.provider === "fal-ai" && "images" in res && Array.isArray(res.images) && res.images[0].url) {
-			if (options?.outputType === "url") {
-				return res.images[0].url;
-			} else {
-				const image = await fetch(res.images[0].url);
-				return await image.blob();
-			}
-		}
-		if (
-			args.provider === "hyperbolic" &&
-			"images" in res &&
-			Array.isArray(res.images) &&
-			res.images[0] &&
-			typeof res.images[0].image === "string"
-		) {
-			if (options?.outputType === "url") {
-				return `data:image/jpeg;base64,${res.images[0].image}`;
-			}
-			const base64Response = await fetch(`data:image/jpeg;base64,${res.images[0].image}`);
-			return await base64Response.blob();
-		}
-		if ("data" in res && Array.isArray(res.data) && res.data[0].b64_json) {
-			const base64Data = res.data[0].b64_json;
-			if (options?.outputType === "url") {
-				return `data:image/jpeg;base64,${base64Data}`;
-			}
-			const base64Response = await fetch(`data:image/jpeg;base64,${base64Data}`);
-			return await base64Response.blob();
-		}
-		if ("output" in res && Array.isArray(res.output)) {
-			if (options?.outputType === "url") {
-				return res.output[0];
-			}
-			const urlResponse = await fetch(res.output[0]);
-			const blob = await urlResponse.blob();
-			return blob;
-		}
-	}
-	const isValidOutput = res && res instanceof Blob;
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Blob");
-	}
-	if (options?.outputType === "url") {
-		const b64 = await res.arrayBuffer().then((buf) => Buffer.from(buf).toString("base64"));
-		return `data:image/jpeg;base64,${b64}`;
-	}
-	return res;
-}
-
-async function pollBflResponse(url: string, outputType?: "url" | "blob"): Promise<Blob> {
-	const urlObj = new URL(url);
-	for (let step = 0; step < 5; step++) {
-		await delay(1000);
-		console.debug(`Polling Black Forest Labs API for the result... ${step + 1}/5`);
-		urlObj.searchParams.set("attempt", step.toString(10));
-		const resp = await fetch(urlObj, { headers: { "Content-Type": "application/json" } });
-		if (!resp.ok) {
-			throw new InferenceOutputError("Failed to fetch result from black forest labs API");
-		}
-		const payload = await resp.json();
-		if (
-			typeof payload === "object" &&
-			payload &&
-			"status" in payload &&
-			typeof payload.status === "string" &&
-			payload.status === "Ready" &&
-			"result" in payload &&
-			typeof payload.result === "object" &&
-			payload.result &&
-			"sample" in payload.result &&
-			typeof payload.result.sample === "string"
-		) {
-			if (outputType === "url") {
-				return payload.result.sample;
-			}
-			const image = await fetch(payload.result.sample);
-			return await image.blob();
-		}
-	}
-	throw new InferenceOutputError("Failed to fetch result from black forest labs API");
+	const { url, info } = await makeRequestOptions(args, providerHelper, { ...options, task: "text-to-image" });
+	return providerHelper.getResponse(res, url, info.headers as Record<string, string>, options?.outputType);
 }

package/src/tasks/cv/textToVideo.ts

@@ -1,74 +1,27 @@
 import type { TextToVideoInput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
-import { isUrl } from "../../lib/isUrl";
-import { pollFalResponse, type FalAiQueueOutput } from "../../providers/fal-ai";
-import type { BaseArgs, InferenceProvider, Options } from "../../types";
-import { omit } from "../../utils/omit";
+import { getProviderHelper } from "../../lib/getProviderHelper";
+import { makeRequestOptions } from "../../lib/makeRequestOptions";
+import type { FalAiQueueOutput } from "../../providers/fal-ai";
+import type { NovitaOutput } from "../../providers/novita";
+import type { ReplicateOutput } from "../../providers/replicate";
+import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
-import { typedInclude } from "../../utils/typedInclude";
 
 export type TextToVideoArgs = BaseArgs & TextToVideoInput;
 
 export type TextToVideoOutput = Blob;
 
-interface ReplicateOutput {
-	output: string;
-}
-
-interface NovitaOutput {
-	video: {
-		video_url: string;
-	};
-}
-
-const SUPPORTED_PROVIDERS = ["fal-ai", "novita", "replicate"] as const satisfies readonly InferenceProvider[];
-
 export async function textToVideo(args: TextToVideoArgs, options?: Options): Promise<TextToVideoOutput> {
-	if (!args.provider || !typedInclude(SUPPORTED_PROVIDERS, args.provider)) {
-		throw new Error(
-			`textToVideo inference is only supported for the following providers: ${SUPPORTED_PROVIDERS.join(", ")}`
-		);
-	}
-
-	const payload =
-		args.provider === "fal-ai" || args.provider === "replicate" || args.provider === "novita"
-			? { ...omit(args, ["inputs", "parameters"]), ...args.parameters, prompt: args.inputs }
-			: args;
-	const { data, requestContext } = await innerRequest<FalAiQueueOutput | ReplicateOutput | NovitaOutput>(payload, {
-		...options,
-		task: "text-to-video",
-	});
-
-	if (args.provider === "fal-ai") {
-		return await pollFalResponse(
-			data as FalAiQueueOutput,
-			requestContext.url,
-			requestContext.info.headers as Record<string, string>
-		);
-	} else if (args.provider === "novita") {
-		const isValidOutput =
-			typeof data === "object" &&
-			!!data &&
-			"video" in data &&
-			typeof data.video === "object" &&
-			!!data.video &&
-			"video_url" in data.video &&
-			typeof data.video.video_url === "string" &&
-			isUrl(data.video.video_url);
-		if (!isValidOutput) {
-			throw new InferenceOutputError("Expected { video: { video_url: string } }");
-		}
-		const urlResponse = await fetch((data as NovitaOutput).video.video_url);
-		return await urlResponse.blob();
-	} else {
-		/// TODO: Replicate: handle the case where the generation request "times out" / is async (ie output is null)
-		/// https://replicate.com/docs/topics/predictions/create-a-prediction
-		const isValidOutput =
-			typeof data === "object" && !!data && "output" in data && typeof data.output === "string" && isUrl(data.output);
-		if (!isValidOutput) {
-			throw new InferenceOutputError("Expected { output: string }");
+	const provider = args.provider ?? "hf-inference";
+	const providerHelper = getProviderHelper(provider, "text-to-video");
+	const { data: response } = await innerRequest<FalAiQueueOutput | ReplicateOutput | NovitaOutput>(
+		args,
+		providerHelper,
+		{
+			...options,
+			task: "text-to-video",
 		}
-		const urlResponse = await fetch(data.output);
-		return await urlResponse.blob();
-	}
+	);
+	const { url, info } = await makeRequestOptions(args, providerHelper, { ...options, task: "text-to-video" });
+	return providerHelper.getResponse(response, url, info.headers as Record<string, string>);
 }

package/src/tasks/cv/zeroShotImageClassification.ts

@@ -1,5 +1,5 @@
 import type { ZeroShotImageClassificationInput, ZeroShotImageClassificationOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options, RequestArgs } from "../../types";
 import { base64FromBytes } from "../../utils/base64FromBytes";
 import { innerRequest } from "../../utils/request";
@@ -44,15 +44,11 @@ export async function zeroShotImageClassification(
 	args: ZeroShotImageClassificationArgs,
 	options?: Options
 ): Promise<ZeroShotImageClassificationOutput> {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "zero-shot-image-classification");
 	const payload = await preparePayload(args);
-	const { data: res } = await innerRequest<ZeroShotImageClassificationOutput>(payload, {
+	const { data: res } = await innerRequest<ZeroShotImageClassificationOutput>(payload, providerHelper, {
 		...options,
 		task: "zero-shot-image-classification",
 	});
-	const isValidOutput =
-		Array.isArray(res) && res.every((x) => typeof x.label === "string" && typeof x.score === "number");
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<{label: string, score: number}>");
-	}
-	return res;
+	return providerHelper.getResponse(res);
 }

package/src/tasks/index.ts

@@ -4,21 +4,23 @@ export * from "./custom/streamingRequest";
 
 // Audio tasks
 export * from "./audio/audioClassification";
+export * from "./audio/audioToAudio";
 export * from "./audio/automaticSpeechRecognition";
 export * from "./audio/textToSpeech";
-export * from "./audio/audioToAudio";
 
 // Computer Vision tasks
 export * from "./cv/imageClassification";
 export * from "./cv/imageSegmentation";
+export * from "./cv/imageToImage";
 export * from "./cv/imageToText";
 export * from "./cv/objectDetection";
 export * from "./cv/textToImage";
-export * from "./cv/imageToImage";
-export * from "./cv/zeroShotImageClassification";
 export * from "./cv/textToVideo";
+export * from "./cv/zeroShotImageClassification";
 
 // Natural Language Processing tasks
+export * from "./nlp/chatCompletion";
+export * from "./nlp/chatCompletionStream";
 export * from "./nlp/featureExtraction";
 export * from "./nlp/fillMask";
 export * from "./nlp/questionAnswering";
@@ -31,13 +33,11 @@ export * from "./nlp/textGenerationStream";
 export * from "./nlp/tokenClassification";
 export * from "./nlp/translation";
 export * from "./nlp/zeroShotClassification";
-export * from "./nlp/chatCompletion";
-export * from "./nlp/chatCompletionStream";
 
 // Multimodal tasks
 export * from "./multimodal/documentQuestionAnswering";
 export * from "./multimodal/visualQuestionAnswering";
 
 // Tabular tasks
-export * from "./tabular/tabularRegression";
 export * from "./tabular/tabularClassification";
+export * from "./tabular/tabularRegression";

package/src/tasks/multimodal/documentQuestionAnswering.ts

@@ -3,11 +3,10 @@ import type {
 	DocumentQuestionAnsweringInputData,
 	DocumentQuestionAnsweringOutput,
 } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options, RequestArgs } from "../../types";
 import { base64FromBytes } from "../../utils/base64FromBytes";
 import { innerRequest } from "../../utils/request";
-import { toArray } from "../../utils/toArray";
 
 /// Override the type to properly set inputs.image as Blob
 export type DocumentQuestionAnsweringArgs = BaseArgs &
@@ -20,6 +19,7 @@ export async function documentQuestionAnswering(
 	args: DocumentQuestionAnsweringArgs,
 	options?: Options
 ): Promise<DocumentQuestionAnsweringOutput[number]> {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "document-question-answering");
 	const reqArgs: RequestArgs = {
 		...args,
 		inputs: {
@@ -30,26 +30,11 @@ export async function documentQuestionAnswering(
 	} as RequestArgs;
 	const { data: res } = await innerRequest<DocumentQuestionAnsweringOutput | DocumentQuestionAnsweringOutput[number]>(
 		reqArgs,
+		providerHelper,
 		{
 			...options,
 			task: "document-question-answering",
 		}
 	);
-	const output = toArray(res);
-	const isValidOutput =
-		Array.isArray(output) &&
-		output.every(
-			(elem) =>
-				typeof elem === "object" &&
-				!!elem &&
-				typeof elem?.answer === "string" &&
-				(typeof elem.end === "number" || typeof elem.end === "undefined") &&
-				(typeof elem.score === "number" || typeof elem.score === "undefined") &&
-				(typeof elem.start === "number" || typeof elem.start === "undefined")
-		);
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<{answer: string, end?: number, score?: number, start?: number}>");
-	}
-
-	return output[0];
+	return providerHelper.getResponse(res);
 }

package/src/tasks/multimodal/visualQuestionAnswering.ts

@@ -3,7 +3,7 @@ import type {
 	VisualQuestionAnsweringInputData,
 	VisualQuestionAnsweringOutput,
 } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options, RequestArgs } from "../../types";
 import { base64FromBytes } from "../../utils/base64FromBytes";
 import { innerRequest } from "../../utils/request";
@@ -19,6 +19,7 @@ export async function visualQuestionAnswering(
 	args: VisualQuestionAnsweringArgs,
 	options?: Options
 ): Promise<VisualQuestionAnsweringOutput[number]> {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "visual-question-answering");
 	const reqArgs: RequestArgs = {
 		...args,
 		inputs: {
@@ -28,18 +29,9 @@ export async function visualQuestionAnswering(
 		},
 	} as RequestArgs;
 
-	const { data: res } = await innerRequest<VisualQuestionAnsweringOutput>(reqArgs, {
+	const { data: res } = await innerRequest<VisualQuestionAnsweringOutput>(reqArgs, providerHelper, {
 		...options,
 		task: "visual-question-answering",
 	});
-
-	const isValidOutput =
-		Array.isArray(res) &&
-		res.every(
-			(elem) => typeof elem === "object" && !!elem && typeof elem?.answer === "string" && typeof elem.score === "number"
-		);
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<{answer: string, score: number}>");
-	}
-	return res[0];
+	return providerHelper.getResponse(res);
 }

package/src/tasks/nlp/chatCompletion.ts

@@ -1,5 +1,5 @@
 import type { ChatCompletionInput, ChatCompletionOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -10,25 +10,10 @@ export async function chatCompletion(
 	args: BaseArgs & ChatCompletionInput,
 	options?: Options
 ): Promise<ChatCompletionOutput> {
-	const { data: res } = await innerRequest<ChatCompletionOutput>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "conversational");
+	const { data: response } = await innerRequest<ChatCompletionOutput>(args, providerHelper, {
 		...options,
-		task: "text-generation",
-		chatCompletion: true,
+		task: "conversational",
 	});
-	const isValidOutput =
-		typeof res === "object" &&
-		Array.isArray(res?.choices) &&
-		typeof res?.created === "number" &&
-		typeof res?.id === "string" &&
-		typeof res?.model === "string" &&
-		/// Together.ai and Nebius do not output a system_fingerprint
-		(res.system_fingerprint === undefined ||
-			res.system_fingerprint === null ||
-			typeof res.system_fingerprint === "string") &&
-		typeof res?.usage === "object";
-
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected ChatCompletionOutput");
-	}
-	return res;
+	return providerHelper.getResponse(response);
 }

package/src/tasks/nlp/chatCompletionStream.ts

@@ -1,4 +1,5 @@
 import type { ChatCompletionInput, ChatCompletionStreamOutput } from "@huggingface/tasks";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerStreamingRequest } from "../../utils/request";
 
@@ -9,9 +10,9 @@ export async function* chatCompletionStream(
 	args: BaseArgs & ChatCompletionInput,
 	options?: Options
 ): AsyncGenerator<ChatCompletionStreamOutput> {
-	yield* innerStreamingRequest<ChatCompletionStreamOutput>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "conversational");
+	yield* innerStreamingRequest<ChatCompletionStreamOutput>(args, providerHelper, {
 		...options,
-		task: "text-generation",
-		chatCompletion: true,
+		task: "conversational",
 	});
 }

package/src/tasks/nlp/featureExtraction.ts

@@ -1,5 +1,5 @@
 import type { FeatureExtractionInput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -17,25 +17,10 @@ export async function featureExtraction(
 	args: FeatureExtractionArgs,
 	options?: Options
 ): Promise<FeatureExtractionOutput> {
-	const { data: res } = await innerRequest<FeatureExtractionOutput>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "feature-extraction");
+	const { data: res } = await innerRequest<FeatureExtractionOutput>(args, providerHelper, {
 		...options,
 		task: "feature-extraction",
 	});
-	let isValidOutput = true;
-
-	const isNumArrayRec = (arr: unknown[], maxDepth: number, curDepth = 0): boolean => {
-		if (curDepth > maxDepth) return false;
-		if (arr.every((x) => Array.isArray(x))) {
-			return arr.every((x) => isNumArrayRec(x as unknown[], maxDepth, curDepth + 1));
-		} else {
-			return arr.every((x) => typeof x === "number");
-		}
-	};
-
-	isValidOutput = Array.isArray(res) && isNumArrayRec(res, 3, 0);
-
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<number[][][] | number[][] | number[] | number>");
-	}
-	return res;
+	return providerHelper.getResponse(res);
 }

package/src/tasks/nlp/fillMask.ts

@@ -1,5 +1,5 @@
 import type { FillMaskInput, FillMaskOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -9,23 +9,10 @@ export type FillMaskArgs = BaseArgs & FillMaskInput;
  * Tries to fill in a hole with a missing word (token to be precise). That’s the base task for BERT models.
  */
 export async function fillMask(args: FillMaskArgs, options?: Options): Promise<FillMaskOutput> {
-	const { data: res } = await innerRequest<FillMaskOutput>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "fill-mask");
+	const { data: res } = await innerRequest<FillMaskOutput>(args, providerHelper, {
 		...options,
 		task: "fill-mask",
 	});
-	const isValidOutput =
-		Array.isArray(res) &&
-		res.every(
-			(x) =>
-				typeof x.score === "number" &&
-				typeof x.sequence === "string" &&
-				typeof x.token === "number" &&
-				typeof x.token_str === "string"
-		);
-	if (!isValidOutput) {
-		throw new InferenceOutputError(
-			"Expected Array<{score: number, sequence: string, token: number, token_str: string}>"
-		);
-	}
-	return res;
+	return providerHelper.getResponse(res);
 }

package/src/tasks/nlp/questionAnswering.ts

@@ -1,5 +1,5 @@
 import type { QuestionAnsweringInput, QuestionAnsweringOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -12,29 +12,14 @@ export async function questionAnswering(
 	args: QuestionAnsweringArgs,
 	options?: Options
 ): Promise<QuestionAnsweringOutput[number]> {
-	const { data: res } = await innerRequest<QuestionAnsweringOutput | QuestionAnsweringOutput[number]>(args, {
-		...options,
-		task: "question-answering",
-	});
-
-	const isValidOutput = Array.isArray(res)
-		? res.every(
-				(elem) =>
-					typeof elem === "object" &&
-					!!elem &&
-					typeof elem.answer === "string" &&
-					typeof elem.end === "number" &&
-					typeof elem.score === "number" &&
-					typeof elem.start === "number"
-		  )
-		: typeof res === "object" &&
-		  !!res &&
-		  typeof res.answer === "string" &&
-		  typeof res.end === "number" &&
-		  typeof res.score === "number" &&
-		  typeof res.start === "number";
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<{answer: string, end: number, score: number, start: number}>");
-	}
-	return Array.isArray(res) ? res[0] : res;
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "question-answering");
+	const { data: res } = await innerRequest<QuestionAnsweringOutput | QuestionAnsweringOutput[number]>(
+		args,
+		providerHelper,
+		{
+			...options,
+			task: "question-answering",
+		}
+	);
+	return providerHelper.getResponse(res);
 }

package/src/tasks/nlp/sentenceSimilarity.ts

@@ -1,5 +1,5 @@
 import type { SentenceSimilarityInput, SentenceSimilarityOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -12,14 +12,10 @@ export async function sentenceSimilarity(
 	args: SentenceSimilarityArgs,
 	options?: Options
 ): Promise<SentenceSimilarityOutput> {
-	const { data: res } = await innerRequest<SentenceSimilarityOutput>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "sentence-similarity");
+	const { data: res } = await innerRequest<SentenceSimilarityOutput>(args, providerHelper, {
 		...options,
 		task: "sentence-similarity",
 	});
-
-	const isValidOutput = Array.isArray(res) && res.every((x) => typeof x === "number");
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected number[]");
-	}
-	return res;
+	return providerHelper.getResponse(res);
 }

package/src/tasks/nlp/summarization.ts

@@ -1,5 +1,5 @@
 import type { SummarizationInput, SummarizationOutput } from "@huggingface/tasks";
-import { InferenceOutputError } from "../../lib/InferenceOutputError";
+import { getProviderHelper } from "../../lib/getProviderHelper";
 import type { BaseArgs, Options } from "../../types";
 import { innerRequest } from "../../utils/request";
 
@@ -9,13 +9,10 @@ export type SummarizationArgs = BaseArgs & SummarizationInput;
  * This task is well known to summarize longer text into shorter text. Be careful, some models have a maximum length of input. That means that the summary cannot handle full books for instance. Be careful when choosing your model.
  */
 export async function summarization(args: SummarizationArgs, options?: Options): Promise<SummarizationOutput> {
-	const { data: res } = await innerRequest<SummarizationOutput[]>(args, {
+	const providerHelper = getProviderHelper(args.provider ?? "hf-inference", "summarization");
+	const { data: res } = await innerRequest<SummarizationOutput[]>(args, providerHelper, {
 		...options,
 		task: "summarization",
 	});
-	const isValidOutput = Array.isArray(res) && res.every((x) => typeof x?.summary_text === "string");
-	if (!isValidOutput) {
-		throw new InferenceOutputError("Expected Array<{summary_text: string}>");
-	}
-	return res?.[0];
+	return providerHelper.getResponse(res);
 }