[#57] Store language information in summarization (#61)

* [#57] Store language information in summarization

Changes:
- Store language information in summarization battle
- Add language filter to the summarization leaderboard

* Remove Apply button

* Update

app.py

@@ -6,6 +6,7 @@ from uuid import uuid4
 
 from firebase_admin import firestore
 import gradio as gr
+import lingua
 
 from leaderboard import build_leaderboard
 from leaderboard import db
@@ -13,6 +14,8 @@ from leaderboard import SUPPORTED_TRANSLATION_LANGUAGES
 import response
 from response import get_responses
 
+detector = lingua.LanguageDetectorBuilder.from_all_languages().build()
+
 
 class VoteOptions(enum.Enum):
   MODEL_A = "Model A is better"
@@ -41,7 +44,12 @@ def vote(vote_button, response_a, response_b, model_a_name, model_b_name,
   }
 
   if category == response.Category.SUMMARIZE.value:
+    language_a = detector.detect_language_of(response_a)
+    language_b = detector.detect_language_of(response_b)
+
     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()
     doc_ref.set(doc)
 
     return outputs

leaderboard.py

@@ -11,6 +11,7 @@ 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
@@ -20,7 +21,7 @@ firebase_admin.initialize_app(credentials.Certificate(get_credentials_json()))
 db = firestore.client()
 
 SUPPORTED_TRANSLATION_LANGUAGES = [
-    "Korean", "English", "Chinese", "Japanese", "Spanish", "French"
+    language.name.capitalize() for language in lingua.Language.all()
 ]
 
 
@@ -49,9 +50,20 @@ def compute_elo(battles, k=4, scale=400, base=10, initial_rating=1000):
   return rating
 
 
-def get_docs(tab: str, source_lang: str = None, target_lang: str = None):
+def get_docs(tab: str,
+             summary_lang: str = None,
+             source_lang: str = None,
+             target_lang: str = None):
   if tab == LeaderboardTab.SUMMARIZATION:
-    return db.collection("arena-summarizations").order_by("timestamp").stream()
+    collection = db.collection("arena-summarizations").order_by("timestamp")
+
+    if summary_lang:
+      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")
@@ -67,8 +79,11 @@ def get_docs(tab: str, source_lang: str = None, target_lang: str = None):
     return collection.stream()
 
 
-def load_elo_ratings(tab, source_lang: str = None, target_lang: str = None):
-  docs = get_docs(tab, source_lang, target_lang)
+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:
@@ -94,28 +109,50 @@ LEADERBOARD_UPDATE_INTERVAL = 600  # 10 minutes
 LEADERBOARD_INFO = "The leaderboard is updated every 10 minutes."
 
 DEFAULT_FILTER_OPTIONS = {
+    "summary_language": "English",
     "source_language": "English",
     "target_language": "Spanish"
 }
 
-filtered_dataframe = gr.DataFrame(
-    headers=["Rank", "Model", "Elo rating"],
-    datatype=["number", "str", "number"],
-    value=lambda: load_elo_ratings(
-        LeaderboardTab.TRANSLATION, DEFAULT_FILTER_OPTIONS[
-            "source_language"], DEFAULT_FILTER_OPTIONS["target_language"]),
-    elem_classes="leaderboard")
 
-
-def update_filtered_leaderboard(source_lang, target_lang):
-  new_value = load_elo_ratings(LeaderboardTab.TRANSLATION, source_lang,
-                               target_lang)
+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)
   return gr.update(value=new_value)
 
 
 def build_leaderboard():
   with gr.Tabs():
     with gr.Tab(LeaderboardTab.SUMMARIZATION.value):
+      with gr.Accordion("Filter", open=False):
+        with gr.Row():
+          languages = [
+              language.name.capitalize() for language in lingua.Language.all()
+          ]
+          summary_language = gr.Dropdown(
+              choices=languages,
+              value=DEFAULT_FILTER_OPTIONS["summary_language"],
+              label="Summary language",
+              interactive=True)
+
+        with gr.Row():
+          filtered_summarization = gr.DataFrame(
+              headers=["Rank", "Model", "Elo rating"],
+              datatype=["number", "str", "number"],
+              value=lambda: load_elo_ratings(
+                  LeaderboardTab.SUMMARIZATION, DEFAULT_FILTER_OPTIONS[
+                      "summary_language"]),
+              elem_classes="leaderboard")
+
+      summary_language.change(fn=update_filtered_leaderboard,
+                              inputs=[
+                                  gr.State(LeaderboardTab.SUMMARIZATION),
+                                  summary_language,
+                                  gr.State(),
+                                  gr.State()
+                              ],
+                              outputs=filtered_summarization)
+
       gr.Dataframe(headers=["Rank", "Model", "Elo rating"],
                    datatype=["number", "str", "number"],
                    value=lambda: load_elo_ratings(LeaderboardTab.SUMMARIZATION),
@@ -137,15 +174,30 @@ def build_leaderboard():
               value=DEFAULT_FILTER_OPTIONS["target_language"],
               interactive=True)
 
+        with gr.Row():
+          filtered_translation = gr.DataFrame(
+              headers=["Rank", "Model", "Elo rating"],
+              datatype=["number", "str", "number"],
+              value=lambda: load_elo_ratings(
+                  LeaderboardTab.TRANSLATION, DEFAULT_FILTER_OPTIONS[
+                      "source_language"], DEFAULT_FILTER_OPTIONS[
+                          "target_language"]),
+              elem_classes="leaderboard")
+
           source_language.change(fn=update_filtered_leaderboard,
-                                 inputs=[source_language, target_language],
-                                 outputs=filtered_dataframe)
+                                 inputs=[
+                                     gr.State(LeaderboardTab.TRANSLATION),
+                                     gr.State(), source_language,
+                                     target_language
+                                 ],
+                                 outputs=filtered_translation)
           target_language.change(fn=update_filtered_leaderboard,
-                                 inputs=[source_language, target_language],
-                                 outputs=filtered_dataframe)
-
-        with gr.Row():
-          filtered_dataframe.render()
+                                 inputs=[
+                                     gr.State(LeaderboardTab.TRANSLATION),
+                                     gr.State(), source_language,
+                                     target_language
+                                 ],
+                                 outputs=filtered_translation)
 
       gr.Dataframe(headers=["Rank", "Model", "Elo rating"],
                    datatype=["number", "str", "number"],

requirements.txt

@@ -54,6 +54,7 @@ Jinja2==3.1.3
 jsonschema==4.21.1
 jsonschema-specifications==2023.12.1
 kiwisolver==1.4.5
+lingua-language-detector==2.0.2
 litellm==1.32.7
 markdown-it-py==3.0.0
 MarkupSafe==2.1.5