@huggingface/inference 2.6.3 → 2.6.5

package/dist/index.d.ts

@@ -1,10 +1,11 @@
+
 export interface Options {
 	/**
 	 * (Default: true) Boolean. If a request 503s and wait_for_model is set to false, the request will be retried with the same parameters but with wait_for_model set to true.
 	 */
 	retry_on_error?: boolean;
 	/**
-	 * (Default: true). Boolean. There is a cache layer on the inference API to speedup requests we have already seen. Most models can use those results as is as models are deterministic (meaning the results will be the same anyway). However if you use a non deterministic model, you can set this parameter to prevent the caching mechanism from being used resulting in a real new query.
+	 * (Default: true). Boolean. There is a cache layer on Inference API (serverless) to speedup requests we have already seen. Most models can use those results as is as models are deterministic (meaning the results will be the same anyway). However if you use a non deterministic model, you can set this parameter to prevent the caching mechanism from being used resulting in a real new query.
 	 */
 	use_cache?: boolean;
 	/**
@@ -24,6 +25,10 @@ export interface Options {
 	 * Custom fetch function to use instead of the default one, for example to use a proxy or edit headers.
 	 */
 	fetch?: typeof fetch;
+	/**
+	 * Abort Controller signal to use for request interruption.
+	 */
+	signal?: AbortSignal;
 
 	/**
 	 * (Default: "same-origin"). String | Boolean. Credentials to use for the request. If this is a string, it will be passed straight on. If it's a boolean, true will be "include" and false will not send credentials at all.
@@ -35,35 +40,51 @@ export type InferenceTask =
 	| "audio-classification"
 	| "audio-to-audio"
 	| "automatic-speech-recognition"
-	| "conversational"
 	| "depth-estimation"
 	| "document-question-answering"
 	| "feature-extraction"
 	| "fill-mask"
+	| "graph-ml"
 	| "image-classification"
+	| "image-feature-extraction"
 	| "image-segmentation"
+	| "image-text-to-text"
+	| "image-to-3d"
 	| "image-to-image"
 	| "image-to-text"
+	| "image-to-video"
+	| "mask-generation"
+	| "multiple-choice"
 	| "object-detection"
-	| "video-classification"
 	| "question-answering"
 	| "reinforcement-learning"
+	| "robotics"
 	| "sentence-similarity"
 	| "summarization"
 	| "table-question-answering"
+	| "table-to-text"
 	| "tabular-classification"
 	| "tabular-regression"
+	| "tabular-to-text"
 	| "text-classification"
 	| "text-generation"
+	| "text-retrieval"
+	| "text-to-3d"
+	| "text-to-audio"
 	| "text-to-image"
 	| "text-to-speech"
 	| "text-to-video"
+	| "text2text-generation"
+	| "time-series-forecasting"
 	| "token-classification"
 	| "translation"
 	| "unconditional-image-generation"
+	| "video-classification"
 	| "visual-question-answering"
+	| "voice-activity-detection"
 	| "zero-shot-classification"
-	| "zero-shot-image-classification";
+	| "zero-shot-image-classification"
+	| "zero-shot-object-detection";
 
 export interface BaseArgs {
 	/**
@@ -73,7 +94,7 @@ export interface BaseArgs {
 	 */
 	accessToken?: string;
 	/**
-	 * The model to use. Can be a full URL for HF inference endpoints.
+	 * The model to use. Can be a full URL for a dedicated inference endpoint.
 	 *
 	 * If not specified, will call huggingface.co/api/tasks to get the default model for the task.
 	 */
@@ -174,7 +195,7 @@ export type TextToSpeechOutput = Blob;
  */
 export function textToSpeech(args: TextToSpeechArgs, options?: Options): Promise<TextToSpeechOutput>;
 /**
- * Primitive to make custom calls to the inference API
+ * Primitive to make custom calls to Inference Endpoints
  */
 export function request<T>(
 	args: RequestArgs,
@@ -479,65 +500,6 @@ export function visualQuestionAnswering(
 	args: VisualQuestionAnsweringArgs,
 	options?: Options
 ): Promise<VisualQuestionAnsweringOutput>;
-export type ConversationalArgs = BaseArgs & {
-	inputs: {
-		/**
-		 * A list of strings corresponding to the earlier replies from the model.
-		 */
-		generated_responses?: string[];
-		/**
-		 * A list of strings corresponding to the earlier replies from the user. Should be of the same length of generated_responses.
-		 */
-		past_user_inputs?: string[];
-		/**
-		 * The last input from the user in the conversation.
-		 */
-		text: string;
-	};
-	parameters?: {
-		/**
-		 * (Default: None). Integer to define the maximum length in tokens of the output summary.
-		 */
-		max_length?: number;
-		/**
-		 * (Default: None). Float (0-120.0). The amount of time in seconds that the query should take maximum. Network can cause some overhead so it will be a soft limit.
-		 */
-		max_time?: number;
-		/**
-		 * (Default: None). Integer to define the minimum length in tokens of the output summary.
-		 */
-		min_length?: number;
-		/**
-		 * (Default: None). Float (0.0-100.0). The more a token is used within generation the more it is penalized to not be picked in successive generation passes.
-		 */
-		repetition_penalty?: number;
-		/**
-		 * (Default: 1.0). Float (0.0-100.0). The temperature of the sampling operation. 1 means regular sampling, 0 means always take the highest score, 100.0 is getting closer to uniform probability.
-		 */
-		temperature?: number;
-		/**
-		 * (Default: None). Integer to define the top tokens considered within the sample operation to create new text.
-		 */
-		top_k?: number;
-		/**
-		 * (Default: None). Float to define the tokens that are within the sample operation of text generation. Add tokens in the sample for more probable to least probable until the sum of the probabilities is greater than top_p.
-		 */
-		top_p?: number;
-	};
-};
-export interface ConversationalOutput {
-	conversation: {
-		generated_responses: string[];
-		past_user_inputs: string[];
-	};
-	generated_text: string;
-	warnings: string[];
-}
-/**
- * This task corresponds to any chatbot like structure. Models tend to have shorter max_length, so please check with caution when using a given model if you need long range dependency or not. Recommended model: microsoft/DialoGPT-large.
- *
- */
-export function conversational(args: ConversationalArgs, options?: Options): Promise<ConversationalOutput>;
 export type FeatureExtractionArgs = BaseArgs & {
 	/**
 	 *  The inputs is a string or a list of strings to get the features from.
@@ -740,64 +702,13 @@ export function textClassification(
 	args: TextClassificationArgs,
 	options?: Options
 ): Promise<TextClassificationOutput>;
-export type TextGenerationArgs = BaseArgs & {
-	/**
-	 * A string to be generated from
-	 */
-	inputs: string;
-	parameters?: {
-		/**
-		 * (Optional: True). Bool. Whether or not to use sampling, use greedy decoding otherwise.
-		 */
-		do_sample?: boolean;
-		/**
-		 * (Default: None). Int (0-250). The amount of new tokens to be generated, this does not include the input length it is a estimate of the size of generated text you want. Each new tokens slows down the request, so look for balance between response times and length of text generated.
-		 */
-		max_new_tokens?: number;
-		/**
-		 * (Default: None). Float (0-120.0). The amount of time in seconds that the query should take maximum. Network can cause some overhead so it will be a soft limit. Use that in combination with max_new_tokens for best results.
-		 */
-		max_time?: number;
-		/**
-		 * (Default: 1). Integer. The number of proposition you want to be returned.
-		 */
-		num_return_sequences?: number;
-		/**
-		 * (Default: None). Float (0.0-100.0). The more a token is used within generation the more it is penalized to not be picked in successive generation passes.
-		 */
-		repetition_penalty?: number;
-		/**
-		 * (Default: True). Bool. If set to False, the return results will not contain the original query making it easier for prompting.
-		 */
-		return_full_text?: boolean;
-		/**
-		 * (Default: 1.0). Float (0.0-100.0). The temperature of the sampling operation. 1 means regular sampling, 0 means always take the highest score, 100.0 is getting closer to uniform probability.
-		 */
-		temperature?: number;
-		/**
-		 * (Default: None). Integer to define the top tokens considered within the sample operation to create new text.
-		 */
-		top_k?: number;
-		/**
-		 * (Default: None). Float to define the tokens that are within the sample operation of text generation. Add tokens in the sample for more probable to least probable until the sum of the probabilities is greater than top_p.
-		 */
-		top_p?: number;
-		/**
-		 * (Default: None). Integer. The maximum number of tokens from the input.
-		 */
-		truncate?: number;
-	};
-};
-export interface TextGenerationOutput {
-	/**
-	 * The continuated string
-	 */
-	generated_text: string;
-}
 /**
  * Use to continue text from a prompt. This is a very generic task. Recommended model: gpt2 (it’s a simple model, but fun to play with).
  */
-export function textGeneration(args: TextGenerationArgs, options?: Options): Promise<TextGenerationOutput>;
+export function textGeneration(
+	args: BaseArgs & TextGenerationInput,
+	options?: Options
+): Promise<TextGenerationOutput>;
 export type TextGenerationStreamFinishReason =
 	/** number of generated tokens == `max_new_tokens` */
 	| "length"
@@ -875,7 +786,7 @@ export interface TextGenerationStreamOutput {
  * Use to continue text from a prompt. Same as `textGeneration` but returns generator that can be read one token at a time
  */
 export function textGenerationStream(
-	args: TextGenerationArgs,
+	args: BaseArgs & TextGenerationInput,
 	options?: Options
 ): AsyncGenerator<TextGenerationStreamOutput>;
 export type TokenClassificationArgs = BaseArgs & {
@@ -934,9 +845,10 @@ export type TranslationArgs = BaseArgs & {
 	/**
 	 * A string to be translated
 	 */
-	inputs: string;
+	inputs: string | string[];
 };
-export interface TranslationOutput {
+export type TranslationOutput = TranslationOutputValue | TranslationOutputValue[];
+export interface TranslationOutputValue {
 	/**
 	 * The string after translation
 	 */
@@ -1050,7 +962,7 @@ export class HfInference {
 	 */
 	textToSpeech(args: Omit<TextToSpeechArgs, 'accessToken'>, options?: Options): Promise<TextToSpeechOutput>;
 	/**
-	 * Primitive to make custom calls to the inference API
+	 * Primitive to make custom calls to Inference Endpoints
 	 */
 	request<T>(
 		args: Omit<RequestArgs, 'accessToken'>,
@@ -1130,11 +1042,6 @@ export class HfInference {
 		args: Omit<VisualQuestionAnsweringArgs, 'accessToken'>,
 		options?: Options
 	): Promise<VisualQuestionAnsweringOutput>;
-	/**
-	 * This task corresponds to any chatbot like structure. Models tend to have shorter max_length, so please check with caution when using a given model if you need long range dependency or not. Recommended model: microsoft/DialoGPT-large.
-	 *
-	 */
-	conversational(args: Omit<ConversationalArgs, 'accessToken'>, options?: Options): Promise<ConversationalOutput>;
 	/**
 	 * This task reads some text and outputs raw float values, that are usually consumed as part of a semantic database/semantic search.
 	 */
@@ -1181,12 +1088,15 @@ export class HfInference {
 	/**
 	 * Use to continue text from a prompt. This is a very generic task. Recommended model: gpt2 (it’s a simple model, but fun to play with).
 	 */
-	textGeneration(args: Omit<TextGenerationArgs, 'accessToken'>, options?: Options): Promise<TextGenerationOutput>;
+	textGeneration(
+		args: Omit<BaseArgs, 'accessToken'> & TextGenerationInput,
+		options?: Options
+	): Promise<TextGenerationOutput>;
 	/**
 	 * Use to continue text from a prompt. Same as `textGeneration` but returns generator that can be read one token at a time
 	 */
 	textGenerationStream(
-		args: Omit<TextGenerationArgs, 'accessToken'>,
+		args: Omit<BaseArgs, 'accessToken'> & TextGenerationInput,
 		options?: Options
 	): AsyncGenerator<TextGenerationStreamOutput>;
 	/**
@@ -1255,7 +1165,7 @@ export class HfInferenceEndpoint {
 	 */
 	textToSpeech(args: Omit<TextToSpeechArgs, 'accessToken' | 'model'>, options?: Options): Promise<TextToSpeechOutput>;
 	/**
-	 * Primitive to make custom calls to the inference API
+	 * Primitive to make custom calls to Inference Endpoints
 	 */
 	request<T>(
 		args: Omit<RequestArgs, 'accessToken' | 'model'>,
@@ -1335,11 +1245,6 @@ export class HfInferenceEndpoint {
 		args: Omit<VisualQuestionAnsweringArgs, 'accessToken' | 'model'>,
 		options?: Options
 	): Promise<VisualQuestionAnsweringOutput>;
-	/**
-	 * This task corresponds to any chatbot like structure. Models tend to have shorter max_length, so please check with caution when using a given model if you need long range dependency or not. Recommended model: microsoft/DialoGPT-large.
-	 *
-	 */
-	conversational(args: Omit<ConversationalArgs, 'accessToken' | 'model'>, options?: Options): Promise<ConversationalOutput>;
 	/**
 	 * This task reads some text and outputs raw float values, that are usually consumed as part of a semantic database/semantic search.
 	 */
@@ -1386,12 +1291,15 @@ export class HfInferenceEndpoint {
 	/**
 	 * Use to continue text from a prompt. This is a very generic task. Recommended model: gpt2 (it’s a simple model, but fun to play with).
 	 */
-	textGeneration(args: Omit<TextGenerationArgs, 'accessToken' | 'model'>, options?: Options): Promise<TextGenerationOutput>;
+	textGeneration(
+		args: Omit<BaseArgs, 'accessToken' | 'model'> & TextGenerationInput,
+		options?: Options
+	): Promise<TextGenerationOutput>;
 	/**
 	 * Use to continue text from a prompt. Same as `textGeneration` but returns generator that can be read one token at a time
 	 */
 	textGenerationStream(
-		args: Omit<TextGenerationArgs, 'accessToken' | 'model'>,
+		args: Omit<BaseArgs, 'accessToken' | 'model'> & TextGenerationInput,
 		options?: Options
 	): AsyncGenerator<TextGenerationStreamOutput>;
 	/**

package/dist/index.js

@@ -1,61 +1,9 @@
 /// <reference path="./index.d.ts" />
-"use strict";
 var __defProp = Object.defineProperty;
-var __getOwnPropDesc = Object.getOwnPropertyDescriptor;
-var __getOwnPropNames = Object.getOwnPropertyNames;
-var __hasOwnProp = Object.prototype.hasOwnProperty;
 var __export = (target, all) => {
   for (var name in all)
     __defProp(target, name, { get: all[name], enumerable: true });
 };
-var __copyProps = (to, from, except, desc) => {
-  if (from && typeof from === "object" || typeof from === "function") {
-    for (let key of __getOwnPropNames(from))
-      if (!__hasOwnProp.call(to, key) && key !== except)
-        __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable });
-  }
-  return to;
-};
-var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod);
-
-// src/index.ts
-var src_exports = {};
-__export(src_exports, {
-  HfInference: () => HfInference,
-  HfInferenceEndpoint: () => HfInferenceEndpoint,
-  InferenceOutputError: () => InferenceOutputError,
-  audioClassification: () => audioClassification,
-  audioToAudio: () => audioToAudio,
-  automaticSpeechRecognition: () => automaticSpeechRecognition,
-  conversational: () => conversational,
-  documentQuestionAnswering: () => documentQuestionAnswering,
-  featureExtraction: () => featureExtraction,
-  fillMask: () => fillMask,
-  imageClassification: () => imageClassification,
-  imageSegmentation: () => imageSegmentation,
-  imageToImage: () => imageToImage,
-  imageToText: () => imageToText,
-  objectDetection: () => objectDetection,
-  questionAnswering: () => questionAnswering,
-  request: () => request,
-  sentenceSimilarity: () => sentenceSimilarity,
-  streamingRequest: () => streamingRequest,
-  summarization: () => summarization,
-  tableQuestionAnswering: () => tableQuestionAnswering,
-  tabularClassification: () => tabularClassification,
-  tabularRegression: () => tabularRegression,
-  textClassification: () => textClassification,
-  textGeneration: () => textGeneration,
-  textGenerationStream: () => textGenerationStream,
-  textToImage: () => textToImage,
-  textToSpeech: () => textToSpeech,
-  tokenClassification: () => tokenClassification,
-  translation: () => translation,
-  visualQuestionAnswering: () => visualQuestionAnswering,
-  zeroShotClassification: () => zeroShotClassification,
-  zeroShotImageClassification: () => zeroShotImageClassification
-});
-module.exports = __toCommonJS(src_exports);
 
 // src/tasks/index.ts
 var tasks_exports = {};
@@ -63,7 +11,6 @@ __export(tasks_exports, {
   audioClassification: () => audioClassification,
   audioToAudio: () => audioToAudio,
   automaticSpeechRecognition: () => automaticSpeechRecognition,
-  conversational: () => conversational,
   documentQuestionAnswering: () => documentQuestionAnswering,
   featureExtraction: () => featureExtraction,
   fillMask: () => fillMask,
@@ -192,7 +139,8 @@ async function makeRequestOptions(args, options) {
       ...otherArgs,
       options: options && otherOptions
     }),
-    credentials
+    credentials,
+    signal: options?.signal
   };
   return { url, info };
 }
@@ -326,7 +274,7 @@ async function* streamingRequest(args, options) {
   const { url, info } = await makeRequestOptions({ ...args, stream: true }, options);
   const response = await (options?.fetch ?? fetch)(url, info);
   if (options?.retry_on_error !== false && response.status === 503 && !options?.wait_for_model) {
-    return streamingRequest(args, {
+    return yield* streamingRequest(args, {
       ...options,
       wait_for_model: true
     });
@@ -584,18 +532,6 @@ async function zeroShotImageClassification(args, options) {
   return res;
 }
 
-// src/tasks/nlp/conversational.ts
-async function conversational(args, options) {
-  const res = await request(args, { ...options, taskHint: "conversational" });
-  const isValidOutput = Array.isArray(res.conversation.generated_responses) && res.conversation.generated_responses.every((x) => typeof x === "string") && Array.isArray(res.conversation.past_user_inputs) && res.conversation.past_user_inputs.every((x) => typeof x === "string") && typeof res.generated_text === "string" && (typeof res.warnings === "undefined" || Array.isArray(res.warnings) && res.warnings.every((x) => typeof x === "string"));
-  if (!isValidOutput) {
-    throw new InferenceOutputError(
-      "Expected {conversation: {generated_responses: string[], past_user_inputs: string[]}, generated_text: string, warnings: string[]}"
-    );
-  }
-  return res;
-}
-
 // src/tasks/nlp/featureExtraction.ts
 async function featureExtraction(args, options) {
   const defaultTask = args.model ? await getDefaultTask(args.model, args.accessToken, options) : void 0;
@@ -765,7 +701,7 @@ async function translation(args, options) {
   if (!isValidOutput) {
     throw new InferenceOutputError("Expected type Array<{translation_text: string}>");
   }
-  return res?.[0];
+  return res?.length === 1 ? res?.[0] : res;
 }
 
 // src/tasks/nlp/zeroShotClassification.ts
@@ -902,15 +838,13 @@ var HfInferenceEndpoint = class {
     }
   }
 };
-// Annotate the CommonJS export names for ESM import in node:
-0 && (module.exports = {
+export {
   HfInference,
   HfInferenceEndpoint,
   InferenceOutputError,
   audioClassification,
   audioToAudio,
   automaticSpeechRecognition,
-  conversational,
   documentQuestionAnswering,
   featureExtraction,
   fillMask,
@@ -937,4 +871,4 @@ var HfInferenceEndpoint = class {
   visualQuestionAnswering,
   zeroShotClassification,
   zeroShotImageClassification
-});
+};

package/package.json

@@ -1,10 +1,10 @@
 {
   "name": "@huggingface/inference",
-  "version": "2.6.3",
-  "packageManager": "pnpm@8.3.1",
+  "version": "2.6.5",
+  "packageManager": "pnpm@8.10.5",
   "license": "MIT",
   "author": "Tim Mikeladze <tim.mikeladze@gmail.com>",
-  "description": "Typescript wrapper for the Hugging Face Inference API",
+  "description": "Typescript wrapper for the Hugging Face Inference Endpoints & Inference API",
   "repository": {
     "type": "git",
     "url": "https://github.com/huggingface/huggingface.js.git"
@@ -30,30 +30,28 @@
   ],
   "source": "src/index.ts",
   "types": "./dist/index.d.ts",
-  "main": "./dist/index.js",
-  "module": "./dist/index.mjs",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
   "exports": {
     "types": "./dist/index.d.ts",
-    "require": "./dist/index.js",
-    "import": "./dist/index.mjs"
+    "require": "./dist/index.cjs",
+    "import": "./dist/index.js"
   },
+  "type": "module",
   "devDependencies": {
     "@types/node": "18.13.0",
-    "ts-node": "^10.9.1",
-    "typescript": "^5.0.4",
-    "vite": "^4.1.4",
-    "vitest": "^0.29.8"
+    "@huggingface/tasks": "^0.6.0"
   },
   "resolutions": {},
   "scripts": {
     "build": "tsup src/index.ts --format cjs,esm --clean && pnpm run dts",
-    "dts": "ts-node scripts/generate-dts.ts",
+    "dts": "tsx scripts/generate-dts.ts",
     "lint": "eslint --quiet --fix --ext .cjs,.ts .",
     "lint:check": "eslint --ext .cjs,.ts .",
     "format": "prettier --write .",
     "format:check": "prettier --check .",
-    "test": "vitest run --config vitest.config.ts",
-    "test:browser": "vitest run --browser.name=chrome --browser.headless --config vitest.config.ts",
-    "type-check": "tsc"
+    "test": "vitest run --config vitest.config.mts",
+    "test:browser": "vitest run --browser.name=chrome --browser.headless --config vitest.config.mts",
+    "check": "tsc"
   }
 }

package/src/lib/getDefaultTask.ts

@@ -2,7 +2,7 @@ import { isUrl } from "./isUrl";
 
 /**
  * We want to make calls to the huggingface hub the least possible, eg if
- * someone is calling the inference API 1000 times per second, we don't want
+ * someone is calling Inference Endpoints 1000 times per second, we don't want
  * to make 1000 calls to the hub to get the task name.
  */
 const taskCache = new Map<string, { task: string; date: Date }>();

package/src/lib/makeRequestOptions.ts

@@ -106,6 +106,7 @@ export async function makeRequestOptions(
 					options: options && otherOptions,
 			  }),
 		credentials,
+		signal: options?.signal,
 	};
 
 	return { url, info };

package/src/tasks/custom/request.ts

@@ -2,7 +2,7 @@ import type { InferenceTask, Options, RequestArgs } from "../../types";
 import { makeRequestOptions } from "../../lib/makeRequestOptions";
 
 /**
- * Primitive to make custom calls to the inference API
+ * Primitive to make custom calls to Inference Endpoints
  */
 export async function request<T>(
 	args: RequestArgs,

package/src/tasks/custom/streamingRequest.ts

@@ -19,7 +19,7 @@ export async function* streamingRequest<T>(
 	const response = await (options?.fetch ?? fetch)(url, info);
 
 	if (options?.retry_on_error !== false && response.status === 503 && !options?.wait_for_model) {
-		return streamingRequest(args, {
+		return yield* streamingRequest(args, {
 			...options,
 			wait_for_model: true,
 		});

package/src/tasks/index.ts

@@ -18,7 +18,6 @@ export * from "./cv/imageToImage";
 export * from "./cv/zeroShotImageClassification";
 
 // Natural Language Processing tasks
-export * from "./nlp/conversational";
 export * from "./nlp/featureExtraction";
 export * from "./nlp/fillMask";
 export * from "./nlp/questionAnswering";

package/src/tasks/nlp/textGeneration.ts

@@ -1,67 +1,15 @@
+import type { TextGenerationInput, TextGenerationOutput } from "@huggingface/tasks/src/tasks/text-generation/inference";
 import { InferenceOutputError } from "../../lib/InferenceOutputError";
 import type { BaseArgs, Options } from "../../types";
 import { request } from "../custom/request";
 
-export type TextGenerationArgs = BaseArgs & {
-	/**
-	 * A string to be generated from
-	 */
-	inputs: string;
-	parameters?: {
-		/**
-		 * (Optional: True). Bool. Whether or not to use sampling, use greedy decoding otherwise.
-		 */
-		do_sample?: boolean;
-		/**
-		 * (Default: None). Int (0-250). The amount of new tokens to be generated, this does not include the input length it is a estimate of the size of generated text you want. Each new tokens slows down the request, so look for balance between response times and length of text generated.
-		 */
-		max_new_tokens?: number;
-		/**
-		 * (Default: None). Float (0-120.0). The amount of time in seconds that the query should take maximum. Network can cause some overhead so it will be a soft limit. Use that in combination with max_new_tokens for best results.
-		 */
-		max_time?: number;
-		/**
-		 * (Default: 1). Integer. The number of proposition you want to be returned.
-		 */
-		num_return_sequences?: number;
-		/**
-		 * (Default: None). Float (0.0-100.0). The more a token is used within generation the more it is penalized to not be picked in successive generation passes.
-		 */
-		repetition_penalty?: number;
-		/**
-		 * (Default: True). Bool. If set to False, the return results will not contain the original query making it easier for prompting.
-		 */
-		return_full_text?: boolean;
-		/**
-		 * (Default: 1.0). Float (0.0-100.0). The temperature of the sampling operation. 1 means regular sampling, 0 means always take the highest score, 100.0 is getting closer to uniform probability.
-		 */
-		temperature?: number;
-		/**
-		 * (Default: None). Integer to define the top tokens considered within the sample operation to create new text.
-		 */
-		top_k?: number;
-		/**
-		 * (Default: None). Float to define the tokens that are within the sample operation of text generation. Add tokens in the sample for more probable to least probable until the sum of the probabilities is greater than top_p.
-		 */
-		top_p?: number;
-		/**
-		 * (Default: None). Integer. The maximum number of tokens from the input.
-		 */
-		truncate?: number;
-	};
-};
-
-export interface TextGenerationOutput {
-	/**
-	 * The continuated string
-	 */
-	generated_text: string;
-}
-
 /**
  * Use to continue text from a prompt. This is a very generic task. Recommended model: gpt2 (it’s a simple model, but fun to play with).
  */
-export async function textGeneration(args: TextGenerationArgs, options?: Options): Promise<TextGenerationOutput> {
+export async function textGeneration(
+	args: BaseArgs & TextGenerationInput,
+	options?: Options
+): Promise<TextGenerationOutput> {
 	const res = await request<TextGenerationOutput[]>(args, {
 		...options,
 		taskHint: "text-generation",

package/src/tasks/nlp/textGenerationStream.ts

@@ -1,6 +1,7 @@
-import type { Options } from "../../types";
+import type { BaseArgs, Options } from "../../types";
 import { streamingRequest } from "../custom/streamingRequest";
-import type { TextGenerationArgs } from "./textGeneration";
+
+import type { TextGenerationInput } from "@huggingface/tasks/src/tasks/text-generation/inference";
 
 export interface TextGenerationStreamToken {
 	/** Token ID from the model tokenizer */
@@ -85,7 +86,7 @@ export interface TextGenerationStreamOutput {
  * Use to continue text from a prompt. Same as `textGeneration` but returns generator that can be read one token at a time
  */
 export async function* textGenerationStream(
-	args: TextGenerationArgs,
+	args: BaseArgs & TextGenerationInput,
 	options?: Options
 ): AsyncGenerator<TextGenerationStreamOutput> {
 	yield* streamingRequest<TextGenerationStreamOutput>(args, {