Add MNIST task (#153)

* Add MNIST task

Author: speed1313

README.md

@@ -36,8 +36,8 @@ uv run --group normal python examples/sample.py \
   --model_id llava-hf/llava-1.5-7b-hf \
   --task_id japanese-heron-bench  \
   --result_dir result  \
-  --metrics "heron-bench" \
-  --judge_model "gpt-4o-2024-11-20" \
+  --metrics heron-bench \
+  --judge_model gpt-4o-2024-11-20 \
   --overwrite
 ```
 
@@ -136,9 +136,8 @@ See `eval_all.sh` for the complete list of model dependencies.
 When adding a new group, remember to configure [conflict](https://docs.astral.sh/uv/concepts/projects/config/#conflicting-dependencies).
 
 ## Browse Predictions with Streamlit
-
 ```bash
-uv run streamlit run scripts/browse_prediction.py --task_id "japanese-heron-bench" --result_dir "result"
+uv run streamlit run scripts/browse_prediction.py -- --task_id japanese-heron-bench --result_dir result --model_list llava-hf/llava-1.5-7b-hf
 ```
 
 ![Streamlit](./assets/streamlit_visualization.png)

examples/sample.py

@@ -14,7 +14,11 @@
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_id", default="llava-hf/llava-1.5-7b-hf")
-    parser.add_argument("--task_id", default="japanese-heron-bench")
+    parser.add_argument(
+        "--task_id",
+        default="japanese-heron-bench",
+        help=f"Task ID to evaluate. Available: {eval_mm.TaskRegistry().get_task_list()}",
+    )
     parser.add_argument("--judge_model", default="gpt-4o-2024-11-20")
     parser.add_argument("--batch_size_for_evaluation", type=int, default=10)
     parser.add_argument("--overwrite", action="store_true")
@@ -27,7 +31,13 @@ def parse_args():
     parser.add_argument("--do_sample", action="store_true", default=False)
     parser.add_argument("--use_cache", action="store_true", default=True)
     parser.add_argument("--max_dataset_len", type=int)
-    parser.add_argument("--metrics", default="llm_as_a_judge_heron_bench")
+    parser.add_argument(
+        "--metrics",
+        type=str,
+        nargs="+",
+        default=["heron-bench"],
+        help=f"Metrics to evaluate. Available: {eval_mm.ScorerRegistry().get_metric_list()}",
+    )
     parser.add_argument(
         "--rotate_choices", action="store_true", help="This option is used in MECHA-ja"
     )
@@ -137,7 +147,6 @@ def save_final_results(preds, task, metrics, scores_by_metric, output_path):
 
 def main():
     args = parse_args()
-    metrics = args.metrics.split(",")
 
     gen_kwargs = GenerationConfig(
         max_new_tokens=args.max_new_tokens,
@@ -163,10 +172,10 @@ def main():
         logger.info("Inference only mode. Skipping evaluation.")
         return
 
-    scores_by_metric, aggregated_metrics = evaluate(args, task, preds, metrics)
+    scores_by_metric, aggregated_metrics = evaluate(args, task, preds, args.metrics)
 
     prediction_path = os.path.join(output_dir, "prediction.jsonl")
-    save_final_results(preds, task, metrics, scores_by_metric, prediction_path)
+    save_final_results(preds, task, args.metrics, scores_by_metric, prediction_path)
 
     evaluation_path = os.path.join(output_dir, "evaluation.jsonl")
     with open(evaluation_path, "w") as f:

scripts/browse_prediction.py

@@ -10,6 +10,7 @@ def parse_args():
     parser = ArgumentParser()
     parser.add_argument("--task_id", type=str, default="japanese-heron-bench")
     parser.add_argument("--result_dir", type=str, default="result")
+    parser.add_argument("--model_list", type=str, nargs="+", default=[])
 
     return parser.parse_args()
 
@@ -23,18 +24,11 @@ def scrollable_text(text):
 if __name__ == "__main__":
     args = parse_args()
 
-    task = eval_mm.tasks.TaskRegistry().get_task_cls(args.task_id)(
-        eval_mm.tasks.TaskConfig()
-    )
+    task = eval_mm.TaskRegistry().load_task(args.task_id)
 
     # Load model prediction
-    model_list = [
-        "google/gemma-3-12b-it",
-        "google/gemma-3-27b-it",
-        "microsoft/Phi-4-multimodal-instruct",
-    ]
     predictions_per_model = {}
-    for model_id in model_list:
+    for model_id in args.model_list:
         prediction_path = os.path.join(
             args.result_dir, args.task_id, model_id, "prediction.jsonl"
         )
@@ -50,8 +44,8 @@ def scrollable_text(text):
 
     SAMPLES_PER_PAGE = 30  # 1ページに表示する件数
     # Question ID, Image, Question, Answer, Prediction_model1, Prediction_model2,..
-    column_width_list = [1, 3, 3, 3] + [4] * len(model_list)
-    st.write(f"# {args.task_id} dataset")
+    column_width_list = [1, 3, 3, 3] + [4] * len(args.model_list)
+    st.write(f"# {args.task_id}")
 
     def show_sample(idx):
         sample = ds[idx]
@@ -64,8 +58,8 @@ def show_sample(idx):
         cols[3].markdown(
             scrollable_text(task.doc_to_answer(sample)), unsafe_allow_html=True
         )
-        for model_id in model_list:
-            cols[4 + model_list.index(model_id)].markdown(
+        for model_id in args.model_list:
+            cols[4 + args.model_list.index(model_id)].markdown(
                 scrollable_text(predictions_per_model[model_id][idx]["text"]),
                 unsafe_allow_html=True,
             )
@@ -93,8 +87,10 @@ def show_sample(idx):
     header_cols[1].markdown("Image")
     header_cols[2].markdown("Question")
     header_cols[3].markdown("Answer")
-    for model_id in model_list:
-        header_cols[4 + model_list.index(model_id)].markdown(f"Prediction ({model_id})")
+    for model_id in args.model_list:
+        header_cols[4 + args.model_list.index(model_id)].markdown(
+            f"Prediction ({model_id})"
+        )
 
     # サンプルを表示
     for idx in range(start_idx, end_idx):

scripts/make_leaderboard.py

@@ -286,7 +286,8 @@ def parse_args():
     parser.add_argument(
         "--task_id_list",
         type=str,
-        help="List of task IDs to include in the leaderboard (e.g. jmmmu,mmmu). If not specified, all tasks will be included.",
+        nargs="+",
+        help=f"List of task IDs to include in the leaderboard. Available: {TASK_ALIAS.keys()}",
     )
     return parser.parse_args()
 

src/eval_mm/tasks/mnist.py

@@ -0,0 +1,43 @@
+from eval_mm.tasks.task import Task
+from datasets import load_dataset, Dataset
+from PIL import Image
+
+
+class MNIST(Task):
+    def __init__(self, config):
+        super().__init__(config)
+
+    @staticmethod
+    def _prepare_dataset() -> Dataset:
+        ds = load_dataset("ylecun/mnist", split="test")
+        ds = ds.map(lambda example, idx: {"question_id": idx}, with_indices=True)
+        return ds
+
+    @staticmethod
+    def doc_to_text(doc) -> str:
+        return "画像に写っている数字は何ですか？ 数字のみを出力してください。"
+
+    @staticmethod
+    def doc_to_visual(doc) -> list[Image.Image]:
+        return [doc["image"]]
+
+    @staticmethod
+    def doc_to_id(doc) -> int:
+        return doc["question_id"]
+
+    @staticmethod
+    def doc_to_answer(doc) -> str:
+        return str(doc["label"])
+
+
+def test_task():
+    from eval_mm.tasks.task import TaskConfig
+
+    task = MNIST(TaskConfig())
+    ds = task.dataset
+    print(ds[0])
+    assert isinstance(task.doc_to_text(ds[0]), str)
+    assert isinstance(task.doc_to_visual(ds[0]), list)
+    assert isinstance(task.doc_to_visual(ds[0])[0], Image.Image)
+    assert isinstance(task.doc_to_id(ds[0]), int)
+    assert isinstance(task.doc_to_answer(ds[0]), str)

src/eval_mm/tasks/task_registry.py

@@ -9,6 +9,7 @@
 from .jic_vqa import JICVQA
 from .mecha_ja import MECHAJa
 from .mmmlu import MMMLU
+from .mnist import MNIST
 from .task import TaskConfig, Task
 
 
@@ -27,6 +28,7 @@ class TaskRegistry:
         "jic-vqa": JICVQA,
         "mecha-ja": MECHAJa,
         "mmmlu": MMMLU,
+        "mnist": MNIST,
     }
 
     @classmethod