Add typings for the TGI / OpenAI chat API

huggingface · Feb 9, 2024 · ad00ad5 · ad00ad5
1 parent 8dda47f
commit ad00ad5
Show file tree

Hide file tree

Showing 2 changed files with 120 additions and 60 deletions.
diff --git a/packages/tasks/src/tasks/text-generation/inference.ts b/packages/tasks/src/tasks/text-generation/inference.ts
@@ -93,20 +93,54 @@ export interface TextGenerationParameters {
 }
 
 /**
- * Outputs for Text Generation inference
+ * Outputs for Text Generation inference when using the text generation API
+ *
+ * Outputs for Text Generation inference when using the messages API. Compatible with Open
+ * AI chat API.
  */
 export interface TextGenerationOutput {
+	choices?: Choice[];
+	/**
+	 * The creation date as a UNIX timestamp.
+	 */
+	created?: number;
 	/**
 	 * When enabled, details about the generation
 	 */
 	details?: TextGenerationOutputDetails;
 	/**
 	 * The generated text
 	 */
-	generated_text: string;
+	generated_text?: string;
+	id?: string;
+	/**
+	 * The name of the model used for this generation.
+	 */
+	model?: string;
+	object?: string;
+	system_fingerprint?: string;
+	[property: string]: unknown;
+}
+
+export interface Choice {
+	delta: string;
+	finish_reason?: FinishReason;
+	index: number;
+	logprobs?: number;
 	[property: string]: unknown;
 }
 
+/**
+ * The reason why the generation was stopped.
+ *
+ * The generated sequence reached the maximum allowed length
+ *
+ * The model generated an end-of-sentence (EOS) token
+ *
+ * One of the sequence in stop_sequences was generated
+ */
+export type FinishReason = "length" | "eos_token" | "stop_sequence";
+
 /**
  * When enabled, details about the generation
  */
