Change num_few_shots enum to handle int type

app.py

@@ -96,7 +96,7 @@ def filter_models(
     size_query: list[str],
     precision_query: list[str],
     add_special_tokens_query: list[str],
-    num_few_shots_query: list[str],
+    num_few_shots_query: list[int],
     version_query: list[str],
     vllm_query: list[str],
 ) -> pd.DataFrame:
@@ -122,7 +122,7 @@ def filter_models(
     df = df[df["Add Special Tokens"].isin(add_special_tokens_query)]
 
     # Filter by number of few-shot examples
-    df = df[df["Few-shot"].astype(str).isin(num_few_shots_query)]
+    df = df[df["Few-shot"].isin(num_few_shots_query)]
 
     # Filter by evaluator version
     df = df[df["llm-jp-eval version"].isin(version_query)]
@@ -175,7 +175,7 @@ def update_table(
     precision_query: list[str],
     size_query: list[str],
     add_special_tokens_query: list[str],
-    num_few_shots_query: list[str],
+    num_few_shots_query: list[int],
     version_query: list[str],
     vllm_query: list[str],
     query: str,
@@ -211,7 +211,7 @@ if len(leaderboard_df) > 0:
         list(NUMERIC_INTERVALS.keys()),
         [i.value.name for i in Precision],
         [i.value.name for i in AddSpecialTokens],
-        [i.value.name for i in NumFewShots],
+        [i.value for i in NumFewShots],
         [i.value.name for i in LLMJpEvalVersion],
         [i.value.name for i in VllmVersion],
     )
@@ -418,8 +418,8 @@ with gr.Blocks() as demo_leaderboard:
             )
             filter_columns_num_few_shots = gr.CheckboxGroup(
                 label="Num Few Shots",
-                choices=[i.value.name for i in NumFewShots],
-                value=[i.value.name for i in NumFewShots],
+                choices=[i.value for i in NumFewShots],
+                value=[i.value for i in NumFewShots],
                 elem_id="filter-columns-num-few-shots",
             )
             filter_columns_version = gr.CheckboxGroup(

src/display/utils.py

@@ -147,18 +147,13 @@ class AddSpecialTokens(Enum):
 
 
 class NumFewShots(Enum):
-    shots_0 = ModelDetails("0")
-    shots_4 = ModelDetails("4")
+    shots_0 = 0
+    shots_4 = 4
 
     @staticmethod
-    def from_str(shots: str) -> "NumFewShots":
+    def from_int(shots: int) -> "NumFewShots":
         try:
-            int_shots = int(shots)
-            if int_shots == 0:
-                return NumFewShots.shots_0
-            if int_shots == 4:
-                return NumFewShots.shots_4
-            raise ValueError
+            return NumFewShots(shots)
         except ValueError:
             raise ValueError(f"Unsupported number of shots: {shots}. Must be either 0 or 4")