page_content_annotations_common.cc 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211
  1. // Copyright 2021 The Chromium Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "components/optimization_guide/core/page_content_annotations_common.h"
  5. #include <algorithm>
  6. #include <ostream>
  7. #include "base/check_op.h"
  8. #include "base/json/json_writer.h"
  9. #include "base/strings/strcat.h"
  10. #include "base/strings/string_number_conversions.h"
  11. #include "base/strings/string_util.h"
  12. #include "base/strings/stringprintf.h"
  13. namespace optimization_guide {
  14. WeightedIdentifier::WeightedIdentifier(int32_t value, double weight)
  15. : value_(value), weight_(weight) {
  16. DCHECK_GE(weight_, 0.0);
  17. DCHECK_LE(weight_, 1.0);
  18. }
  19. WeightedIdentifier::WeightedIdentifier(const WeightedIdentifier&) = default;
  20. WeightedIdentifier::~WeightedIdentifier() = default;
  21. bool WeightedIdentifier::operator==(const WeightedIdentifier& other) const {
  22. constexpr double kWeightTolerance = 1e-6;
  23. return this->value_ == other.value_ &&
  24. std::abs(this->weight_ - other.weight_) <= kWeightTolerance;
  25. }
  26. std::string WeightedIdentifier::ToString() const {
  27. return base::StringPrintf("WeightedIdentifier{%d,%f}", value(), weight());
  28. }
  29. base::Value WeightedIdentifier::AsValue() const {
  30. base::Value::Dict wi;
  31. wi.Set("value", value());
  32. wi.Set("weight", weight());
  33. return base::Value(std::move(wi));
  34. }
  35. std::ostream& operator<<(std::ostream& stream, const WeightedIdentifier& ws) {
  36. stream << ws.ToString();
  37. return stream;
  38. }
  39. BatchAnnotationResult::BatchAnnotationResult() = default;
  40. BatchAnnotationResult::BatchAnnotationResult(const BatchAnnotationResult&) =
  41. default;
  42. BatchAnnotationResult::~BatchAnnotationResult() = default;
  43. bool BatchAnnotationResult::HasOutputForType() const {
  44. switch (type()) {
  45. case AnnotationType::kUnknown:
  46. return false;
  47. case AnnotationType::kPageTopics:
  48. return !!topics();
  49. case AnnotationType::kContentVisibility:
  50. return !!visibility_score();
  51. case AnnotationType::kPageEntities:
  52. return !!entities();
  53. }
  54. }
  55. base::Value BatchAnnotationResult::AsValue() const {
  56. base::Value::Dict result;
  57. result.Set("input", input());
  58. result.Set("type", AnnotationTypeToString(type()));
  59. if (topics()) {
  60. base::Value::List list;
  61. for (const auto& wi : *topics()) {
  62. list.Append(wi.AsValue());
  63. }
  64. result.Set("topics", std::move(list));
  65. }
  66. if (entities()) {
  67. base::Value::List list;
  68. for (const auto& md : *entities()) {
  69. list.Append(md.AsValue());
  70. }
  71. result.Set("entities", std::move(list));
  72. }
  73. if (visibility_score()) {
  74. result.Set("visibility_score", *visibility_score());
  75. }
  76. return base::Value(std::move(result));
  77. }
  78. std::string BatchAnnotationResult::ToJSON() const {
  79. std::string json;
  80. if (base::JSONWriter::Write(AsValue(), &json)) {
  81. return json;
  82. }
  83. return std::string();
  84. }
  85. std::string BatchAnnotationResult::ToString() const {
  86. std::string output = "nullopt";
  87. if (topics_) {
  88. std::vector<std::string> all_weighted_ids;
  89. for (const WeightedIdentifier& wi : *topics_) {
  90. all_weighted_ids.push_back(wi.ToString());
  91. }
  92. output = "{" + base::JoinString(all_weighted_ids, ",") + "}";
  93. } else if (entities_) {
  94. std::vector<std::string> all_entities;
  95. for (const ScoredEntityMetadata& md : *entities_) {
  96. all_entities.push_back(md.ToString());
  97. }
  98. output = "{" + base::JoinString(all_entities, ",") + "}";
  99. } else if (visibility_score_) {
  100. output = base::NumberToString(*visibility_score_);
  101. }
  102. return base::StringPrintf(
  103. "BatchAnnotationResult{"
  104. "\"%s\", "
  105. "type: %s, "
  106. "output: %s}",
  107. input_.c_str(), AnnotationTypeToString(type_).c_str(), output.c_str());
  108. }
  109. std::ostream& operator<<(std::ostream& stream,
  110. const BatchAnnotationResult& result) {
  111. stream << result.ToString();
  112. return stream;
  113. }
  114. // static
  115. BatchAnnotationResult BatchAnnotationResult::CreatePageTopicsResult(
  116. const std::string& input,
  117. absl::optional<std::vector<WeightedIdentifier>> topics) {
  118. BatchAnnotationResult result;
  119. result.input_ = input;
  120. result.topics_ = topics;
  121. result.type_ = AnnotationType::kPageTopics;
  122. // Always sort the result (if present) by the given score.
  123. if (result.topics_) {
  124. std::sort(result.topics_->begin(), result.topics_->end(),
  125. [](const WeightedIdentifier& a, const WeightedIdentifier& b) {
  126. return a.weight() < b.weight();
  127. });
  128. }
  129. return result;
  130. }
  131. // static
  132. BatchAnnotationResult BatchAnnotationResult::CreatePageEntitiesResult(
  133. const std::string& input,
  134. absl::optional<std::vector<ScoredEntityMetadata>> entities) {
  135. BatchAnnotationResult result;
  136. result.input_ = input;
  137. result.entities_ = entities;
  138. result.type_ = AnnotationType::kPageEntities;
  139. // Always sort the result (if present) by the given score.
  140. if (result.entities_) {
  141. std::sort(result.entities_->begin(), result.entities_->end(),
  142. [](const ScoredEntityMetadata& a, const ScoredEntityMetadata& b) {
  143. return a.score < b.score;
  144. });
  145. }
  146. return result;
  147. }
  148. // static
  149. BatchAnnotationResult BatchAnnotationResult::CreateContentVisibilityResult(
  150. const std::string& input,
  151. absl::optional<double> visibility_score) {
  152. BatchAnnotationResult result;
  153. result.input_ = input;
  154. result.visibility_score_ = visibility_score;
  155. result.type_ = AnnotationType::kContentVisibility;
  156. return result;
  157. }
  158. // static
  159. BatchAnnotationResult BatchAnnotationResult::CreateEmptyAnnotationsResult(
  160. const std::string& input) {
  161. BatchAnnotationResult result;
  162. result.input_ = input;
  163. return result;
  164. }
  165. bool BatchAnnotationResult::operator==(
  166. const BatchAnnotationResult& other) const {
  167. return this->input_ == other.input_ && this->type_ == other.type_ &&
  168. this->topics_ == other.topics_ && this->entities_ == other.entities_ &&
  169. this->visibility_score_ == other.visibility_score_;
  170. }
  171. std::vector<BatchAnnotationResult> CreateEmptyBatchAnnotationResults(
  172. const std::vector<std::string>& inputs) {
  173. std::vector<BatchAnnotationResult> results;
  174. results.reserve(inputs.size());
  175. for (const std::string& input : inputs) {
  176. results.emplace_back(
  177. BatchAnnotationResult::CreateEmptyAnnotationsResult(input));
  178. }
  179. return results;
  180. }
  181. } // namespace optimization_guide