diff --git a/mediapipe/tasks/python/components/containers/embedding_result.py b/mediapipe/tasks/python/components/containers/embedding_result.py index 8ddbb3ae5..63c315603 100644 --- a/mediapipe/tasks/python/components/containers/embedding_result.py +++ b/mediapipe/tasks/python/components/containers/embedding_result.py @@ -26,42 +26,6 @@ _EmbeddingProto = embeddings_pb2.Embedding _EmbeddingResultProto = embeddings_pb2.EmbeddingResult -@dataclasses.dataclass -class FloatEmbedding: - """Defines a dense floating-point embedding. - - Attributes: - values: A NumPy array indicating the raw output of the embedding layer. - """ - - values: np.ndarray - - @classmethod - @doc_controls.do_not_generate_docs - def create_from_pb2(cls, pb2_obj: _FloatEmbeddingProto) -> 'FloatEmbedding': - """Creates a `FloatEmbedding` object from the given protobuf object.""" - return FloatEmbedding(values=np.array(pb2_obj.values, dtype=float)) - - -@dataclasses.dataclass -class QuantizedEmbedding: - """Defines a dense scalar-quantized embedding. - - Attributes: - values: A NumPy array indicating the raw output of the embedding layer. - """ - - values: np.ndarray - - @classmethod - @doc_controls.do_not_generate_docs - def create_from_pb2( - cls, pb2_obj: _QuantizedEmbeddingProto) -> 'QuantizedEmbedding': - """Creates a `QuantizedEmbedding` object from the given protobuf object.""" - return QuantizedEmbedding( - values=np.array(bytearray(pb2_obj.values), dtype=np.uint8)) - - @dataclasses.dataclass class Embedding: """Embedding result for a given embedder head. diff --git a/mediapipe/tasks/python/test/text/text_embedder_test.py b/mediapipe/tasks/python/test/text/text_embedder_test.py index 8ef616150..e1f1a45e8 100644 --- a/mediapipe/tasks/python/test/text/text_embedder_test.py +++ b/mediapipe/tasks/python/test/text/text_embedder_test.py @@ -15,7 +15,6 @@ import enum import os -from unittest import mock import numpy as np from absl.testing import absltest @@ -30,8 +29,6 @@ from mediapipe.tasks.python.text import text_embedder ImageEmbedderResult = embedding_result_module.EmbeddingResult _BaseOptions = base_options_module.BaseOptions _EmbedderOptions = embedder_options_module.EmbedderOptions -_FloatEmbedding = embedding_result_module.FloatEmbedding -_QuantizedEmbedding = embedding_result_module.QuantizedEmbedding _Embedding = embedding_result_module.Embedding _TextEmbedder = text_embedder.TextEmbedder _TextEmbedderOptions = text_embedder.TextEmbedderOptions