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));
} 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());
return absl::OkStatus();

View File

@ -120,7 +120,7 @@ class EmbeddingAggregationCalculatorTest : public tflite_shims::testing::Test {
CalculatorGraph calculator_graph_;
};
TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) {
TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutAggregation) {
EmbeddingResult embedding = ParseTextProtoOrDie<EmbeddingResult>(
R"pb(embeddings { head_index: 0 })pb");
@ -129,10 +129,12 @@ TEST_F(EmbeddingAggregationCalculatorTest, SucceedsWithoutTimestamps) {
MP_ASSERT_OK(Send(embedding));
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(Send(ParseTextProtoOrDie<EmbeddingResult>(R"pb(embeddings {
head_index: 0

View File

@ -246,7 +246,7 @@ class PostprocessingTest : public tflite_shims::testing::Test {
absl::make_unique<std::vector<Tensor>>();
};
TEST_F(PostprocessingTest, SucceedsWithoutTimestamps) {
TEST_F(PostprocessingTest, SucceedsWithoutAggregation) {
// Build graph.
proto::EmbedderOptions options;
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));
// 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(0).head_index(), 0);
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.
proto::EmbedderOptions options;
MP_ASSERT_OK_AND_ASSIGN(auto poller, BuildGraph(kMobileNetV3Embedder, options,