1 // Copyright 2018 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 
5 #ifndef CHROME_BROWSER_RESOURCE_COORDINATOR_TAB_RANKER_TAB_SCORE_PREDICTOR_H_
6 #define CHROME_BROWSER_RESOURCE_COORDINATOR_TAB_RANKER_TAB_SCORE_PREDICTOR_H_
7 
8 #include <map>
9 #include <memory>
10 
11 #include "base/compiler_specific.h"
12 #include "base/macros.h"
13 #include "base/optional.h"
14 
15 namespace assist_ranker {
16 class ExamplePreprocessorConfig;
17 class RankerExample;
18 }  // namespace assist_ranker
19 
20 namespace tab_ranker {
21 
22 namespace tfnative_model {
23 struct FixedAllocations;
24 }  // namespace tfnative_model
25 
26 namespace pairwise_model {
27 struct FixedAllocations;
28 }  // namespace pairwise_model
29 
30 struct TabFeatures;
31 
32 // These values are persisted to logs. Entries should not be renumbered and
33 // numeric values should never be reused.
34 enum class TabRankerResult {
35   kSuccess = 0,
36   kPreprocessorInitializationFailed = 1,
37   kPreprocessorOtherError = 2,
38   kUnrecognizableScorer = 3,
39   kMaxValue = kUnrecognizableScorer
40 };
41 
42 // Makes predictions using the tab reactivation DNN classifier. Background tabs
43 // are scored based on how likely they are to be reactivated.
44 class TabScorePredictor {
45  public:
46   enum ScorerType {
47     kMRUScorer = 0,
48     kMLScorer = 1,
49     kPairwiseScorer = 2,
50     kFrecencyScorer = 3,
51     kMaxValue = kFrecencyScorer
52   };
53   TabScorePredictor();
54   ~TabScorePredictor();
55 
56   // Scores the tab using the tab reactivation model. A higher score indicates
57   // the tab is more likely to be reactivated than a lower score. A lower score
58   // indicates the tab is more likely to be closed.
59   TabRankerResult ScoreTab(const TabFeatures& tab,
60                            float* score) WARN_UNUSED_RESULT;
61 
62   // Scores multiple tabs.
63   // Input is a map from an id (lifecycle_unit id) to the TabFeatures of that
64   // tab.
65   // Returns a map from an id to its predicted reactivation score.
66   // If the scoring fails at any step, it will set
67   // std::numeric_limits<float>::max() as the reactivation score for that tab.
68   std::map<int32_t, float> ScoreTabs(
69       const std::map<int32_t, base::Optional<TabFeatures>>& tabs);
70 
71  private:
72   friend class ScoreTabsWithPairwiseScorerTest;
73 
74   // Loads the preprocessor config if not already loaded.
75   void LazyInitialize();
76 
77   // Calculates reactivation score of a single tab with mru feature.
78   TabRankerResult ScoreTabWithMRUScorer(const TabFeatures& tab, float* score);
79   // Calculates reactivation score of a single tab with ml model.
80   TabRankerResult ScoreTabWithMLScorer(const TabFeatures& tab, float* score);
81   // Preprocess and inferences on the |example|.
82   TabRankerResult PredictWithPreprocess(assist_ranker::RankerExample* example,
83                                         float* score);
84   // Calculates the relative reaction score between tab1 and tab2.
85   // For pairwise model, the ml model is applied to the pair(tab1, tab2).
86   // For non-pairwise model, the score is the difference of reactivation
87   // scores on these two tabs.
88   TabRankerResult ScoreTabsPairs(const TabFeatures& tab1,
89                                  const TabFeatures& tab2,
90                                  float* score);
91   std::map<int32_t, float> ScoreTabsWithPairwiseScorer(
92       const std::map<int32_t, base::Optional<TabFeatures>>& tabs);
93   TabRankerResult ScoreTabWithFrecencyScorer(const TabFeatures& tab,
94                                              float* score);
95 
96   std::unique_ptr<assist_ranker::ExamplePreprocessorConfig>
97       preprocessor_config_;
98 
99   // Fixed-size working memory provided to the inferencing function. Lazy
100   // initialized once so it isn't reallocated for every inference.
101   std::unique_ptr<tfnative_model::FixedAllocations> tfnative_alloc_;
102   std::unique_ptr<pairwise_model::FixedAllocations> pairwise_alloc_;
103 
104   const float discard_count_penalty_ = 0.0f;
105   const float mru_scorer_penalty_ = 1.0f;
106   const ScorerType type_ = kMLScorer;
107 
108   DISALLOW_COPY_AND_ASSIGN(TabScorePredictor);
109 };
110 
111 }  // namespace tab_ranker
112 
113 #endif  // CHROME_BROWSER_RESOURCE_COORDINATOR_TAB_RANKER_TAB_SCORE_PREDICTOR_H_
114