Internal MediaPipe Tasks change
PiperOrigin-RevId: 525182282
This commit is contained in:
parent
3e0ed2ced0
commit
64d1e74c20
|
@ -32,6 +32,7 @@ _TextEmbedderOptions = text_embedder.TextEmbedderOptions
|
|||
|
||||
_BERT_MODEL_FILE = 'mobilebert_embedding_with_metadata.tflite'
|
||||
_REGEX_MODEL_FILE = 'regex_one_embedding_with_metadata.tflite'
|
||||
_USE_MODEL_FILE = 'universal_sentence_encoder_qa_with_metadata.tflite'
|
||||
_TEST_DATA_DIR = 'mediapipe/tasks/testdata/text'
|
||||
# Tolerance for embedding vector coordinate values.
|
||||
_EPSILON = 1e-4
|
||||
|
@ -138,6 +139,24 @@ class TextEmbedderTest(parameterized.TestCase):
|
|||
16,
|
||||
(0.549632, 0.552879),
|
||||
),
|
||||
(
|
||||
False,
|
||||
False,
|
||||
_USE_MODEL_FILE,
|
||||
ModelFileType.FILE_NAME,
|
||||
0.851961,
|
||||
100,
|
||||
(1.422951, 1.404664),
|
||||
),
|
||||
(
|
||||
True,
|
||||
False,
|
||||
_USE_MODEL_FILE,
|
||||
ModelFileType.FILE_CONTENT,
|
||||
0.851961,
|
||||
100,
|
||||
(0.127049, 0.125416),
|
||||
),
|
||||
)
|
||||
def test_embed(self, l2_normalize, quantize, model_name, model_file_type,
|
||||
expected_similarity, expected_size, expected_first_values):
|
||||
|
@ -213,6 +232,24 @@ class TextEmbedderTest(parameterized.TestCase):
|
|||
16,
|
||||
(0.549632, 0.552879),
|
||||
),
|
||||
(
|
||||
False,
|
||||
False,
|
||||
_USE_MODEL_FILE,
|
||||
ModelFileType.FILE_NAME,
|
||||
0.851961,
|
||||
100,
|
||||
(1.422951, 1.404664),
|
||||
),
|
||||
(
|
||||
True,
|
||||
False,
|
||||
_USE_MODEL_FILE,
|
||||
ModelFileType.FILE_CONTENT,
|
||||
0.851961,
|
||||
100,
|
||||
(0.127049, 0.125416),
|
||||
),
|
||||
)
|
||||
def test_embed_in_context(self, l2_normalize, quantize, model_name,
|
||||
model_file_type, expected_similarity, expected_size,
|
||||
|
@ -251,6 +288,7 @@ class TextEmbedderTest(parameterized.TestCase):
|
|||
@parameterized.parameters(
|
||||
# TODO: The similarity should likely be lower
|
||||
(_BERT_MODEL_FILE, 0.980880),
|
||||
(_USE_MODEL_FILE, 0.780334),
|
||||
)
|
||||
def test_embed_with_different_themes(self, model_file, expected_similarity):
|
||||
# Creates embedder.
|
||||
|
|
Loading…
Reference in New Issue
Block a user