Repository: caikit/caikit
Pull request #237 (2023-06-14 19:13:39)

Text generation stream
## Description

This PR adds the interfaces for Text Generation and Text Generation Streaming tasks

**NOTE:** These interfaces intentionally align closely with [Hugging Face Detail Parameters](https://huggingface.co/docs/api-inference/detailed_parameters#text-generation-task) and [Text Generation Inference](https://github.com/huggingface/text-generation-inference/blob/main/docs/openapi.json)


Patch:
diff --git a/caikit/interfaces/__init__.py b/caikit/interfaces/__init__.py
index 4bd73eb3..ddffa296 100644
--- a/caikit/interfaces/__init__.py
+++ b/caikit/interfaces/__init__.py
@@ -1,3 +1,17 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 """This library defines the taxonomy of Data Model objects and Tasks for the
 entire CAIKit project. Data objects and tasks are grouped domain, making for a
 three-level hierarchy for models:
@@ -8,10 +22,5 @@ This library intentionally does NOT define any implementations, as those are
 left to the domain implementation libraries.
 """
 
-# First Party
-import import_tracker
-
-# Import each domain with lazy import errors
-with import_tracker.lazy_import_errors():
-    # Local
-    from . import common, runtime
+# NOTE: We do not proactively import sub-modules in the interfaces here as they
+#   may contain optional dependencies
diff --git a/caikit/interfaces/nlp/__init__.py b/caikit/interfaces/nlp/__init__.py
new file mode 100644
index 00000000..4d71e7d1
--- /dev/null
+++ b/caikit/interfaces/nlp/__init__.py
@@ -0,0 +1,16 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Local
+from . import data_model, tasks
diff --git a/caikit/interfaces/nlp/data_model/__init__.py b/caikit/interfaces/nlp/data_model/__init__.py
new file mode 100644
index 00000000..f334e111
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/__init__.py
@@ -0,0 +1,24 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Local
+from . import package, text_generation
+from .package import NLP_PACKAGE
+from .text_generation import (
+    FinishReason,
+    GeneratedTextResult,
+    GeneratedTextStreamResult,
+    GeneratedToken,
+    TokenStreamDetails,
+)
diff --git a/caikit/interfaces/nlp/data_model/package.py b/caikit/interfaces/nlp/data_model/package.py
new file mode 100644
index 00000000..fffb5f1f
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/package.py
@@ -0,0 +1,19 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Local
+from ....core.data_model import CAIKIT_DATA_MODEL
+
+# Shared package name constant for all NLP data model objects
+NLP_PACKAGE = f"{CAIKIT_DATA_MODEL}.nlp"
diff --git a/caikit/interfaces/nlp/data_model/text_generation.py b/caikit/interfaces/nlp/data_model/text_generation.py
new file mode 100644
index 00000000..12d97a29
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/text_generation.py
@@ -0,0 +1,72 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Data structures for text generation representations"""
+
+# Standard
+from enum import Enum
+from typing import Optional
+
+# Third Party
+import numpy as np
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+import alog
+
+# Local
+from ....core import DataObjectBase, dataobject
+from ...common.data_model import ProducerId
+from .package import NLP_PACKAGE
+
+log = alog.use_channel("DATAM")
+
+
+@dataobject(package=NLP_PACKAGE)
+class FinishReason(Enum):
+    NOT_FINISHED = 0
+    MAX_TOKENS = 1
+    EOS_TOKEN = 2
+    CANCELLED = 3
+    TIME_LIMIT = 4
+    STOP_SEQUENCE = 5
+
+
+@dataobject(package=NLP_PACKAGE)
+class GeneratedTextResult(DataObjectBase):
+    generated_text: Annotated[str, FieldNumber(1)]
+    generated_tokens: Annotated[int, FieldNumber(2)]
+    finish_reason: Annotated[FinishReason, FieldNumber(3)]
+    producer_id: Annotated[ProducerId, FieldNumber(4)]
+
+
+@dataobject(package=NLP_PACKAGE)
+class GeneratedToken(DataObjectBase):
+    text: Annotated[str, FieldNumber(1)]
+    id: Annotated[Optional[np.uint32], FieldNumber(2)]
+    logprob: Annotated[Optional[float], FieldNumber(3)]
+    special: Annotated[Optional[bool], FieldNumber(4)]
+
+
+@dataobject(package=NLP_PACKAGE)
+class TokenStreamDetails(DataObjectBase):
+    finish_reason: Annotated[FinishReason, FieldNumber(1)]
+    generated_tokens: Annotated[np.uint32, FieldNumber(2)]
+    seed: Annotated[np.uint64, FieldNumber(3)]
+
+
+@dataobject(package=NLP_PACKAGE)
+class GeneratedTextStreamResult(DataObjectBase):
+    token: Annotated[GeneratedToken, FieldNumber(1)]
+    generated_text: Annotated[Optional[str], FieldNumber(2)]
+    details: Annotated[Optional[TokenStreamDetails], FieldNumber(3)]
diff --git a/caikit/interfaces/nlp/tasks.py b/caikit/interfaces/nlp/tasks.py
new file mode 100644
index 00000000..1cd6227a
--- /dev/null
+++ b/caikit/interfaces/nlp/tasks.py
@@ -0,0 +1,43 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module holds the Task definitions for all common NLP tasks
+"""
+
+# Standard
+from typing import Iterable
+
+# Local
+from ...core import TaskBase, task
+from .data_model.text_generation import GeneratedTextResult, GeneratedTextStreamResult
+
+
+@task(
+    required_parameters={"inputs": str},
+    output_type=GeneratedTextResult,
+)
+class TextGenerationTask(TaskBase):
+    """The Text Generation Task is responsible for taking input prompting text
+    and generating additional text from that prompt.
+    """
+
+
+@task(
+    required_parameters={"inputs": str},
+    output_type=Iterable[GeneratedTextStreamResult],
+)
+class TextGenerationStreamTask(TaskBase):
+    """The Text Generation Stream Task is responsible for taking input prompting
+    text and iteratively producing output tokens in a stream.
+    """
diff --git a/pyproject.toml b/pyproject.toml
index 793f8342..68eea578 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,7 +22,6 @@ dependencies = [
     "grpcio-reflection>=1.35.0,<2.0",
     "grpcio>=1.35.0,<2.0,!=1.55.0",
     "ijson>=3.1.4,<3.3.0",
-    "import-tracker>=3.1.5,<4",
     "munch>=2.5.0,<4.0",
     "numpy>=1.20,<2",
     "prometheus_client>=0.12.0,<1.0",


Test patch:
diff --git a/tests/interfaces/__init__.py b/tests/interfaces/__init__.py
new file mode 100644
index 00000000..2068258b
--- /dev/null
+++ b/tests/interfaces/__init__.py
@@ -0,0 +1,13 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/tests/interfaces/nlp/__init__.py b/tests/interfaces/nlp/__init__.py
new file mode 100644
index 00000000..2068258b
--- /dev/null
+++ b/tests/interfaces/nlp/__init__.py
@@ -0,0 +1,13 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/tests/interfaces/nlp/test_nlp_tasks.py b/tests/interfaces/nlp/test_nlp_tasks.py
new file mode 100644
index 00000000..4ca103c5
--- /dev/null
+++ b/tests/interfaces/nlp/test_nlp_tasks.py
@@ -0,0 +1,86 @@
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for the NLP task definitions"""
+# Standard
+from typing import Type
+
+# Third Party
+import pytest
+
+# Local
+from caikit.core import ModuleBase, TaskBase, module
+from caikit.interfaces.nlp import tasks as nlp_tasks
+from tests.core.helpers import *
+
+## Helpers #####################################################################
+
+
+## Tests #######################################################################
+
+
+class InvalidType:
+    pass
+
+
+@pytest.mark.parametrize(
+    "task", (nlp_tasks.TextGenerationTask, nlp_tasks.TextGenerationStreamTask)
+)
+def test_tasks(reset_globals, task: Type[TaskBase]):
+    """Common tests for all tasks"""
+    # Only support single required param named "inputs"
+    assert set(task.get_required_parameters().keys()) == {"inputs"}
+    input_type = task.get_required_parameters()["inputs"]
+    output_type = task.get_output_type()
+
+    # Version with the right signature and nothing else
+    @module(id="foo1", name="Foo", version="0.0.0", task=task)
+    class Foo1(ModuleBase):
+        def run(self, inputs: input_type) -> output_type:
+            return output_type()
+
+    # Version with the right signature plus extra args
+    @module(id="foo2", name="Foo", version="0.0.0", task=task)
+    class Foo2(ModuleBase):
+        def run(
+            self,
+            inputs: input_type,
+            workit: bool,
+            makeit: bool,
+            doit: bool,
+        ) -> output_type:
+            return output_type()
+
+    # Version with missing required argument
+    with pytest.raises(TypeError):
+
+        @module(id="foo3", name="Foo", version="0.0.0", task=task)
+        class Foo3(ModuleBase):
+            def run(self, other_name: str) -> output_type:
+                return output_type()
+
+    # Version with bad required argument type
+    with pytest.raises(TypeError):
+
+        @module(id="foo4", name="Foo", version="0.0.0", task=task)
+        class Foo4(ModuleBase):
+            def run(self, inputs: InvalidType) -> output_type:
+                return output_type()
+
+    # Version with bad return type
+    with pytest.raises(TypeError):
+
+        @module(id="foo", name="Foo", version="0.0.0", task=task)
+        class Foo(ModuleBase):
+            def run(self, inputs: input_type) -> InvalidType:
+                return "hi there"