diff --git a/packages/tasks/src/tasks/text-generation/inference.ts b/packages/tasks/src/tasks/text-generation/inference.ts index b3c002c6e5..825d7d368f 100644 --- a/packages/tasks/src/tasks/text-generation/inference.ts +++ b/packages/tasks/src/tasks/text-generation/inference.ts @@ -93,9 +93,17 @@ 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 */ @@ -103,10 +111,36 @@ export interface TextGenerationOutput { /** * 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 index dc68a3aada..4f2e53c9a4 100644 --- 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",