Repository: caikit/caikit (language: python)

Issue:
Add nlp classification interfaces to caikit.interfaces
## Description

We want to move [classification interfaces](https://github.com/caikit/caikit-nlp/blob/main/caikit_nlp/data_model/classification.py) to `caikit.interfaces` and provide a "combined" task and output to enable milestone https://github.com/caikit/caikit-nlp/milestone/5 (but can't link milestones across repos). This output should have text generation output and token classification output fields.

## Acceptance Criteria

<!-- Remove any that don't apply -->

- [ ] Unit tests cover new/changed code
- [ ] Interfaces added
- [ ] Type of [semantic version](https://semver.org/) change is identified



Discussion / hints:

Fix patch:
diff --git a/caikit/interfaces/nlp/data_model/__init__.py b/caikit/interfaces/nlp/data_model/__init__.py
index f334e1112..ff1e6d715 100644
--- a/caikit/interfaces/nlp/data_model/__init__.py
+++ b/caikit/interfaces/nlp/data_model/__init__.py
@@ -13,8 +13,19 @@
 # limitations under the License.
 
 # Local
-from . import package, text_generation
+from . import classification, package, text, text_generation
+from .classification import (
+    ClassificationResult,
+    ClassificationResults,
+    ClassificationTrainRecord,
+    ClassifiedGeneratedTextResult,
+    ClassifiedGeneratedTextStreamResult,
+    TokenClassificationResult,
+    TokenClassificationResults,
+    TokenClassificationStreamResult,
+)
 from .package import NLP_PACKAGE
+from .text import Token, TokenizationResults, TokenizationStreamResult
 from .text_generation import (
     FinishReason,
     GeneratedTextResult,
diff --git a/caikit/interfaces/nlp/data_model/classification.py b/caikit/interfaces/nlp/data_model/classification.py
new file mode 100644
index 000000000..56666e242
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/classification.py
@@ -0,0 +1,136 @@
+# 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 classification representations"""
+
+# Standard
+from typing import List, 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 .package import NLP_PACKAGE
+from .text_generation import FinishReason
+
+log = alog.use_channel("DATAM")
+
+
+@dataobject(package=NLP_PACKAGE)
+class ClassificationTrainRecord(DataObjectBase):
+    """A classification training record consisting of a single train instance."""
+
+    text: Annotated[str, FieldNumber(1)]  # Text to be classified
+    labels: Annotated[
+        List[str], FieldNumber(2)
+    ]  # Class labels to be learnt for the text
+
+
+@dataobject(package=NLP_PACKAGE)
+class ClassificationResult(DataObjectBase):
+    """A single classification prediction."""
+
+    label: Annotated[str, FieldNumber(1)]  # Predicted relevant class name
+    score: Annotated[
+        float, FieldNumber(2)
+    ]  # The confidence-like score of this prediction in [0, 1]
+
+
+@dataobject(package=NLP_PACKAGE)
+class ClassificationResults(DataObjectBase):
+    """Classification results generated from a text and consisting multiple classes."""
+
+    results: Annotated[
+        List[ClassificationResult], FieldNumber(1)
+    ]  # List of classifications for a text
+
+
+# NOTE: This is meant to align with the HuggingFace token classification task:
+# https://huggingface.co/docs/transformers/tasks/token_classification#inference
+# The field `word` does not necessarily correspond to a single "word",
+# and `entity` may not always be applicable beyond "entity" in the NER
+# (named entity recognition) sense
+@dataobject(package=NLP_PACKAGE)
+class TokenClassificationResult(DataObjectBase):
+    """A single token classification prediction."""
+
+    start: Annotated[int, FieldNumber(1)]  # Beginning offset of the token
+    end: Annotated[int, FieldNumber(2)]  # Ending offset of the token
+    word: Annotated[str, FieldNumber(3)]  # Text referenced by this token
+    entity: Annotated[
+        str, FieldNumber(4)
+    ]  # Predicted relevant class name for the token
+    entity_group: Annotated[str, FieldNumber(5)]  # Aggregate label, if applicable
+    score: Annotated[
+        float, FieldNumber(6)
+    ]  # The confidence-like score of this classification prediction in [0, 1]
+    token_count: Annotated[
+        Optional[int], FieldNumber(7)
+    ]  # Length of tokens in the text
+
+
+@dataobject(package=NLP_PACKAGE)
+class TokenClassificationResults(DataObjectBase):
+    """Token classification results generated from a text and consisting multiple classes."""
+
+    results: Annotated[List[TokenClassificationResult], FieldNumber(1)]
+
+
+@dataobject(package=NLP_PACKAGE)
+class TokenClassificationStreamResult(TokenClassificationResults):
+    """
+    Streaming token classification results that indicates up to where in stream is processed.
+    """
+
+    processed_index: Annotated[
+        int, FieldNumber(2)
+    ]  # Result index up to which text is processed
+
+
+@dataobject(package=NLP_PACKAGE)
+class ClassifiedGeneratedTextResult(DataObjectBase):
+    """Classification result on text produced by a text generation model, contains
+    information from the original text generation output as well as the result of
+    classification on the generated text.
+    """
+
+    text: Annotated[str, FieldNumber(1)]  # The generated text
+    token_classification_results: Annotated[
+        Optional[List[TokenClassificationResult]], FieldNumber(2)
+    ]  # Token classification results for this generated text
+    finish_reason: Annotated[
+        Optional[FinishReason], FieldNumber(3)
+    ]  # Reason as to why text generation stopped
+    token_count: Annotated[
+        Optional[int], FieldNumber(4)
+    ]  # Length of generated tokens sequence
+    seed: Annotated[
+        Optional[np.uint64], FieldNumber(5)
+    ]  # The random seed used for text generation
+
+
+@dataobject(package=NLP_PACKAGE)
+class ClassifiedGeneratedTextStreamResult(ClassifiedGeneratedTextResult):
+    """
+    Streaming classification on generated text result that indicates up to where in stream
+    is processed.
+    """
+
+    processed_index: Annotated[
+        int, FieldNumber(6)
+    ]  # Result index up to which text is processed
diff --git a/caikit/interfaces/nlp/data_model/text.py b/caikit/interfaces/nlp/data_model/text.py
new file mode 100644
index 000000000..87dd5f15b
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/text.py
@@ -0,0 +1,57 @@
+# 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 representations"""
+
+# Standard
+from typing import List
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+import alog
+
+# Local
+from ....core import DataObjectBase, dataobject
+from .package import NLP_PACKAGE
+
+log = alog.use_channel("DATAM")
+
+
+@dataobject(package=NLP_PACKAGE)
+class Token(DataObjectBase):
+    """Tokens here are the basic units of text. Tokens can be characters, words,
+    sub-words, or other segments of text or code, depending on the method of
+    tokenization chosen or the task being implemented.
+    """
+
+    start: Annotated[int, FieldNumber(1)]  # Beginning offset of the token
+    end: Annotated[int, FieldNumber(2)]  # Ending offset of the token
+    text: Annotated[str, FieldNumber(3)]  # Text referenced by this token
+
+
+@dataobject(package=NLP_PACKAGE)
+class TokenizationResults(DataObjectBase):
+    """Tokenization result generated from a text."""
+
+    results: Annotated[List[Token], FieldNumber(1)]
+
+
+@dataobject(package=NLP_PACKAGE)
+class TokenizationStreamResult(TokenizationResults):
+    """
+    Streaming tokenization result that indicates up to where in stream is processed.
+    """
+
+    processed_index: Annotated[
+        int, FieldNumber(2)
+    ]  # Result index up to which text is processed
diff --git a/caikit/interfaces/nlp/tasks.py b/caikit/interfaces/nlp/tasks.py
index 3321467c0..c8e16689b 100644
--- a/caikit/interfaces/nlp/tasks.py
+++ b/caikit/interfaces/nlp/tasks.py
@@ -20,6 +20,14 @@
 
 # Local
 from ...core import TaskBase, task
+from .data_model.classification import (
+    ClassificationResults,
+    ClassifiedGeneratedTextResult,
+    ClassifiedGeneratedTextStreamResult,
+    TokenClassificationResults,
+    TokenClassificationStreamResult,
+)
+from .data_model.text import TokenizationResults, TokenizationStreamResult
 from .data_model.text_generation import GeneratedTextResult, GeneratedTextStreamResult
 
 
@@ -32,3 +40,45 @@ class TextGenerationTask(TaskBase):
     """The Text Generation Task is responsible for taking input prompting text
     and generating additional text from that prompt.
     """
+
+
+@task(
+    required_parameters={"text": str},
+    output_type=ClassificationResults,
+)
+class TextClassificationTask(TaskBase):
+    """The text classification task is responsible for assigning a label or class to text."""
+
+
+@task(
+    unary_parameters={"text": str},
+    streaming_parameters={"text_stream": Iterable[str]},
+    unary_output_type=TokenClassificationResults,
+    streaming_output_type=Iterable[TokenClassificationStreamResult],
+)
+class TokenClassificationTask(TaskBase):
+    """The token classification task is responsible for assigning a label to individual
+    tokens in a document.
+    """
+
+
+@task(
+    unary_parameters={"text": str},
+    streaming_parameters={"text_stream": Iterable[str]},
+    unary_output_type=TokenizationResults,
+    streaming_output_type=Iterable[TokenizationStreamResult],
+)
+class TokenizationTask(TaskBase):
+    """The tokenization task is responsible for splitting a document into tokens."""
+
+
+@task(
+    unary_parameters={"text": str},
+    unary_output_type=ClassifiedGeneratedTextResult,
+    streaming_output_type=Iterable[ClassifiedGeneratedTextStreamResult],
+)
+class ClassificationWithTextGenerationTask(TaskBase):
+    """The classification with text generation task is responsible for taking
+    input prompting text, generating additional text from that prompt and classifying
+    the generated text based on detectors.
+    """


Test patch:
diff --git a/tests/interfaces/nlp/test_classification.py b/tests/interfaces/nlp/test_classification.py
new file mode 100644
index 000000000..7100a3929
--- /dev/null
+++ b/tests/interfaces/nlp/test_classification.py
@@ -0,0 +1,229 @@
+# 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 caikit.interfaces.nlp.data_model import (
+    ClassificationResult,
+    ClassificationResults,
+    ClassificationTrainRecord,
+    ClassifiedGeneratedTextResult,
+    TokenClassificationResult,
+    TokenClassificationResults,
+)
+
+## Setup #########################################################################
+
+classification1 = ClassificationResult(label="temperature", score=0.71)
+
+classification2 = ClassificationResult(label="conditions", score=0.98)
+
+classification_result = ClassificationResults(
+    results=[classification1, classification2]
+)
+
+token_classification1 = TokenClassificationResult(
+    start=0, end=5, word="moose", entity="animal", score=0.8
+)
+token_classification2 = TokenClassificationResult(
+    start=7, end=12, word="goose", entity="animal", score=0.7
+)
+token_classification_result = TokenClassificationResults(
+    results=[token_classification1, token_classification2]
+)
+
+classification_train_record = ClassificationTrainRecord(
+    text="It is 20 degrees today", labels=["temperature"]
+)
+
+classification_generated_text_result = ClassifiedGeneratedTextResult(
+    text="moose goose foo bar",
+    token_classification_results=[token_classification1, token_classification2],
+)
+
+## Tests ########################################################################
+
+### ClassificationResult
+
+
+def test_classification_all_fields_accessible():
+    classification_result = ClassificationResult(label="temperature", score=0.71)
+    assert classification_result.label == "temperature"
+    assert classification_result.score == 0.71
+
+
+def test_classification_from_proto_and_back():
+    new = ClassificationResult.from_proto(classification1.to_proto())
+    assert new.label == "temperature"
+    assert new.score == 0.71
+
+
+def test_classification_from_json_and_back():
+    new = ClassificationResult.from_json(classification1.to_json())
+    assert new.label == "temperature"
+    assert new.score == 0.71
+
+
+### ClassificationResults
+
+
+def test_classification_result_all_fields_accessible():
+    classification_result = ClassificationResults(results=[classification1])
+    assert classification_result.results[0].label == "temperature"
+    assert classification_result.results[0].score == 0.71
+
+
+def test_classification_result_from_proto_and_back():
+    new = ClassificationResults.from_proto(classification_result.to_proto())
+    assert new.results[0].label == "temperature"
+    assert new.results[0].score == 0.71
+    assert new.results[1].label == "conditions"
+    assert new.results[1].score == 0.98
+
+
+def test_classification_result_from_json_and_back():
+    new = ClassificationResults.from_json(classification_result.to_json())
+    assert new.results[0].label == "temperature"
+    assert new.results[0].score == 0.71
+    assert new.results[1].label == "conditions"
+    assert new.results[1].score == 0.98
+
+
+### TokenClassificationResult
+
+
+def test_token_classification_all_fields_accessible():
+    token_classification = TokenClassificationResult(
+        start=0,
+        end=28,
+        word="The cow jumped over the moon",
+        entity="neutral",
+        score=0.6,
+    )
+    assert token_classification.start == 0
+    assert token_classification.end == 28
+    assert token_classification.word == "The cow jumped over the moon"
+    assert token_classification.entity == "neutral"
+    assert token_classification.score == 0.6
+
+
+def test_classification_from_proto_and_back():
+    new = TokenClassificationResult.from_proto(token_classification1.to_proto())
+    assert new.start == 0
+    assert new.word == "moose"
+    assert new.score == 0.8
+
+
+def test_classification_from_json_and_back():
+    new = TokenClassificationResult.from_json(token_classification1.to_json())
+    assert new.start == 0
+    assert new.word == "moose"
+    assert new.score == 0.8
+
+
+### TokenClassificationResults
+
+
+def test_token_classification_result_all_fields_accessible():
+    token_classification_result = TokenClassificationResults(
+        results=[token_classification1]
+    )
+    assert token_classification_result.results[0].start == 0
+    assert token_classification_result.results[0].word == "moose"
+    assert token_classification_result.results[0].score == 0.8
+
+
+def test_token_classification_result_from_proto_and_back():
+    new = TokenClassificationResults.from_proto(token_classification_result.to_proto())
+    assert new.results[0].start == 0
+    assert new.results[0].word == "moose"
+    assert new.results[0].score == 0.8
+    assert new.results[1].end == 12
+    assert new.results[1].entity == "animal"
+
+
+def test_classification_result_from_json_and_back():
+    new = TokenClassificationResults.from_json(token_classification_result.to_json())
+    assert new.results[0].start == 0
+    assert new.results[0].word == "moose"
+    assert new.results[0].score == 0.8
+    assert new.results[1].end == 12
+    assert new.results[1].entity == "animal"
+
+
+### ClassificationTrainRecord
+
+
+def test_classification_train_record_all_fields_accessible():
+    classification_train_record = ClassificationTrainRecord(
+        text="It is 20 degrees today", labels=["temperature"]
+    )
+    assert classification_train_record.text == "It is 20 degrees today"
+    assert classification_train_record.labels == ["temperature"]
+
+
+def test_classification_train_record_from_proto_and_back():
+    new = ClassificationTrainRecord.from_proto(classification_train_record.to_proto())
+    assert new.text == "It is 20 degrees today"
+    assert new.labels == ["temperature"]
+
+
+def test_classification_train_record_from_json_and_back():
+    new = ClassificationTrainRecord.from_json(classification_train_record.to_json())
+    assert new.text == "It is 20 degrees today"
+    assert new.labels == ["temperature"]
+
+
+### ClassifiedGeneratedTextResult
+
+
+def test_classification_generated_text_result_all_fields_accessible():
+    classification_generated_text_result = ClassifiedGeneratedTextResult(
+        text="moose goose foo bar",
+        token_classification_results=[token_classification1, token_classification2],
+    )
+    assert classification_generated_text_result.text == "moose goose foo bar"
+    assert (
+        classification_generated_text_result.token_classification_results[0]
+        == token_classification1
+    )
+    assert (
+        classification_generated_text_result.token_classification_results[1]
+        == token_classification2
+    )
+
+
+def test_classification_generated_text_result_from_proto_and_back():
+    new = ClassifiedGeneratedTextResult.from_proto(
+        classification_generated_text_result.to_proto()
+    )
+    assert new.text == "moose goose foo bar"
+    assert new.token_classification_results[0].start == 0
+    assert new.token_classification_results[0].word == "moose"
+    assert new.token_classification_results[0].score == 0.8
+    assert new.token_classification_results[1].start == 7
+    assert new.token_classification_results[1].word == "goose"
+    assert new.token_classification_results[1].score == 0.7
+
+
+def test_classification_generated_text_result_from_json_and_back():
+    new = ClassifiedGeneratedTextResult.from_json(
+        classification_generated_text_result.to_json()
+    )
+    assert new.text == "moose goose foo bar"
+    assert new.token_classification_results[0].start == 0
+    assert new.token_classification_results[0].word == "moose"
+    assert new.token_classification_results[0].score == 0.8
+    assert new.token_classification_results[1].start == 7
+    assert new.token_classification_results[1].word == "goose"
+    assert new.token_classification_results[1].score == 0.7
diff --git a/tests/interfaces/nlp/test_text.py b/tests/interfaces/nlp/test_text.py
new file mode 100644
index 000000000..e6234df28
--- /dev/null
+++ b/tests/interfaces/nlp/test_text.py
@@ -0,0 +1,68 @@
+# 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 caikit.interfaces.nlp.data_model import Token, TokenizationResults
+
+## Setup #########################################################################
+
+dummy_token = Token(start=0, end=11, text="Hello World")
+dummy_token2 = Token(start=0, end=7, text="foo bar")
+tokenization_result = TokenizationResults([dummy_token, dummy_token2])
+
+## Tests ########################################################################
+
+### Token
+
+
+def test_token_all_fields_accessible():
+    token = Token(start=0, end=11, text="Hello World")
+    assert token.start == 0
+    assert token.end == 11
+    assert token.text == "Hello World"
+
+
+def test_token_from_proto_and_back():
+    new = Token.from_proto(dummy_token.to_proto())
+    assert new.start == 0
+    assert new.end == 11
+    assert new.text == "Hello World"
+
+
+def test_token_from_json_and_back():
+    new = Token.from_json(dummy_token.to_json())
+    assert new.start == 0
+    assert new.end == 11
+    assert new.text == "Hello World"
+
+
+### TokenizationResults
+
+
+def test_tokenization_result_all_fields_accessible():
+    tokenization_result = TokenizationResults([dummy_token, dummy_token2])
+    assert tokenization_result.results[0] == dummy_token
+    assert tokenization_result.results[1] == dummy_token2
+
+
+def test_tokenization_result_from_proto_and_back():
+    new = TokenizationResults.from_proto(tokenization_result.to_proto())
+    assert new.results[0] == dummy_token
+    assert new.results[1] == dummy_token2
+
+
+def test_tokenization_result_from_json_and_back():
+    new = TokenizationResults.from_json(tokenization_result.to_json())
+    assert new.results[0] == dummy_token
+    assert new.results[1] == dummy_token2