components/segmentation_platform/internal/metadata/metadata_writer.cc - chromium/src.git - Git at Google

 // Copyright 2022 The Chromium Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 #ifdef UNSAFE_BUFFERS_BUILD
 // TODO(crbug.com/40285824): Remove this and convert code to safer constructs.
 #pragma allow_unsafe_buffers
 #endif

 #include "components/segmentation_platform/internal/metadata/metadata_writer.h"

 #include <cstddef>
 #include <optional>
 #include <vector>

 #include "base/check.h"
 #include "base/metrics/metrics_hashes.h"
 #include "base/strings/strcat.h"
 #include "components/segmentation_platform/public/constants.h"
 #include "components/segmentation_platform/public/proto/model_metadata.pb.h"

 namespace segmentation_platform {

 namespace {

 void FillCustomInput(const MetadataWriter::CustomInput feature,
                      proto::CustomInput& input) {
   input.set_tensor_length(feature.tensor_length);
   input.set_fill_policy(feature.fill_policy);
   for (size_t i = 0; i < feature.default_values_size; ++i) {
     input.add_default_value(feature.default_values[i]);
   }
   if (feature.name) {
     input.set_name(feature.name);
   }

   for (size_t i = 0; i < feature.arg_size; ++i) {
     (*input.mutable_additional_args())[feature.arg[i].first] =
         std::string(feature.arg[i].second);
   }
 }

 template <typename StringVector>
 void PopulateMultiClassClassifier(
     proto::Predictor::MultiClassClassifier* multi_class_classifier,
     const StringVector& class_labels,
     int top_k_outputs) {
   multi_class_classifier->set_top_k_outputs(top_k_outputs);
   for (const auto& class_label : class_labels) {
     multi_class_classifier->mutable_class_labels()->Add(
         std::string(class_label));
   }
 }

 }  // namespace

 MetadataWriter::MetadataWriter(proto::SegmentationModelMetadata* metadata)
     : metadata_(metadata) {}
 MetadataWriter::~MetadataWriter() = default;

 void MetadataWriter::AddUmaFeatures(const UMAFeature features[],
                                     size_t features_size,
                                     bool is_output) {
   for (size_t i = 0; i < features_size; i++) {
     const auto& feature = features[i];
     proto::UMAFeature* uma_feature;
     if (is_output) {
       auto* training_output =
           metadata_->mutable_training_outputs()->add_outputs();
       uma_feature =
           training_output->mutable_uma_output()->mutable_uma_feature();
     } else {
       auto* input_feature = metadata_->add_input_features();
       uma_feature = input_feature->mutable_uma_feature();
     }
     uma_feature->set_type(feature.signal_type);
     uma_feature->set_name(feature.name);
     uma_feature->set_name_hash(base::HashMetricName(feature.name));
     uma_feature->set_bucket_count(feature.bucket_count);
     uma_feature->set_tensor_length(feature.tensor_length);
     uma_feature->set_aggregation(feature.aggregation);

     for (size_t j = 0; j < feature.enum_ids_size; j++) {
       uma_feature->add_enum_ids(feature.accepted_enum_ids[j]);
     }

     for (size_t j = 0; j < feature.default_values_size; j++) {
       uma_feature->add_default_values(feature.default_values[j]);
     }
   }
 }

 proto::SqlFeature* MetadataWriter::AddSqlFeature(const SqlFeature& feature) {
   proto::SqlFeature* proto =
       metadata_->add_input_features()->mutable_sql_feature();
   proto->set_sql(feature.sql);
   for (size_t ev = 0; ev < feature.events_size; ++ev) {
     const auto& event = feature.events[ev];
     auto* ukm_event = proto->mutable_signal_filter()->add_ukm_events();
     ukm_event->set_event_hash(event.event_hash.GetUnsafeValue());
     for (size_t m = 0; m < event.metrics_size; ++m) {
       ukm_event->mutable_metric_hash_filter()->Add(
           event.metrics[m].GetUnsafeValue());
     }
   }
   return proto;
 }

 proto::SqlFeature* MetadataWriter::AddSqlFeature(
     const SqlFeature& feature,
     const BindValues& bind_values) {
   auto* proto = AddSqlFeature(feature);

   unsigned index = 0;
   for (const auto& it : bind_values) {
     auto* value = proto->add_bind_values();
     for (unsigned i = index; i < index + it.second.tensor_length; ++i) {
       value->add_bind_field_index(i);
     }
     index += it.second.tensor_length;
     value->set_param_type(it.first);
     FillCustomInput(it.second, *value->mutable_value());
   }
   return proto;
 }

 proto::CustomInput* MetadataWriter::AddCustomInput(const CustomInput& feature) {
   proto::CustomInput* proto =
       metadata_->add_input_features()->mutable_custom_input();
   FillCustomInput(feature, *proto);
   return proto;
 }

 void MetadataWriter::AddDiscreteMappingEntries(
     const std::string& key,
     const std::pair<float, int>* mappings,
     size_t mappings_size) {
   auto* discrete_mappings = metadata_->mutable_discrete_mappings();
   for (size_t i = 0; i < mappings_size; i++) {
     auto* discrete_mapping_entry = (*discrete_mappings)[key].add_entries();
     discrete_mapping_entry->set_min_result(mappings[i].first);
     discrete_mapping_entry->set_rank(mappings[i].second);
   }
 }

 void MetadataWriter::AddBooleanSegmentDiscreteMapping(const std::string& key) {
   const int selected_rank = 1;
   const float model_score = 1;
   const std::pair<float, int> mappings[]{{model_score, selected_rank}};
   AddDiscreteMappingEntries(key, mappings, 1);
 }

 void MetadataWriter::AddBooleanSegmentDiscreteMappingWithSubsegments(
     const std::string& key,
     float threshold,
     int max_value) {
   DCHECK_GT(threshold, 0);
   // Should record at least 2 subsegments.
   DCHECK_GT(max_value, 1);
   const int selected_rank = 1;
   const std::pair<float, int> mappings[]{{threshold, selected_rank}};
   AddDiscreteMappingEntries(key, mappings, 1);

   std::vector<std::pair<float, int>> subsegment_mapping;
   for (int i = 1; i <= max_value; ++i) {
     subsegment_mapping.emplace_back(i, i);
   }
   AddDiscreteMappingEntries(
       base::StrCat({key, kSubsegmentDiscreteMappingSuffix}),
       subsegment_mapping.data(), subsegment_mapping.size());
 }

 void MetadataWriter::SetSegmentationMetadataConfig(
     proto::TimeUnit time_unit,
     uint64_t bucket_duration,
     int64_t signal_storage_length,
     int64_t min_signal_collection_length,
     int64_t result_time_to_live) {
   metadata_->set_time_unit(time_unit);
   metadata_->set_bucket_duration(bucket_duration);
   metadata_->set_signal_storage_length(signal_storage_length);
   metadata_->set_min_signal_collection_length(min_signal_collection_length);
   metadata_->set_result_time_to_live(result_time_to_live);
 }

 void MetadataWriter::SetDefaultSegmentationMetadataConfig(
     int min_signal_collection_length_days,
     int signal_storage_length_days) {
   SetSegmentationMetadataConfig(proto::TimeUnit::DAY, /*bucket_duration=*/1,
                                 signal_storage_length_days,
                                 min_signal_collection_length_days,
                                 /*result_time_to_live=*/1);
 }

 void MetadataWriter::AddOutputConfigForBinaryClassifier(
     float threshold,
     const std::string& positive_label,
     const std::string& negative_label) {
   proto::Predictor::BinaryClassifier* binary_classifier =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_binary_classifier();

   binary_classifier->set_threshold(threshold);
   binary_classifier->set_positive_label(positive_label);
   binary_classifier->set_negative_label(negative_label);
 }

 void MetadataWriter::SetIgnorePreviousModelTTLInOutputConfig() {
   metadata_->mutable_output_config()->set_ignore_previous_model_ttl(true);
 }

 void MetadataWriter::AddOutputConfigForMultiClassClassifier(
     base::span<const char* const> class_labels,
     int top_k_outputs,
     std::optional<float> threshold) {
   proto::Predictor::MultiClassClassifier* multi_class_classifier =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_multi_class_classifier();

   PopulateMultiClassClassifier(multi_class_classifier, class_labels,
                                top_k_outputs);
   if (threshold.has_value()) {
     multi_class_classifier->set_threshold(threshold.value());
   }
 }

 void MetadataWriter::AddOutputConfigForMultiClassClassifier(
     const std::vector<std::string>& class_labels,
     int top_k_outputs,
     std::optional<float> threshold) {
   proto::Predictor::MultiClassClassifier* multi_class_classifier =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_multi_class_classifier();

   PopulateMultiClassClassifier(multi_class_classifier, class_labels,
                                top_k_outputs);
   if (threshold.has_value()) {
     multi_class_classifier->set_threshold(threshold.value());
   }
 }

 void MetadataWriter::AddOutputConfigForMultiClassClassifier(
     base::span<const char* const> class_labels,
     int top_k_outputs,
     const base::span<float> per_class_thresholds) {
   CHECK_EQ(class_labels.size(), per_class_thresholds.size());
   proto::Predictor::MultiClassClassifier* multi_class_classifier =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_multi_class_classifier();

   PopulateMultiClassClassifier(multi_class_classifier, class_labels,
                                top_k_outputs);

   for (float per_class_threshold : per_class_thresholds) {
     multi_class_classifier->add_class_thresholds(per_class_threshold);
   }
 }

 void MetadataWriter::AddOutputConfigForBinnedClassifier(
     const std::vector<std::pair<float, std::string>>& bins,
     std::string underflow_label) {
   proto::Predictor::BinnedClassifier* binned_classifier =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_binned_classifier();

   binned_classifier->set_underflow_label(underflow_label);
   for (const std::pair<float, std::string>& bin : bins) {
     proto::Predictor::BinnedClassifier::Bin* current_bin =
         binned_classifier->add_bins();
     current_bin->set_min_range(bin.first);
     current_bin->set_label(bin.second);
   }
 }

 void MetadataWriter::AddOutputConfigForGenericPredictor(
     const std::vector<std::string>& labels) {
   proto::Predictor::GenericPredictor* generic_predictor =
       metadata_->mutable_output_config()
           ->mutable_predictor()
           ->mutable_generic_predictor();
   generic_predictor->mutable_output_labels()->Assign(labels.begin(),
                                                      labels.end());
 }

 void MetadataWriter::AddPredictedResultTTLInOutputConfig(
     std::vector<std::pair<std::string, std::int64_t>> top_label_to_ttl_list,
     int64_t default_ttl,
     proto::TimeUnit time_unit) {
   proto::PredictedResultTTL* predicted_result_ttl =
       metadata_->mutable_output_config()->mutable_predicted_result_ttl();
   predicted_result_ttl->set_time_unit(time_unit);
   predicted_result_ttl->set_default_ttl(default_ttl);
   auto* top_label_to_ttl_map =
       predicted_result_ttl->mutable_top_label_to_ttl_map();
   for (const std::pair<std::string, int64_t>& label_to_ttl :
        top_label_to_ttl_list) {
     (*top_label_to_ttl_map)[label_to_ttl.first] = label_to_ttl.second;
   }
 }

 void MetadataWriter::AddDelayTrigger(uint64_t delay_sec) {
   auto* config =
       metadata_->mutable_training_outputs()->mutable_trigger_config();
   auto* trigger = config->add_observation_trigger();
   trigger->set_delay_sec(delay_sec);
   config->set_decision_type(proto::TrainingOutputs::TriggerConfig::ONDEMAND);
 }

 void MetadataWriter::AddFromInputContext(const char* custom_input_name,
                                          const char* additional_args_name) {
   proto::CustomInput* custom_input = AddCustomInput(MetadataWriter::CustomInput{
       .tensor_length = 1,
       .fill_policy = proto::CustomInput::FILL_FROM_INPUT_CONTEXT,
       .name = custom_input_name});
   (*custom_input->mutable_additional_args())["name"] = additional_args_name;
 }

 }  // namespace segmentation_platform
	// Copyright 2022 The Chromium Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	#ifdef UNSAFE_BUFFERS_BUILD
	// TODO(crbug.com/40285824): Remove this and convert code to safer constructs.
	#pragma allow_unsafe_buffers
	#endif

	#include "components/segmentation_platform/internal/metadata/metadata_writer.h"

	#include <cstddef>
	#include <optional>
	#include <vector>

	#include "base/check.h"
	#include "base/metrics/metrics_hashes.h"
	#include "base/strings/strcat.h"
	#include "components/segmentation_platform/public/constants.h"
	#include "components/segmentation_platform/public/proto/model_metadata.pb.h"

	namespace segmentation_platform {

	namespace {

	void FillCustomInput(const MetadataWriter::CustomInput feature,
	proto::CustomInput& input) {
	input.set_tensor_length(feature.tensor_length);
	input.set_fill_policy(feature.fill_policy);
	for (size_t i = 0; i < feature.default_values_size; ++i) {
	input.add_default_value(feature.default_values[i]);
	}
	if (feature.name) {
	input.set_name(feature.name);
	}

	for (size_t i = 0; i < feature.arg_size; ++i) {
	(*input.mutable_additional_args())[feature.arg[i].first] =
	std::string(feature.arg[i].second);
	}
	}

	template <typename StringVector>
	void PopulateMultiClassClassifier(
	proto::Predictor::MultiClassClassifier* multi_class_classifier,
	const StringVector& class_labels,
	int top_k_outputs) {
	multi_class_classifier->set_top_k_outputs(top_k_outputs);
	for (const auto& class_label : class_labels) {
	multi_class_classifier->mutable_class_labels()->Add(
	std::string(class_label));
	}
	}

	} // namespace

	MetadataWriter::MetadataWriter(proto::SegmentationModelMetadata* metadata)
	: metadata_(metadata) {}
	MetadataWriter::~MetadataWriter() = default;

	void MetadataWriter::AddUmaFeatures(const UMAFeature features[],
	size_t features_size,
	bool is_output) {
	for (size_t i = 0; i < features_size; i++) {
	const auto& feature = features[i];
	proto::UMAFeature* uma_feature;
	if (is_output) {
	auto* training_output =
	metadata_->mutable_training_outputs()->add_outputs();
	uma_feature =
	training_output->mutable_uma_output()->mutable_uma_feature();
	} else {
	auto* input_feature = metadata_->add_input_features();
	uma_feature = input_feature->mutable_uma_feature();
	}
	uma_feature->set_type(feature.signal_type);
	uma_feature->set_name(feature.name);
	uma_feature->set_name_hash(base::HashMetricName(feature.name));
	uma_feature->set_bucket_count(feature.bucket_count);
	uma_feature->set_tensor_length(feature.tensor_length);
	uma_feature->set_aggregation(feature.aggregation);

	for (size_t j = 0; j < feature.enum_ids_size; j++) {
	uma_feature->add_enum_ids(feature.accepted_enum_ids[j]);
	}

	for (size_t j = 0; j < feature.default_values_size; j++) {
	uma_feature->add_default_values(feature.default_values[j]);
	}
	}
	}

	proto::SqlFeature* MetadataWriter::AddSqlFeature(const SqlFeature& feature) {
	proto::SqlFeature* proto =
	metadata_->add_input_features()->mutable_sql_feature();
	proto->set_sql(feature.sql);
	for (size_t ev = 0; ev < feature.events_size; ++ev) {
	const auto& event = feature.events[ev];
	auto* ukm_event = proto->mutable_signal_filter()->add_ukm_events();
	ukm_event->set_event_hash(event.event_hash.GetUnsafeValue());
	for (size_t m = 0; m < event.metrics_size; ++m) {
	ukm_event->mutable_metric_hash_filter()->Add(
	event.metrics[m].GetUnsafeValue());
	}
	}
	return proto;
	}

	proto::SqlFeature* MetadataWriter::AddSqlFeature(
	const SqlFeature& feature,
	const BindValues& bind_values) {
	auto* proto = AddSqlFeature(feature);

	unsigned index = 0;
	for (const auto& it : bind_values) {
	auto* value = proto->add_bind_values();
	for (unsigned i = index; i < index + it.second.tensor_length; ++i) {
	value->add_bind_field_index(i);
	}
	index += it.second.tensor_length;
	value->set_param_type(it.first);
	FillCustomInput(it.second, *value->mutable_value());
	}
	return proto;
	}

	proto::CustomInput* MetadataWriter::AddCustomInput(const CustomInput& feature) {
	proto::CustomInput* proto =
	metadata_->add_input_features()->mutable_custom_input();
	FillCustomInput(feature, *proto);
	return proto;
	}

	void MetadataWriter::AddDiscreteMappingEntries(
	const std::string& key,
	const std::pair<float, int>* mappings,
	size_t mappings_size) {
	auto* discrete_mappings = metadata_->mutable_discrete_mappings();
	for (size_t i = 0; i < mappings_size; i++) {
	auto* discrete_mapping_entry = (*discrete_mappings)[key].add_entries();
	discrete_mapping_entry->set_min_result(mappings[i].first);
	discrete_mapping_entry->set_rank(mappings[i].second);
	}
	}

	void MetadataWriter::AddBooleanSegmentDiscreteMapping(const std::string& key) {
	const int selected_rank = 1;
	const float model_score = 1;
	const std::pair<float, int> mappings[]{{model_score, selected_rank}};
	AddDiscreteMappingEntries(key, mappings, 1);
	}

	void MetadataWriter::AddBooleanSegmentDiscreteMappingWithSubsegments(
	const std::string& key,
	float threshold,
	int max_value) {
	DCHECK_GT(threshold, 0);
	// Should record at least 2 subsegments.
	DCHECK_GT(max_value, 1);
	const int selected_rank = 1;
	const std::pair<float, int> mappings[]{{threshold, selected_rank}};
	AddDiscreteMappingEntries(key, mappings, 1);

	std::vector<std::pair<float, int>> subsegment_mapping;
	for (int i = 1; i <= max_value; ++i) {
	subsegment_mapping.emplace_back(i, i);
	}
	AddDiscreteMappingEntries(
	base::StrCat({key, kSubsegmentDiscreteMappingSuffix}),
	subsegment_mapping.data(), subsegment_mapping.size());
	}

	void MetadataWriter::SetSegmentationMetadataConfig(
	proto::TimeUnit time_unit,
	uint64_t bucket_duration,
	int64_t signal_storage_length,
	int64_t min_signal_collection_length,
	int64_t result_time_to_live) {
	metadata_->set_time_unit(time_unit);
	metadata_->set_bucket_duration(bucket_duration);
	metadata_->set_signal_storage_length(signal_storage_length);
	metadata_->set_min_signal_collection_length(min_signal_collection_length);
	metadata_->set_result_time_to_live(result_time_to_live);
	}

	void MetadataWriter::SetDefaultSegmentationMetadataConfig(
	int min_signal_collection_length_days,
	int signal_storage_length_days) {
	SetSegmentationMetadataConfig(proto::TimeUnit::DAY, /bucket_duration=/1,
	signal_storage_length_days,
	min_signal_collection_length_days,
	/result_time_to_live=/1);
	}

	void MetadataWriter::AddOutputConfigForBinaryClassifier(
	float threshold,
	const std::string& positive_label,
	const std::string& negative_label) {
	proto::Predictor::BinaryClassifier* binary_classifier =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_binary_classifier();

	binary_classifier->set_threshold(threshold);
	binary_classifier->set_positive_label(positive_label);
	binary_classifier->set_negative_label(negative_label);
	}

	void MetadataWriter::SetIgnorePreviousModelTTLInOutputConfig() {
	metadata_->mutable_output_config()->set_ignore_previous_model_ttl(true);
	}

	void MetadataWriter::AddOutputConfigForMultiClassClassifier(
	base::span<const char* const> class_labels,
	int top_k_outputs,
	std::optional<float> threshold) {
	proto::Predictor::MultiClassClassifier* multi_class_classifier =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_multi_class_classifier();

	PopulateMultiClassClassifier(multi_class_classifier, class_labels,
	top_k_outputs);
	if (threshold.has_value()) {
	multi_class_classifier->set_threshold(threshold.value());
	}
	}

	void MetadataWriter::AddOutputConfigForMultiClassClassifier(
	const std::vector<std::string>& class_labels,
	int top_k_outputs,
	std::optional<float> threshold) {
	proto::Predictor::MultiClassClassifier* multi_class_classifier =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_multi_class_classifier();

	PopulateMultiClassClassifier(multi_class_classifier, class_labels,
	top_k_outputs);
	if (threshold.has_value()) {
	multi_class_classifier->set_threshold(threshold.value());
	}
	}

	void MetadataWriter::AddOutputConfigForMultiClassClassifier(
	base::span<const char* const> class_labels,
	int top_k_outputs,
	const base::span<float> per_class_thresholds) {
	CHECK_EQ(class_labels.size(), per_class_thresholds.size());
	proto::Predictor::MultiClassClassifier* multi_class_classifier =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_multi_class_classifier();

	PopulateMultiClassClassifier(multi_class_classifier, class_labels,
	top_k_outputs);

	for (float per_class_threshold : per_class_thresholds) {
	multi_class_classifier->add_class_thresholds(per_class_threshold);
	}
	}

	void MetadataWriter::AddOutputConfigForBinnedClassifier(
	const std::vector<std::pair<float, std::string>>& bins,
	std::string underflow_label) {
	proto::Predictor::BinnedClassifier* binned_classifier =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_binned_classifier();

	binned_classifier->set_underflow_label(underflow_label);
	for (const std::pair<float, std::string>& bin : bins) {
	proto::Predictor::BinnedClassifier::Bin* current_bin =
	binned_classifier->add_bins();
	current_bin->set_min_range(bin.first);
	current_bin->set_label(bin.second);
	}
	}

	void MetadataWriter::AddOutputConfigForGenericPredictor(
	const std::vector<std::string>& labels) {
	proto::Predictor::GenericPredictor* generic_predictor =
	metadata_->mutable_output_config()
	->mutable_predictor()
	->mutable_generic_predictor();
	generic_predictor->mutable_output_labels()->Assign(labels.begin(),
	labels.end());
	}

	void MetadataWriter::AddPredictedResultTTLInOutputConfig(
	std::vector<std::pair<std::string, std::int64_t>> top_label_to_ttl_list,
	int64_t default_ttl,
	proto::TimeUnit time_unit) {
	proto::PredictedResultTTL* predicted_result_ttl =
	metadata_->mutable_output_config()->mutable_predicted_result_ttl();
	predicted_result_ttl->set_time_unit(time_unit);
	predicted_result_ttl->set_default_ttl(default_ttl);
	auto* top_label_to_ttl_map =
	predicted_result_ttl->mutable_top_label_to_ttl_map();
	for (const std::pair<std::string, int64_t>& label_to_ttl :
	top_label_to_ttl_list) {
	(*top_label_to_ttl_map)[label_to_ttl.first] = label_to_ttl.second;
	}
	}

	void MetadataWriter::AddDelayTrigger(uint64_t delay_sec) {
	auto* config =
	metadata_->mutable_training_outputs()->mutable_trigger_config();
	auto* trigger = config->add_observation_trigger();
	trigger->set_delay_sec(delay_sec);
	config->set_decision_type(proto::TrainingOutputs::TriggerConfig::ONDEMAND);
	}

	void MetadataWriter::AddFromInputContext(const char* custom_input_name,
	const char* additional_args_name) {
	proto::CustomInput* custom_input = AddCustomInput(MetadataWriter::CustomInput{
	.tensor_length = 1,
	.fill_policy = proto::CustomInput::FILL_FROM_INPUT_CONTEXT,
	.name = custom_input_name});
	(*custom_input->mutable_additional_args())["name"] = additional_args_name;
	}

	} // namespace segmentation_platform