[Segmentation] Allow tensor uploading thru metadata
Change-Id: I6f8530cc92623d35fd639eb6060fae0ba148337f
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3846605
Reviewed-by: Siddhartha S <[email protected]>
Commit-Queue: Min Qin <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1038883}
diff --git a/components/segmentation_platform/internal/segmentation_ukm_helper.cc b/components/segmentation_platform/internal/segmentation_ukm_helper.cc
index ab51818..324e752f 100644
--- a/components/segmentation_platform/internal/segmentation_ukm_helper.cc
+++ b/components/segmentation_platform/internal/segmentation_ukm_helper.cc
@@ -115,8 +115,9 @@
// Add inputs to ukm message.
if (!AddInputsToUkm(&execution_result, segment_id, model_version,
- input_tensor))
+ input_tensor)) {
return ukm::kInvalidSourceId;
+ }
// TODO(xingliu): Also record continuous outputs for model execution.
execution_result.SetPredictionResult(FloatToInt64(result))
@@ -162,9 +163,6 @@
SegmentId segment_id,
int64_t model_version,
const std::vector<float>& input_tensor) {
- if (!allowed_segment_ids_.contains(static_cast<int>(segment_id)))
- return false;
-
if (input_tensor.size() > ARRAY_SIZE(kSegmentationUkmInputMethods)) {
// Don't record UKM if there are too many tensors.
stats::RecordTooManyInputTensors(input_tensor.size());
@@ -202,6 +200,17 @@
return true;
}
+bool SegmentationUkmHelper::CanUploadTensors(
+ const proto::SegmentInfo& segment_info) const {
+ if (!base::FeatureList::IsEnabled(
+ features::kSegmentationStructuredMetricsFeature)) {
+ return false;
+ }
+ return segment_info.model_metadata().upload_tensors() ||
+ allowed_segment_ids_.contains(
+ static_cast<int>(segment_info.segment_id()));
+}
+
// static
int64_t SegmentationUkmHelper::FloatToInt64(float f) {
// Encode the float number in IEEE754 double precision.
