[#37] Store ELO ratings in DB after calculation (#112)

* [#37] Store ELO ratings in DB after calculation

This change adds logic to store ELO ratings in the database after they are calculated.

Previously, we calculated and loaded the ratings without storing them. Now, we store them for future use.

This change doesn't affect current operations as we are not using the stored data yet. However, it sets the groundwork for future optimizations where we will use the stored ratings to avoid recalculating ELO scores for already calculated battles.

* update

* update

* review

* fix

* fix

* review

README.md

@@ -49,6 +49,9 @@ Get Involved: [Discuss and contribute on GitHub](https://github.com/yanolja/aren
 
    ```shell
    CREDENTIALS_PATH=<your crednetials path> \
+   RATINGS_COLLECTION=<your collection> \
+   SUMMARIZATIONS_COLLECTION=<your collection> \
+   TRANSLATIONS_COLLECTION=<your collection> \
    OPENAI_API_KEY=<your key> \
    ANTHROPIC_API_KEY=<your key> \
    MISTRAL_API_KEY=<your key> \

app.py

@@ -8,8 +8,8 @@ from firebase_admin import firestore
 import gradio as gr
 import lingua
 
+from db import db
 from leaderboard import build_leaderboard
-from leaderboard import db
 from leaderboard import SUPPORTED_LANGUAGES
 from model import check_models
 from model import supported_models
@@ -50,6 +50,7 @@ def vote(vote_button, response_a, response_b, model_a_name, model_b_name,
     language_a = detector.detect_language_of(response_a)
     language_b = detector.detect_language_of(response_b)
 
+    # TODO(#37): Move DB operations to db.py.
     doc_ref = db.collection("arena-summarizations").document(doc_id)
     doc["model_a_response_language"] = language_a.name.lower()
     doc["model_b_response_language"] = language_b.name.lower()

db.py

@@ -0,0 +1,114 @@
+"""
+This module handles the management of the database.
+"""
+from dataclasses import dataclass
+import enum
+import os
+from typing import List
+
+import firebase_admin
+from firebase_admin import credentials
+from firebase_admin import firestore
+from google.cloud.firestore_v1 import base_query
+import gradio as gr
+
+from credentials import get_credentials_json
+
+
+def get_required_env(name: str) -> str:
+  value = os.getenv(name)
+  if value is None:
+    raise ValueError(f"Environment variable {name} is not set")
+  return value
+
+
+RATINGS_COLLECTION = get_required_env("RATINGS_COLLECTION")
+SUMMARIZATIONS_COLLECTION = get_required_env("SUMMARIZATIONS_COLLECTION")
+TRANSLATIONS_COLLECTION = get_required_env("TRANSLATIONS_COLLECTION")
+
+if gr.NO_RELOAD:
+  firebase_admin.initialize_app(credentials.Certificate(get_credentials_json()))
+  db = firestore.client()
+
+
+class Category(enum.Enum):
+  SUMMARIZATION = "summarization"
+  TRANSLATION = "translation"
+
+
+@dataclass
+class Rating:
+  model: str
+  rating: int
+
+
+def get_ratings(category: Category, source_lang: str | None,
+                target_lang: str | None) -> List[Rating] | None:
+  doc_id = "#".join([category.value] +
+                    [lang for lang in (source_lang, target_lang) if lang])
+  # TODO(#37): Make it more clear what fields are in the document.
+  doc_dict = db.collection(RATINGS_COLLECTION).document(doc_id).get().to_dict()
+  if doc_dict is None:
+    return None
+
+  # TODO(#37): Return the timestamp as well.
+  doc_dict.pop("timestamp")
+
+  return [Rating(model, rating) for model, rating in doc_dict.items()]
+
+
+def set_ratings(category: Category, ratings: List[Rating], source_lang: str,
+                target_lang: str | None):
+  source_lang_lowercase = source_lang.lower()
+  target_lang_lowercase = target_lang.lower() if target_lang else None
+
+  doc_id = "#".join([category.value, source_lang_lowercase] +
+                    ([target_lang_lowercase] if target_lang_lowercase else []))
+  doc_ref = db.collection(RATINGS_COLLECTION).document(doc_id)
+
+  new_ratings = {rating.model: rating.rating for rating in ratings}
+  new_ratings["timestamp"] = firestore.SERVER_TIMESTAMP
+  doc_ref.set(new_ratings, merge=True)
+
+
+@dataclass
+class Battle:
+  model_a: str
+  model_b: str
+  winner: str
+
+
+def get_battles(category: Category, source_lang: str | None,
+                target_lang: str | None) -> List[Battle]:
+  source_lang_lowercase = source_lang.lower() if source_lang else None
+  target_lang_lowercase = target_lang.lower() if target_lang else None
+
+  if category == Category.SUMMARIZATION:
+    collection = db.collection(SUMMARIZATIONS_COLLECTION).order_by("timestamp")
+
+    if source_lang_lowercase:
+      collection = collection.where(filter=base_query.FieldFilter(
+          "model_a_response_language", "==", source_lang_lowercase)).where(
+              filter=base_query.FieldFilter("model_b_response_language", "==",
+                                            source_lang_lowercase))
+
+  elif category == Category.TRANSLATION:
+    collection = db.collection(TRANSLATIONS_COLLECTION).order_by("timestamp")
+
+    if source_lang_lowercase:
+      collection = collection.where(filter=base_query.FieldFilter(
+          "source_language", "==", source_lang_lowercase))
+
+    if target_lang_lowercase:
+      collection = collection.where(filter=base_query.FieldFilter(
+          "target_language", "==", target_lang_lowercase))
+
+  else:
+    raise ValueError(f"Invalid category: {category}")
+
+  docs = collection.stream()
+  battles = []
+  for doc in docs:
+    data = doc.to_dict()
+    battles.append(Battle(data["model_a"], data["model_b"], data["winner"]))
+  return battles

leaderboard.py

@@ -5,21 +5,13 @@ It provides a leaderboard component.
 from collections import defaultdict
 import enum
 import math
-from typing import Tuple
+from typing import Dict, List, Tuple
 
-import firebase_admin
-from firebase_admin import credentials
-from firebase_admin import firestore
-from google.cloud.firestore_v1 import base_query
 import gradio as gr
 import lingua
-import pandas as pd
 
-from credentials import get_credentials_json
-
-if gr.NO_RELOAD:
-  firebase_admin.initialize_app(credentials.Certificate(get_credentials_json()))
-  db = firestore.client()
+import db
+from db import get_battles
 
 SUPPORTED_LANGUAGES = [
     language.name.capitalize() for language in lingua.Language.all()
@@ -34,11 +26,16 @@ class LeaderboardTab(enum.Enum):
 
 
 # Ref: https://colab.research.google.com/drive/1RAWb22-PFNI-X1gPVzc927SGUdfr6nsR?usp=sharing#scrollTo=QLGc6DwxyvQc pylint: disable=line-too-long
-def compute_elo(battles, k=4, scale=400, base=10, initial_rating=1000):
+def compute_elo(battles: List[db.Battle],
+                k=4,
+                scale=400,
+                base=10,
+                initial_rating=1000) -> Dict[str, int]:
   rating = defaultdict(lambda: initial_rating)
 
-  for model_a, model_b, winner in battles[["model_a", "model_b",
-                                           "winner"]].itertuples(index=False):
+  for battle in battles:
+    model_a, model_b, winner = battle.model_a, battle.model_b, battle.winner
+
     rating_a = rating[model_a]
     rating_b = rating[model_b]
 
@@ -50,71 +47,41 @@ def compute_elo(battles, k=4, scale=400, base=10, initial_rating=1000):
     rating[model_a] += k * (scored_point_a - expected_score_a)
     rating[model_b] += k * (1 - scored_point_a - expected_score_b)
 
-  return rating
-
-
-def get_docs(tab: str,
-             summary_lang: str = None,
-             source_lang: str = None,
-             target_lang: str = None):
-  if tab == LeaderboardTab.SUMMARIZATION:
-    collection = db.collection("arena-summarizations").order_by("timestamp")
-
-    if summary_lang and (not summary_lang == ANY_LANGUAGE):
-      collection = collection.where(filter=base_query.FieldFilter(
-          "model_a_response_language", "==", summary_lang.lower())).where(
-              filter=base_query.FieldFilter("model_b_response_language", "==",
-                                            summary_lang.lower()))
-
-    return collection.stream()
-
-  if tab == LeaderboardTab.TRANSLATION:
-    collection = db.collection("arena-translations").order_by("timestamp")
+  return {model: math.floor(rating + 0.5) for model, rating in rating.items()}
 
-    if source_lang and (not source_lang == ANY_LANGUAGE):
-      collection = collection.where(filter=base_query.FieldFilter(
-          "source_language", "==", source_lang.lower()))
 
-    if target_lang and (not target_lang == ANY_LANGUAGE):
-      collection = collection.where(filter=base_query.FieldFilter(
-          "target_language", "==", target_lang.lower()))
+def load_elo_ratings(tab, source_lang: str, target_lang: str | None):
+  category = db.Category.SUMMARIZATION if tab == LeaderboardTab.SUMMARIZATION else db.Category.TRANSLATION
 
-    return collection.stream()
-
-
-def load_elo_ratings(tab,
-                     summary_lang: str = None,
-                     source_lang: str = None,
-                     target_lang: str = None):
-  docs = get_docs(tab, summary_lang, source_lang, target_lang)
-
-  battles = []
-  for doc in docs:
-    data = doc.to_dict()
-    battles.append({
-        "model_a": data["model_a"],
-        "model_b": data["model_b"],
-        "winner": data["winner"]
-    })
+  # TODO(#37): Call db.get_ratings and return the ratings if exists.
 
+  battles = get_battles(category,
+                        None if source_lang == ANY_LANGUAGE else source_lang,
+                        None if target_lang == ANY_LANGUAGE else target_lang)
   if not battles:
     return
 
-  battles = pd.DataFrame(battles)
-  ratings = compute_elo(battles)
+  computed_ratings = compute_elo(battles)
+
+  db.set_ratings(
+      category,
+      [db.Rating(model, rating) for model, rating in computed_ratings.items()],
+      source_lang, target_lang)
 
-  sorted_ratings = sorted(ratings.items(), key=lambda x: x[1], reverse=True)
+  sorted_ratings = sorted(
+      computed_ratings.items(),
+      key=lambda x: x[1],  # rating
+      reverse=True)
 
   rank = 0
   last_rating = None
   rating_rows = []
   for index, (model, rating) in enumerate(sorted_ratings):
-    int_rating = math.floor(rating + 0.5)
-    if int_rating != last_rating:
+    if rating != last_rating:
       rank = index + 1
 
-    rating_rows.append([rank, model, int_rating])
-    last_rating = int_rating
+    rating_rows.append([rank, model, rating])
+    last_rating = rating
 
   return rating_rows
 
@@ -123,9 +90,9 @@ LEADERBOARD_UPDATE_INTERVAL = 600  # 10 minutes
 LEADERBOARD_INFO = "The leaderboard is updated every 10 minutes."
 
 
-def update_filtered_leaderboard(tab, summary_lang: str, source_lang: str,
-                                target_lang: str):
-  new_value = load_elo_ratings(tab, summary_lang, source_lang, target_lang)
+def update_filtered_leaderboard(tab: str, source_lang: str,
+                                target_lang: str | None):
+  new_value = load_elo_ratings(tab, source_lang, target_lang)
   return gr.update(value=new_value)
 
 
@@ -149,14 +116,15 @@ def build_leaderboard():
           headers=["Rank", "Model", "Elo rating"],
           datatype=["number", "str", "number"],
           value=lambda: load_elo_ratings(LeaderboardTab.SUMMARIZATION,
-                                         ANY_LANGUAGE),
+                                         ANY_LANGUAGE, None),
           elem_classes="leaderboard",
           visible=False)
 
       original_summarization = gr.Dataframe(
           headers=["Rank", "Model", "Elo rating"],
           datatype=["number", "str", "number"],
-          value=lambda: load_elo_ratings(LeaderboardTab.SUMMARIZATION),
+          value=lambda: load_elo_ratings(LeaderboardTab.SUMMARIZATION,
+                                         ANY_LANGUAGE, None),
           every=LEADERBOARD_UPDATE_INTERVAL,
           elem_classes="leaderboard")
       gr.Markdown(LEADERBOARD_INFO)
@@ -165,7 +133,6 @@ def build_leaderboard():
           fn=update_filtered_leaderboard,
           inputs=[
               gr.State(LeaderboardTab.SUMMARIZATION), summary_language,
-              gr.State(None),
               gr.State(None)
           ],
           outputs=filtered_summarization).then(
@@ -197,7 +164,8 @@ def build_leaderboard():
       original_translation = gr.Dataframe(
           headers=["Rank", "Model", "Elo rating"],
           datatype=["number", "str", "number"],
-          value=lambda: load_elo_ratings(LeaderboardTab.TRANSLATION),
+          value=lambda: load_elo_ratings(LeaderboardTab.TRANSLATION,
+                                         ANY_LANGUAGE, ANY_LANGUAGE),
           every=LEADERBOARD_UPDATE_INTERVAL,
           elem_classes="leaderboard")
       gr.Markdown(LEADERBOARD_INFO)
@@ -205,8 +173,8 @@ def build_leaderboard():
       source_language.change(
           fn=update_filtered_leaderboard,
           inputs=[
-              gr.State(LeaderboardTab.TRANSLATION),
-              gr.State(None), source_language, target_language
+              gr.State(LeaderboardTab.TRANSLATION), source_language,
+              target_language
           ],
           outputs=filtered_translation).then(
               fn=toggle_leaderboard,
@@ -215,8 +183,8 @@ def build_leaderboard():
       target_language.change(
           fn=update_filtered_leaderboard,
           inputs=[
-              gr.State(LeaderboardTab.TRANSLATION),
-              gr.State(None), source_language, target_language
+              gr.State(LeaderboardTab.TRANSLATION), source_language,
+              target_language
           ],
           outputs=filtered_translation).then(
               fn=toggle_leaderboard,

response.py

@@ -11,7 +11,7 @@ from uuid import uuid4
 from firebase_admin import firestore
 import gradio as gr
 
-from leaderboard import db
+from db import db
 from model import ContextWindowExceededError
 from model import Model
 from model import supported_models
@@ -22,7 +22,7 @@ logging.basicConfig()
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-
+# TODO(#37): Move DB operations to db.py.
 def get_history_collection(category: str):
   if category == Category.SUMMARIZE.value:
     return db.collection("arena-summarization-history")