Fix/Proxy: record search boost

koursaros-ai · Dec 4, 2019 · 7a120ce · 7a120ce
1 parent 14446e6
commit 7a120ce
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/nboost/proxy.py b/nboost/proxy.py
@@ -144,13 +144,11 @@ def record_topk_and_choices(self, topk: int = None, choices: list = None):
 
     @stats.vars_context
     def record_mrrs(self, upstream_mrr: float = None, model_mrr: float = None):
-        """Add the upstream mrr and model mrr to the stats"""
+        """Add the upstream mrr, model mrr, and search boost to the stats"""
         with suppress(ZeroDivisionError):
-            self.record_search_boost(search_boost=model_mrr / upstream_mrr)
-
-    @stats.vars_context
-    def record_search_boost(self, search_boost: float = None):
-        """Record the quotient of model mrr and upstream mrr"""
+            var = self.stats.record['vars']
+            var['search_boost'] = {
+                'avg': var['model_mrr']['avg'] / var['upstream_mrr']['avg']}
 
     @stats.time_context
     def server_connect(self, server_socket: socket.socket) -> None: