EmbeddingAggregationCalculator should fill in the timestamp_ms field of the embedding results in the stream mode.

Per user feedback, the consistency between the packet timestamp and the timestamp field of the embedding result helps reducing the confusion.

PiperOrigin-RevId: 501572379
This commit is contained in:
Jiuqiang Tang 2023-01-12 08:27:57 -08:00 committed by Copybara-Service
parent 9cbb76939d
commit 5c74ed2ae5
3 changed files with 12 additions and 7 deletions

View File

@ -120,7 +120,9 @@ absl::Status EmbeddingAggregationCalculator::Process(CalculatorContext* cc) {
} }
kTimestampedEmbeddingsOut(cc).Send(std::move(results)); kTimestampedEmbeddingsOut(cc).Send(std::move(results));
} else { } else {
kEmbeddingsOut(cc).Send(kEmbeddingsIn(cc)); auto result = kEmbeddingsIn(cc).Get();
result.set_timestamp_ms(cc->InputTimestamp().Value() / 1000);
kEmbeddingsOut(cc).Send(result);
} }
RET_CHECK(cached_embeddings_.empty()); RET_CHECK(cached_embeddings_.empty());
return absl::OkStatus(); return absl::OkStatus();

View File

@ -120,7 +120,7 @@ class EmbeddingAggregationCalculatorTest : public tflite_shims::testing::Test {
CalculatorGraph calculator_graph_; CalculatorGraph calculator_graph_;
}; };
TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) { TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutAggregation) {
EmbeddingResult embedding = ParseTextProtoOrDie<EmbeddingResult>( EmbeddingResult embedding = ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings { head_index: 0 })pb"); R"pb(embeddings { head_index: 0 })pb");
@ -129,10 +129,12 @@ TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) {
MP_ASSERT_OK(Send(embedding)); MP_ASSERT_OK(Send(embedding));
MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult<EmbeddingResult>(poller)); MP_ASSERT_OK_AND_ASSIGN(auto result, GetResult<EmbeddingResult>(poller));
EXPECT_THAT(result, EqualsProto(embedding)); EXPECT_THAT(result, EqualsProto(ParseTextProtoOrDie<EmbeddingResult>(
R"pb(timestamp_ms: 0
embeddings { head_index: 0 })pb")));
} }
TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithTimestamps) { TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithAggregation) {
MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true)); MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(/*connect_timestamps=*/true));
MP_ASSERT_OK(Send(ParseTextProtoOrDie<EmbeddingResult>(R"pb(embeddings { MP_ASSERT_OK(Send(ParseTextProtoOrDie<EmbeddingResult>(R"pb(embeddings {
head_index: 0 head_index: 0

View File

@ -246,7 +246,7 @@ class PostprocessingTest : public tflite_shims::testing::Test {
absl::make_unique<std::vector<Tensor>>(); absl::make_unique<std::vector<Tensor>>();
}; };
TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) { TEST_F(PostprocessingTest, SucceedsWithoutAggregation) {
// Build graph. // Build graph.
proto::EmbedderOptions options; proto::EmbedderOptions options;
MP_ASSERT_OK_AND_ASSIGN(auto poller, MP_ASSERT_OK_AND_ASSIGN(auto poller,
@ -261,7 +261,8 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
MP_ASSERT_OK_AND_ASSIGN(auto results, GetResult<EmbeddingResult>(poller)); MP_ASSERT_OK_AND_ASSIGN(auto results, GetResult<EmbeddingResult>(poller));
// Validate results. // Validate results.
EXPECT_FALSE(results.has_timestamp_ms()); EXPECT_TRUE(results.has_timestamp_ms());
EXPECT_EQ(results.timestamp_ms(), 0);
EXPECT_EQ(results.embeddings_size(), 1); EXPECT_EQ(results.embeddings_size(), 1);
EXPECT_EQ(results.embeddings(0).head_index(), 0); EXPECT_EQ(results.embeddings(0).head_index(), 0);
EXPECT_EQ(results.embeddings(0).head_name(), "feature"); EXPECT_EQ(results.embeddings(0).head_name(), "feature");
@ -273,7 +274,7 @@ TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
} }
} }
TEST_F(PostprocessingTest, SucceedsWithTimestamps) { TEST_F(PostprocessingTest, SucceedsWithAggregation) {
// Build graph. // Build graph.
proto::EmbedderOptions options; proto::EmbedderOptions options;
MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(kMobileNetV3Embedder, options, MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(kMobileNetV3Embedder, options,