Repository: caikit/caikit
Pull request #573 (2023-11-09 21:18:27)

Embedding service data models and tasks for caikit/interfaces
Embedding services includes
* embedding (vectors)
* sentence similarity (scores)
* rerank (top-n documents)

<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:
Adds data models and tasks in caikit/interfaces to be used for embedding services.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [x] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/interfaces/common/data_model/__init__.py b/caikit/interfaces/common/data_model/__init__.py
index 188143db..e65c6fae 100644
--- a/caikit/interfaces/common/data_model/__init__.py
+++ b/caikit/interfaces/common/data_model/__init__.py
@@ -25,8 +25,15 @@ Domain agnostic data model objects
 
 # Local
 # Import individual packages
-from . import primitive_sequences, producer
+from . import primitive_sequences, producer, vectors
 from .file import File
 from .primitive_sequences import BoolSequence, FloatSequence, IntSequence, StrSequence
 from .producer import ProducerId
 from .stream_sources import FileReference, ListOfFileReferences
+from .vectors import (
+    ListOfVector1D,
+    NpFloat32Sequence,
+    NpFloat64Sequence,
+    PyFloatSequence,
+    Vector1D,
+)
diff --git a/caikit/interfaces/common/data_model/vectors.py b/caikit/interfaces/common/data_model/vectors.py
new file mode 100644
index 00000000..5189761d
--- /dev/null
+++ b/caikit/interfaces/common/data_model/vectors.py
@@ -0,0 +1,198 @@
+# 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 embedding vector representations
+"""
+# Standard
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Union
+import json
+
+# Third Party
+from google.protobuf import json_format
+import numpy as np
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+import alog
+
+# Local
+from caikit.core import DataObjectBase, dataobject
+from caikit.core.data_model import PACKAGE_COMMON
+from caikit.core.exceptions import error_handler
+
+log = alog.use_channel("DATAM")
+error = error_handler.get(log)
+
+
+@dataobject(PACKAGE_COMMON)
+@dataclass
+class PyFloatSequence(DataObjectBase):
+    values: Annotated[List[float], FieldNumber(1)] = field(default_factory=list)
+
+
+@dataobject(PACKAGE_COMMON)
+@dataclass
+class NpFloat32Sequence(DataObjectBase):
+    values: Annotated[List[np.float32], FieldNumber(1)]
+
+    @classmethod
+    def from_proto(cls, proto):
+        values = np.asarray(proto.values, dtype=np.float32)
+        return cls(values)
+
+
+@dataobject(PACKAGE_COMMON)
+@dataclass
+class NpFloat64Sequence(DataObjectBase):
+    values: Annotated[List[np.float64], FieldNumber(1)]
+
+    @classmethod
+    def from_proto(cls, proto):
+        values = np.asarray(proto.values, dtype=np.float64)
+        return cls(values)
+
+
+@dataobject(PACKAGE_COMMON)
+@dataclass
+class Vector1D(DataObjectBase):
+    """Data representation for a 1 dimension vector of float-type data."""
+
+    data: Annotated[
+        Union[
+            PyFloatSequence,
+            NpFloat32Sequence,
+            NpFloat64Sequence,
+        ],
+        FieldNumber(1),
+    ]
+
+    def __post_init__(self):
+        error.value_check(
+            "<NLP92989048E>",
+            hasattr(self.data, "values"),
+            ValueError("Vector1D requires a float sequence data object with values."),
+        )
+
+    @classmethod
+    def from_vector(cls, vector):
+        dtype = getattr(vector, "dtype", False)
+        if dtype is None:
+            data = PyFloatSequence(vector)
+        elif dtype == np.float32:
+            data = NpFloat32Sequence(vector)
+        elif dtype == np.float64:
+            data = NpFloat64Sequence(vector)
+        else:
+            data = PyFloatSequence(vector)
+        return cls(data=data)
+
+    @classmethod
+    def from_json(cls, json_str: Union[Dict[str, Any], str]) -> "Vector1D":
+        """JSON does not have different float types. Move data into data_pyfloatsequence"""
+
+        json_obj = json.loads(json_str) if isinstance(json_str, str) else json_str
+        data = json_obj.pop("data")
+        if data is not None:
+            json_obj["data_pyfloatsequence"] = data
+
+        json_str = json.dumps(json_obj)
+        try:
+            # Parse given JSON into google.protobufs.pyext.cpp_message.GeneratedProtocolMessageType
+            parsed_proto = json_format.Parse(
+                json_str, cls.get_proto_class()(), ignore_unknown_fields=False
+            )
+
+            # Use from_proto to return the DataBase object from the parsed proto
+            return cls.from_proto(parsed_proto)
+
+        except json_format.ParseError as ex:
+            error("<NLP39795399E>", ValueError(ex))
+
+    def to_dict(self) -> dict:
+        """to_dict is needed to make things serializable"""
+        values = self.data.values if self.data.values is not None else []
+        return {
+            "data": {
+                # coerce numpy.ndarray and numpy.float32 into JSON serializable list of floats
+                "values": values.tolist()
+                if isinstance(values, np.ndarray)
+                else values
+            }
+        }
+
+    @classmethod
+    def from_proto(cls, proto):
+        """Wrap the data in an appropriate float sequence, wrapped by this class"""
+        woo = proto.WhichOneof("data")
+        if woo is None:
+            return cls(PyFloatSequence())
+
+        woo_data = getattr(proto, woo)
+        if woo == "data_npfloat64sequence":
+            ret = cls(NpFloat64Sequence.from_proto(woo_data))
+        elif woo == "data_npfloat32sequence":
+            ret = cls(NpFloat32Sequence.from_proto(woo_data))
+        else:
+            ret = cls(PyFloatSequence.from_proto(woo_data))
+        return ret
+
+    def fill_proto(self, proto):
+        """Fill in the data in an appropriate data_<float type sequence>"""
+        values = self.data.values
+        if values is not None and len(values) > 0:
+            sample = values[0]
+            error.type_check(
+                "<NLP47515960E>", float, np.float32, np.float64, sample=sample
+            )
+            if isinstance(sample, np.float64):
+                proto.data_npfloat64sequence.values.extend(values)
+            elif isinstance(sample, np.float32):
+                proto.data_npfloat32sequence.values.extend(values)
+            else:
+                proto.data_pyfloatsequence.values.extend(values)
+
+        return proto
+
+
+@dataobject(PACKAGE_COMMON)
+class ListOfVector1D(DataObjectBase):
+    """Data representation for an embedding matrix holding 2D vectors"""
+
+    vectors: Annotated[List[Vector1D], FieldNumber(1)]
+
+    def __post_init__(self):
+        error.type_check("<NLP94336739E>", list, vectors=self.vectors)
+        error.type_check_all("<NLP94783841E>", Vector1D, vectors=self.vectors)
+
+    @classmethod
+    def from_json(cls, json_str: Union[Dict[str, Any], str]) -> "ListOfVector1D":
+        """Fill in the vector data in an appropriate data_<float type sequence>"""
+
+        json_obj = json.loads(json_str) if isinstance(json_str, str) else json_str
+        for v in json_obj["vectors"]:
+            data = v.pop("data")
+            if data is not None:
+                v["data_pyfloatsequence"] = data
+        json_str = json.dumps(json_obj)
+        try:
+            # Parse given JSON into google.protobufs.pyext.cpp_message.GeneratedProtocolMessageType
+            parsed_proto = json_format.Parse(
+                json_str, cls.get_proto_class()(), ignore_unknown_fields=False
+            )
+
+            # Use from_proto to return the DataBase object from the parsed proto
+            return cls.from_proto(parsed_proto)
+
+        except json_format.ParseError as ex:
+            error("<NLP85443094E>", ValueError(ex))
diff --git a/caikit/interfaces/nlp/data_model/__init__.py b/caikit/interfaces/nlp/data_model/__init__.py
index ff1e6d71..90905954 100644
--- a/caikit/interfaces/nlp/data_model/__init__.py
+++ b/caikit/interfaces/nlp/data_model/__init__.py
@@ -13,7 +13,15 @@
 # limitations under the License.
 
 # Local
-from . import classification, package, text, text_generation
+from . import (
+    classification,
+    embedding_vectors,
+    package,
+    reranker,
+    sentence_similarity,
+    text,
+    text_generation,
+)
 from .classification import (
     ClassificationResult,
     ClassificationResults,
@@ -24,7 +32,14 @@ from .classification import (
     TokenClassificationResults,
     TokenClassificationStreamResult,
 )
+from .embedding_vectors import EmbeddingResult, EmbeddingResults
 from .package import NLP_PACKAGE
+from .reranker import RerankResult, RerankResults, RerankScore, RerankScores
+from .sentence_similarity import (
+    SentenceSimilarityResult,
+    SentenceSimilarityResults,
+    SentenceSimilarityScores,
+)
 from .text import Token, TokenizationResults, TokenizationStreamResult
 from .text_generation import (
     FinishReason,
diff --git a/caikit/interfaces/nlp/data_model/embedding_vectors.py b/caikit/interfaces/nlp/data_model/embedding_vectors.py
new file mode 100644
index 00000000..30163f14
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/embedding_vectors.py
@@ -0,0 +1,47 @@
+# 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 embedding vector representations
+"""
+# Standard
+from dataclasses import dataclass
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+import alog
+
+# Local
+from ...common.data_model import ListOfVector1D, ProducerId, Vector1D
+from caikit.core import DataObjectBase, dataobject
+from caikit.core.exceptions import error_handler
+
+log = alog.use_channel("DATAM")
+error = error_handler.get(log)
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+@dataclass
+class EmbeddingResult(DataObjectBase):
+    """Result from text embedding task"""
+
+    result: Annotated[Vector1D, FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+@dataclass
+class EmbeddingResults(DataObjectBase):
+    """Results from text embeddings task"""
+
+    results: Annotated[ListOfVector1D, FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
diff --git a/caikit/interfaces/nlp/data_model/reranker.py b/caikit/interfaces/nlp/data_model/reranker.py
new file mode 100644
index 00000000..be700caf
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/reranker.py
@@ -0,0 +1,66 @@
+# 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.
+
+# Standard
+from typing import List, Optional
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+
+# Local
+from ...common.data_model import ProducerId
+from caikit.core import DataObjectBase, dataobject
+from caikit.core.data_model.json_dict import JsonDict
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class RerankScore(DataObjectBase):
+    """The score for one document (one query)"""
+
+    document: Annotated[Optional[JsonDict], FieldNumber(1)]
+    index: Annotated[int, FieldNumber(2)]
+    score: Annotated[float, FieldNumber(3)]
+    text: Annotated[Optional[str], FieldNumber(4)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class RerankScores(DataObjectBase):
+    """Scores for a query in a rerank task.
+    This is a list of n ReRankScore where n is based on top_n documents and each score indicates
+    the relevance of that document for this query. Results are ordered most-relevant first.
+    """
+
+    query: Annotated[Optional[str], FieldNumber(1)]
+    scores: Annotated[List[RerankScore], FieldNumber(2)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class RerankResult(DataObjectBase):
+    """Result for one query in a rerank task.
+    This is a list of n ReRankScore where n is based on top_n documents and each score indicates
+    the relevance of that document for this query. Results are ordered most-relevant first.
+    """
+
+    result: Annotated[RerankScores, FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class RerankResults(DataObjectBase):
+    """Results list for rerank tasks (supporting multiple queries).
+    For multiple queries, each one has a RerankQueryResult (ranking the documents for that query).
+    """
+
+    results: Annotated[List[RerankScores], FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
diff --git a/caikit/interfaces/nlp/data_model/sentence_similarity.py b/caikit/interfaces/nlp/data_model/sentence_similarity.py
new file mode 100644
index 00000000..5c908a1d
--- /dev/null
+++ b/caikit/interfaces/nlp/data_model/sentence_similarity.py
@@ -0,0 +1,52 @@
+# 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 embedding vector representations
+"""
+# Standard
+from typing import List
+
+# First Party
+from py_to_proto.dataclass_to_proto import Annotated, FieldNumber
+import alog
+
+# Local
+from ...common.data_model import ProducerId
+from caikit.core import DataObjectBase, dataobject
+from caikit.core.exceptions import error_handler
+
+log = alog.use_channel("DATAM")
+error = error_handler.get(log)
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class SentenceSimilarityScores(DataObjectBase):
+    """Scores for a sentence similarity task"""
+
+    scores: Annotated[List[float], FieldNumber(1)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class SentenceSimilarityResult(DataObjectBase):
+    """Result for sentence similarity task"""
+
+    result: Annotated[SentenceSimilarityScores, FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
+
+
+@dataobject(package="caikit_data_model.caikit_nlp")
+class SentenceSimilarityResults(DataObjectBase):
+    """Results list for sentence similarity tasks"""
+
+    results: Annotated[List[SentenceSimilarityScores], FieldNumber(1)]
+    producer_id: Annotated[ProducerId, FieldNumber(2)]
diff --git a/caikit/interfaces/nlp/tasks.py b/caikit/interfaces/nlp/tasks.py
index c8e16689..73f5ee53 100644
--- a/caikit/interfaces/nlp/tasks.py
+++ b/caikit/interfaces/nlp/tasks.py
@@ -16,10 +16,12 @@ This module holds the Task definitions for all common NLP tasks
 """
 
 # Standard
-from typing import Iterable
+from typing import Iterable, List
 
 # Local
 from ...core import TaskBase, task
+from ...core.data_model.json_dict import JsonDict
+from .data_model import SentenceSimilarityResult, SentenceSimilarityResults
 from .data_model.classification import (
     ClassificationResults,
     ClassifiedGeneratedTextResult,
@@ -27,6 +29,8 @@ from .data_model.classification import (
     TokenClassificationResults,
     TokenClassificationStreamResult,
 )
+from .data_model.embedding_vectors import EmbeddingResult, EmbeddingResults
+from .data_model.reranker import RerankResult, RerankResults
 from .data_model.text import TokenizationResults, TokenizationStreamResult
 from .data_model.text_generation import GeneratedTextResult, GeneratedTextStreamResult
 
@@ -82,3 +86,85 @@ class ClassificationWithTextGenerationTask(TaskBase):
     input prompting text, generating additional text from that prompt and classifying
     the generated text based on detectors.
     """
+
+
+@task(
+    required_parameters={"text": str},
+    output_type=EmbeddingResult,
+)
+class EmbeddingTask(TaskBase):
+    """Return a text embedding for the input text string"""
+
+
+@task(
+    required_parameters={"texts": List[str]},
+    output_type=EmbeddingResults,
+)
+class EmbeddingTasks(TaskBase):
+    """Return a text embedding for each text string in the input list"""
+
+
+@task(
+    required_parameters={
+        "documents": List[JsonDict],
+        "query": str,
+    },
+    output_type=RerankResult,
+)
+class RerankTask(TaskBase):
+    """Returns an ordered list ranking the most relevant documents for the query
+
+    Required parameters:
+        query: The search query
+        documents: JSON documents containing "text" or alternative "_text" to search
+    Returns:
+        The top_n documents in order of relevance (most relevant first).
+        For each, a score and document index (position in input) is returned.
+        The original document JSON is returned depending on optional args.
+        The top_n optional parameter limits the results when used.
+    """
+
+
+@task(
+    required_parameters={
+        "documents": List[JsonDict],
+        "queries": List[str],
+    },
+    output_type=RerankResults,
+)
+class RerankTasks(TaskBase):
+    """Returns an ordered list for each query ranking the most relevant documents for the query
+
+    Required parameters:
+        queries: The search queries
+        documents: JSON documents containing "text" or alternative "_text" to search
+    Returns:
+        Results in order of the queries.
+        In each query result:
+            The query text is optionally included for visual convenience.
+            The top_n documents in order of relevance (most relevant first).
+            For each, a score and document index (position in input) is returned.
+            The original document JSON is returned depending on optional args.
+            The top_n optional parameter limits the results when used.
+    """
+
+
+@task(
+    required_parameters={"source_sentence": str, "sentences": List[str]},
+    output_type=SentenceSimilarityResult,
+)
+class SentenceSimilarityTask(TaskBase):
+    """Compare the source_sentence to each of the sentences.
+    Result contains a list of scores in the order of the input sentences.
+    """
+
+
+@task(
+    required_parameters={"source_sentences": List[str], "sentences": List[str]},
+    output_type=SentenceSimilarityResults,
+)
+class SentenceSimilarityTasks(TaskBase):
+    """Compare each of the source_sentences to each of the sentences.
+    Returns a list of results in the order of the source_sentences.
+    Each result contains a list of scores in the order of the input sentences.
+    """


Test patch:
diff --git a/tests/interfaces/common/test_vectors.py b/tests/interfaces/common/test_vectors.py
new file mode 100644
index 00000000..aa315508
--- /dev/null
+++ b/tests/interfaces/common/test_vectors.py
@@ -0,0 +1,152 @@
+# 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.
+"""Test for embedding vectors
+"""
+# Standard
+from collections import namedtuple
+
+# Third Party
+import numpy as np
+import pytest
+
+# Local
+from caikit.interfaces.common import data_model as dm
+
+## Setup #########################################################################
+
+DUMMY_VECTOR_SHAPE = (5,)
+RANDOM_SEED = 77
+np.random.seed(RANDOM_SEED)
+random_number_generator = np.random.default_rng()
+
+# To tests the limits of our type-checking, this can replace our legit data objects
+TRICK_SEQUENCE = namedtuple("Trick", "values")
+
+
+@pytest.fixture
+def simple_array_of_floats():
+    return [1.1, 2.2]
+
+
+@pytest.fixture
+def simple_array_of_ints():
+    return ["foo", 1, 2, 3, 4]
+
+
+@pytest.fixture
+def random_numpy_vector1d_float32():
+    return random_number_generator.random(DUMMY_VECTOR_SHAPE, dtype=np.float32)
+
+
+@pytest.fixture
+def random_numpy_vector1d_float64():
+    return random_number_generator.random(DUMMY_VECTOR_SHAPE, dtype=np.float64)
+
+
+@pytest.fixture
+def random_python_vector1d_float(random_numpy_vector1d_float32):
+    return random_numpy_vector1d_float32.tolist()
+
+
+## Tests ########################################################################
+
+
+@pytest.mark.parametrize(
+    "sequence",
+    [
+        dm.PyFloatSequence(),
+        dm.NpFloat32Sequence(),
+        dm.NpFloat64Sequence(),
+        TRICK_SEQUENCE(values=None),
+    ],
+    ids=type,
+)
+def test_empty_sequences(sequence):
+    """No type check error with empty sequences"""
+    new_dm_from_init = dm.Vector1D(sequence)
+    assert isinstance(new_dm_from_init.data, type(sequence))
+    assert new_dm_from_init.data.values is None
+
+    # Test proto
+    proto_from_dm = new_dm_from_init.to_proto()
+    new_dm_from_proto = dm.Vector1D.from_proto(proto_from_dm)
+    assert isinstance(new_dm_from_proto, dm.Vector1D)
+    assert new_dm_from_proto.data.values is None
+
+    # Test json
+    json_from_dm = new_dm_from_init.to_json()
+    new_dm_from_json = dm.Vector1D.from_json(json_from_dm)
+    assert isinstance(new_dm_from_json, dm.Vector1D)
+    assert new_dm_from_json.data.values == []
+
+
+def test_vector1d_iterator_error():
+    """Cannot just shove in an iterator and expect it to work"""
+    with pytest.raises(ValueError):
+        dm.Vector1D(data=[1.1, 2.2, 3.3])
+
+
+def _assert_array_check(new_array, data_values, float_type):
+    for value in new_array.data.values:
+        assert isinstance(value, float_type)
+    np.testing.assert_array_equal(new_array.data.values, data_values)
+
+
+@pytest.mark.parametrize(
+    "float_seq_class, random_values, float_type",
+    [
+        (dm.PyFloatSequence, "random_python_vector1d_float", float),
+        (dm.NpFloat32Sequence, "random_numpy_vector1d_float32", np.float32),
+        (dm.NpFloat64Sequence, "random_numpy_vector1d_float64", np.float64),
+        (
+            TRICK_SEQUENCE,
+            "simple_array_of_floats",
+            float,
+        ),  # Sneaky but tests corner cases for now
+    ],
+)
+def test_vector1d_dm(float_seq_class, random_values, float_type, request):
+
+    # Test init
+    fixture_values = request.getfixturevalue(random_values)
+    dm_init = dm.Vector1D(data=float_seq_class(fixture_values))
+    _assert_array_check(dm_init, fixture_values, float_type)
+
+    # Test proto
+    dm_to_proto = dm_init.to_proto()
+    dm_from_proto = dm.Vector1D.from_proto(dm_to_proto)
+    _assert_array_check(dm_from_proto, fixture_values, float_type)
+
+    # Test json
+    dm_to_json = dm_init.to_json()
+    dm_from_json = dm.Vector1D.from_json(dm_to_json)
+    _assert_array_check(
+        dm_from_json, fixture_values, float
+    )  # NOTE: always float after json
+
+
+@pytest.mark.parametrize(
+    "float_seq_class, random_values, float_type",
+    [
+        (dm.PyFloatSequence, "random_python_vector1d_float", float),
+        (dm.NpFloat32Sequence, "random_numpy_vector1d_float32", np.float32),
+        (dm.NpFloat64Sequence, "random_numpy_vector1d_float64", np.float64),
+    ],
+)
+def test_vector1d_dm_from_vector(float_seq_class, random_values, float_type, request):
+    fixture_values = request.getfixturevalue(random_values)
+    v = dm.Vector1D.from_vector(fixture_values)
+    assert isinstance(v.data, float_seq_class)
+    assert isinstance(v.data.values[0], float_type)
+    _assert_array_check(v, fixture_values, float_type)
diff --git a/tests/interfaces/nlp/test_reranker.py b/tests/interfaces/nlp/test_reranker.py
new file mode 100644
index 00000000..16840ef2
--- /dev/null
+++ b/tests/interfaces/nlp/test_reranker.py
@@ -0,0 +1,188 @@
+# 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.
+"""Test for reranker
+"""
+
+# Standard
+import random
+import string
+
+# Third Party
+import pytest
+
+# Local
+from caikit.interfaces.nlp import data_model as dm
+
+## Setup #########################################################################
+
+
+@pytest.fixture
+def input_document():
+    return {
+        "text": "this is the input text",
+        "_text": "alternate _text here",
+        "title": "some title attribute here",
+        "anything": "another string attribute",
+        "str_test": "test string",
+        "int_test": 1234,
+        "float_test": 9876.4321,
+    }
+
+
+@pytest.fixture
+def input_random_document():
+    return {
+        "text": "".join(random.choices(string.printable, k=100)),
+        "random_str": "".join(random.choices(string.printable, k=100)),
+        "random_int": random.randint(-99999, 99999),
+        "random_float": random.uniform(-99999, 99999),
+    }
+
+
+@pytest.fixture
+def input_documents(input_document, input_random_document):
+    return [input_document, input_random_document]
+
+
+@pytest.fixture
+def input_score(input_document):
+    return {
+        "document": input_document,
+        "index": 1234,
+        "score": 9876.54321,
+        "text": "this is the input text",
+    }
+
+
+@pytest.fixture
+def input_random_score(input_random_document):
+    return {
+        "document": input_random_document,
+        "index": random.randint(-99999, 99999),
+        "score": random.uniform(-99999, 99999),
+        "text": "".join(random.choices(string.printable, k=100)),
+    }
+
+
+@pytest.fixture
+def input_random_score_3():
+    return {
+        "document": {"text": "random foo3"},
+        "index": random.randint(-99999, 99999),
+        "score": random.uniform(-99999, 99999),
+        "text": "".join(random.choices(string.printable, k=100)),
+    }
+
+
+@pytest.fixture
+def input_scores(input_score, input_random_score):
+    return [dm.RerankScore(**input_score), dm.RerankScore(**input_random_score)]
+
+
+@pytest.fixture
+def input_scores2(input_random_score, input_random_score_3):
+    return [
+        dm.RerankScore(**input_random_score),
+        dm.RerankScore(**input_random_score_3),
+    ]
+
+
+@pytest.fixture
+def input_result_1(input_scores):
+    return {"result": dm.RerankScores(query="foo", scores=input_scores)}
+
+
+@pytest.fixture
+def input_result_2(input_scores2):
+    return {"result": dm.RerankScores(query="bar", scores=input_scores2)}
+
+
+@pytest.fixture
+def input_results(input_scores, input_scores2):
+    return {
+        "results": [
+            dm.RerankScores(query="foo", scores=input_scores),
+            dm.RerankScores(query="bar", scores=input_scores2),
+        ]
+    }
+
+
+@pytest.fixture
+def input_sentence_similarity_scores_1():
+    return {"scores": [random.uniform(-99999, 99999) for _ in range(10)]}
+
+
+@pytest.fixture
+def input_sentence_similarity_result(input_sentence_similarity_scores_1):
+    return {"result": dm.SentenceSimilarityScores(**input_sentence_similarity_scores_1)}
+
+
+@pytest.fixture
+def input_sentence_similarity_scores_2():
+    return {"scores": [random.uniform(-99999, 99999) for _ in range(10)]}
+
+
+@pytest.fixture
+def input_sentence_similarities_scores(
+    input_sentence_similarity_scores_1, input_sentence_similarity_scores_2
+):
+    return [
+        dm.SentenceSimilarityScores(**input_sentence_similarity_scores_1),
+        dm.SentenceSimilarityScores(**input_sentence_similarity_scores_2),
+    ]
+
+
+@pytest.fixture
+def input_sentence_similarity_results(input_sentence_similarities_scores):
+    return {"results": input_sentence_similarities_scores}
+
+
+## Tests ########################################################################
+
+
+@pytest.mark.parametrize(
+    "data_object, inputs",
+    [
+        (dm.RerankScore, "input_score"),
+        (dm.RerankScore, "input_random_score"),
+        (dm.RerankResult, "input_result_1"),
+        (dm.RerankResults, "input_results"),
+        (dm.SentenceSimilarityResult, "input_sentence_similarity_result"),
+        (dm.SentenceSimilarityResults, "input_sentence_similarity_results"),
+    ],
+)
+def test_data_object(data_object, inputs, request):
+    # Init data object
+    fixture_values = request.getfixturevalue(inputs)
+    new_do_from_init = data_object(**fixture_values)
+    assert isinstance(new_do_from_init, data_object)
+    assert_fields_match(new_do_from_init, fixture_values)
+
+    # Test to/from proto
+    proto_from_dm = new_do_from_init.to_proto()
+    new_do_from_proto = data_object.from_proto(proto_from_dm)
+    assert isinstance(new_do_from_proto, data_object)
+    assert_fields_match(new_do_from_proto, fixture_values)
+    assert new_do_from_init == new_do_from_proto
+
+    # Test to/from json
+    json_from_dm = new_do_from_init.to_json()
+    new_do_from_json = data_object.from_json(json_from_dm)
+    assert isinstance(new_do_from_json, data_object)
+    assert_fields_match(new_do_from_json, fixture_values)
+    assert new_do_from_init == new_do_from_json
+
+
+def assert_fields_match(data_object, inputs):
+    assert all(getattr(data_object, key) == value for key, value in inputs.items())