feat(model): add text generation task (#136)

Because

- support text generation task

This commit

- add text generation task input/output definitions

openapiv2/openapiv2.swagger.yaml

@@ -3312,6 +3312,7 @@ definitions:
       - TASK_INSTANCE_SEGMENTATION
       - TASK_SEMANTIC_SEGMENTATION
       - TASK_TEXT_TO_IMAGE
+      - TASK_TEXT_GENERATION
     default: TASK_UNSPECIFIED
     description: |-
       - TASK_UNSPECIFIED: Task: UNSPECIFIED
@@ -3322,6 +3323,7 @@ definitions:
        - TASK_INSTANCE_SEGMENTATION: Task: INSTANCE SEGMENTATION
        - TASK_SEMANTIC_SEGMENTATION: Task: SEMANTIC SEGMENTATION
        - TASK_TEXT_TO_IMAGE: Task: TEXT TO IMAGE
+       - TASK_TEXT_GENERATION: Task: TEXT GENERATION
     title: Task enumerates the task type of a model instance
   PipelineMode:
     type: string
@@ -5539,6 +5541,9 @@ definitions:
       text_to_image:
         $ref: "#/definitions/v1alphaTextToImageInput"
         title: The text to image input
+      text_generation:
+        $ref: "#/definitions/v1alphaTextGenerationInput"
+        title: The text generation input
       unspecified:
         $ref: "#/definitions/v1alphaUnspecifiedInput"
         title: The unspecified task input
@@ -5610,6 +5615,49 @@ definitions:
         title: List of generated images
         readOnly: true
     title: TextToImageOutput represents the output of text to image task
+  v1alphaTextGenerationInput:
+    type: object
+    properties:
+      prompt:
+        type: string
+        title: The prompt text
+        readOnly: true
+      output_len:
+        type: string
+        format: int64
+        title: The maximum number of tokens for model to generate
+        readOnly: true
+      bad_words_list:
+        type: array
+        items:
+          type: string
+        title: The words to avoid being generated by the model
+        readOnly: true
+      stop_words_list:
+        type: array
+        items:
+          type: string
+        title: The trigger words to stop generation
+        readOnly: true
+      topk:
+        type: string
+        format: int64
+        title: Top k for sampling
+        readOnly: true
+      seed:
+        type: string
+        format: int64
+        title: The seed
+        readOnly: true
+    title: TextToImageInput represents the input of text generation task
+  v1alphaTextGenerationOutput:
+    type: object
+    properties:
+      text:
+        type: string
+        title: The text generated by model
+        readOnly: true
+    title: TextToImageOutput represents the output of text generation task
   v1alphaTriggerModelInstanceBinaryFileUploadResponse:
     type: object
     properties:
@@ -5998,6 +6046,10 @@ definitions:
         $ref: "#/definitions/v1alphaTextToImageOutput"
         title: The text to image output
         readOnly: true
+      text_generation:
+        $ref: "#/definitions/v1alphaTextGenerationOutput"
+        title: The text generation output
+        readOnly: true
       unspecified:
         $ref: "#/definitions/v1alphaUnspecifiedOutput"
         title: The unspecified task output
@@ -6067,6 +6119,10 @@ definitions:
         $ref: "#/definitions/v1alphaTextToImageOutput"
         title: The text to image output
         readOnly: true
+      text_generation:
+        $ref: "#/definitions/v1alphaTextGenerationOutput"
+        title: The text generation output
+        readOnly: true
       unspecified:
         $ref: "#/definitions/v1alphaUnspecifiedOutput"
         title: The unspecified task output

vdp/model/v1alpha/model.proto

@@ -21,6 +21,7 @@ import "vdp/model/v1alpha/task_ocr.proto";
 import "vdp/model/v1alpha/task_instance_segmentation.proto";
 import "vdp/model/v1alpha/task_semantic_segmentation.proto";
 import "vdp/model/v1alpha/task_text_to_image.proto";
+import "vdp/model/v1alpha/task_text_generation.proto";
 import "vdp/model/v1alpha/task_unspecified.proto";
 
 // Model represents a model
@@ -109,6 +110,8 @@ message ModelInstance {
     TASK_SEMANTIC_SEGMENTATION = 6;
     // Task: TEXT TO IMAGE
     TASK_TEXT_TO_IMAGE = 7;
+    // Task: TEXT Generation
+    TASK_TEXT_GENERATION = 8;
   }
 
   // State enumerates a model instance state
@@ -514,8 +517,10 @@ message TaskInput {
     SemanticSegmentationInput semantic_segmentation = 6;
     // The text to image input
     TextToImageInput text_to_image = 7;
+    // The text generation input
+    TextGenerationInput text_generation = 8;
     // The unspecified task input
-    UnspecifiedInput unspecified = 8;
+    UnspecifiedInput unspecified = 9;
   }
 }
 
@@ -541,8 +546,11 @@ message TaskOutput {
     // The text to image output
     TextToImageOutput text_to_image = 7
         [ (google.api.field_behavior) = OUTPUT_ONLY ];
+    // The text generation output
+    TextGenerationOutput text_generation = 8
+        [ (google.api.field_behavior) = OUTPUT_ONLY ];
     // The unspecified task output
-    UnspecifiedOutput unspecified = 8
+    UnspecifiedOutput unspecified = 9
         [ (google.api.field_behavior) = OUTPUT_ONLY ];
   }
 }

vdp/model/v1alpha/task_text_generation.proto

@@ -0,0 +1,28 @@
+syntax = "proto3";
+
+package vdp.model.v1alpha;
+
+// Google api
+import "google/api/field_behavior.proto";
+
+// TextGenerationInput represents the input of text generation task
+message TextGenerationInput {
+  // The prompt text
+  string prompt = 1 [ (google.api.field_behavior) = REQUIRED ];
+  // The maximum number of tokens for model to generate
+  optional int64 output_len = 2 [ (google.api.field_behavior) = OPTIONAL ];
+  // The words not to avoid being generated by the model
+  optional string bad_words_list = 3 [ (google.api.field_behavior) = OPTIONAL ];
+  // The trigger words to stop generation
+  optional string stop_words_list = 4 [ (google.api.field_behavior) = OPTIONAL ];
+  // Top k for sampling
+  optional int64 topk = 5 [ (google.api.field_behavior) = OPTIONAL ];
+  // The seed
+  optional int64 seed = 6 [ (google.api.field_behavior) = OPTIONAL ];
+}
+
+// TextGenerationOutput represents the output of text generation task
+message TextGenerationOutput {
+  // The text generated by the model
+  string text = 1 [ (google.api.field_behavior) = OUTPUT_ONLY ];
+}