From 6f166464d873c922eae7da6748b316639091acf9 Mon Sep 17 00:00:00 2001
From: Wing Lian <wing.lian@gmail.com>
Date: Sat, 19 Aug 2023 21:26:32 -0400
Subject: [PATCH] include metrics in callback

---
 src/axolotl/utils/callbacks.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/axolotl/utils/callbacks.py b/src/axolotl/utils/callbacks.py
index 6bb42d5aa..89411676f 100644
--- a/src/axolotl/utils/callbacks.py
+++ b/src/axolotl/utils/callbacks.py
@@ -4,7 +4,7 @@ from __future__ import annotations
 
 import logging
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict
 
 import evaluate
 import numpy as np
@@ -154,6 +154,7 @@ def mmlu_eval_callback_factory(trainer, tokenizer):
             args: AxolotlTrainingArguments,
             state: TrainerState,  # pylint: disable=unused-argument
             control: TrainerControl,  # pylint: disable=unused-argument
+            metrics: Dict[str, float],  # pylint: disable=unused-argument
             **kwargs,  # pylint: disable=unused-argument
         ):
             data_loader = trainer.get_eval_dataloader(mmlu_dataset)