@@ -115,9 +149,6 @@ export interface TextGenerationOutputDetails {
 	 * Details about additional sequences when best_of is provided
 	 */
 	best_of_sequences?: TextGenerationSequenceDetails[];
-	/**
-	 * The reason why the generation was stopped.
-	 */
 	finish_reason: FinishReason;
 	/**
 	 * The number of generated tokens
@@ -136,9 +167,6 @@ export interface TextGenerationOutputDetails {
 }
 
 export interface TextGenerationSequenceDetails {
-	/**
-	 * The reason why the generation was stopped.
-	 */
 	finish_reason: FinishReason;
 	/**
 	 * The generated text
@@ -160,15 +188,6 @@ export interface TextGenerationSequenceDetails {
 	[property: string]: unknown;
 }
 
-/**
- * The generated sequence reached the maximum allowed length
- *
- * The model generated an end-of-sentence (EOS) token
- *
- * One of the sequence in stop_sequences was generated
- */
-export type FinishReason = "length" | "eos_token" | "stop_sequence";
-
 export interface PrefillToken {
 	id: number;
 	logprob: number;

diff --git a/packages/tasks/src/tasks/text-generation/spec/output.json b/packages/tasks/src/tasks/text-generation/spec/output.json
@@ -1,49 +1,96 @@
 {
 	"$id": "/inference/schemas/text-generation/output.json",
 	"$schema": "http://json-schema.org/draft-06/schema#",
-	"description": "Outputs for Text Generation inference",
 	"title": "TextGenerationOutput",
-	"type": "object",
-	"properties": {
-		"generated_text": {
-			"type": "string",
-			"description": "The generated text"
-		},
-		"details": {
-			"description": "When enabled, details about the generation",
-			"title": "TextGenerationOutputDetails",
-			"allOf": [
-				{ "$ref": "#/$defs/SequenceDetails" },
-				{
-					"type": "object",
-					"properties": {
-						"best_of_sequences": {
-							"type": "array",
-							"description": "Details about additional sequences when best_of is provided",
-							"items": {
-								"allOf": [
-									{ "$ref": "#/$defs/SequenceDetails" },
-									{
-										"type": "object",
-										"properties": {
-											"generated_text": {
-												"type": "integer",
-												"description": "The generated text"
+	"anyOf": [{ "$ref": "#/$defs/TextGenerationOutputText" }, { "$ref": "#/$defs/TextGenerationOutputMessages" }],
+	"$defs": {
+		"TextGenerationOutputText": {
+			"type": "object",
+			"title": "TextGenerationOutputText",
+			"description": "Outputs for Text Generation inference when using the text generation API",
+			"properties": {
+				"generated_text": {
+					"type": "string",
+					"description": "The generated text"
+				},
+				"details": {
+					"description": "When enabled, details about the generation",
+					"title": "TextGenerationOutputDetails",
+					"allOf": [
+						{ "$ref": "#/$defs/SequenceDetails" },
+						{
+							"type": "object",
+							"properties": {
+								"best_of_sequences": {
+									"type": "array",
+									"description": "Details about additional sequences when best_of is provided",
+									"items": {
+										"allOf": [
+											{ "$ref": "#/$defs/SequenceDetails" },
+											{
+												"type": "object",
+												"properties": {
+													"generated_text": {
+														"type": "integer",
+														"description": "The generated text"
+													}
+												},
+												"required": ["generated_text"]
 											}
-										},
-										"required": ["generated_text"]
+										]
 									}
-								]
+								}
 							}
 						}
-					}
+					]
 				}
+			},
+			"required": ["generated_text"]
+		},
+		"TextGenerationOutputMessages": {
+			"type": "object",
+			"title": "TextGenerationOutputMessages",
+			"description": "Outputs for Text Generation inference when using the messages API. Compatible with Open AI chat API.",
+			"properties": {
+				"choices": {
+					"type": "array",
+					"items": {
+						"type": "object",
+						"properties": {
+							"delta": {
+								"type": "string",
+								"properties": {
+									"content": { "type": "string" },
+									"role": { "type": "string" }
+								}
+							},
+							"finish_reason": { "$ref": "#/$defs/FinishReason" },
+							"index": { "type": "integer" },
+							"logprobs": { "type": "number" }
+						},
+						"required": ["delta", "index"]
+					}
+				},
+				"created": {
+					"type": "integer",
+					"description": "The creation date as a UNIX timestamp."
+				},
+				"id": { "type": "string" },
+				"model": { "type": "string", "description": "The name of the model used for this generation." },
+				"object": { "type": "string" },
+				"system_fingerprint": { "type": "string" }
+			},
+			"required": ["choices", "created", "id", "model", "object", "system_fingerprint"]
+		},
+		"FinishReason": {
+			"type": "string",
+			"description": "The reason why the generation was stopped.",
+			"oneOf": [
+				{ "const": "length", "description": "The generated sequence reached the maximum allowed length" },
+				{ "const": "eos_token", "description": "The model generated an end-of-sentence (EOS) token" },
+				{ "const": "stop_sequence", "description": "One of the sequence in stop_sequences was generated" }
 			]
-		}
-	},
-	"required": ["generated_text"],
-
-	"$defs": {
+		},
 		"Token": {
 			"type": "object",
 			"title": "Token",
@@ -70,13 +117,7 @@
 			"title": "TextGenerationSequenceDetails",
 			"properties": {
 				"finish_reason": {
-					"type": "string",
-					"description": "The reason why the generation was stopped.",
-					"oneOf": [
-						{ "const": "length", "description": "The generated sequence reached the maximum allowed length" },
-						{ "const": "eos_token", "description": "The model generated an end-of-sentence (EOS) token" },
-						{ "const": "stop_sequence", "description": "One of the sequence in stop_sequences was generated" }
-					]
+					"$ref": "#/$defs/FinishReason"
 				},
 				"generated_tokens": {
 					"type": "integer",