fix wandb so mypy doesn't complain

.mypy.ini

@@ -8,6 +8,9 @@ ignore_missing_imports = True
 [mypy-axolotl.monkeypatch.*]
 ignore_errors = True
 
+[mypy-axolotl.utils.callbacks]
+disable_error_code = attr-defined
+
 [mypy-flash_attn.*]
 ignore_missing_imports = True
 

requirements.txt

@@ -30,3 +30,4 @@ scipy
 scikit-learn==1.2.2
 pynvml
 art
+wandb

src/axolotl/utils/callbacks.py

@@ -11,7 +11,6 @@
 import pandas as pd
 import torch
 import torch.distributed as dist
-import wandb
 from datasets import load_dataset
 from optimum.bettertransformer import BetterTransformer
 from tqdm import tqdm
@@ -25,6 +24,7 @@
 )
 from transformers.trainer_utils import PREFIX_CHECKPOINT_DIR, IntervalStrategy
 
+import wandb
 from axolotl.utils.bench import log_gpu_memory_usage
 from axolotl.utils.distributed import (
     barrier,
@@ -367,7 +367,7 @@ def on_evaluate(
                 output_scores=False,
             )
 
-            def logits_to_tokens(logits) -> str:
+            def logits_to_tokens(logits) -> torch.Tensor:
                 probabilities = torch.softmax(logits, dim=-1)
                 # Get the predicted token ids (the ones with the highest probability)
                 predicted_token_ids = torch.argmax(probabilities, dim=-1)