components/browsing_topics/annotator.h

   1 // Copyright 2023 The Chromium Authors
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #ifndef COMPONENTS_BROWSING_TOPICS_ANNOTATOR_H_
   6 #define COMPONENTS_BROWSING_TOPICS_ANNOTATOR_H_
   7
   8 #include <string>
   9 #include <vector>
  10
  11 #include "base/functional/callback.h"
  12 #include "components/optimization_guide/core/bert_model_handler.h"
  13 #include "third_party/abseil-cpp/absl/types/optional.h"
  14
  15 namespace browsing_topics {
  16
  17 // An annotation for a single input to the Annotator.
  18 struct Annotation {
  19   explicit Annotation(const std::string& input);
  20   Annotation(const Annotation& other);
  21   ~Annotation();
  22
  23   std::string input;
  24   std::vector<int32_t> topics;
  25 };
  26
  27 // An ML-based annotator for using in Browsing Topics. Models are delivered by
  28 // the OptimizationGuide service and processed using the TFLite Support Library.
  29 // Support for a Google-provided lookup table is also included that
  30 // short-circuits the ML model, allowing for better performance on popular pages
  31 // as well as human labels when the ML model is inaccurate.
  32 class Annotator {
  33  public:
  34   using BatchAnnotationCallback =
  35       base::OnceCallback<void(const std::vector<Annotation>&)>;
  36
  37   Annotator() = default;
  38   virtual ~Annotator() = default;
  39
  40   // This is the main entry point for callers.
  41   virtual void BatchAnnotate(BatchAnnotationCallback callback,
  42                              const std::vector<std::string>& inputs) = 0;
  43
  44   // Runs |callback| when the model is ready to execute. If the model is ready
  45   // now, the callback is run immediately.
  46   virtual void NotifyWhenModelAvailable(base::OnceClosure callback) = 0;
  47
  48   // Returns the model info if the model file is available.
  49   virtual absl::optional<optimization_guide::ModelInfo>
  50   GetBrowsingTopicsModelInfo() const = 0;
  51 };
  52
  53 }  // namespace browsing_topics
  54
  55 #endif  // COMPONENTS_BROWSING_TOPICS_ANNOTATOR_H_