fix for old python

okdshin · Dec 9, 2023 · 4c0a2be · 4c0a2be
1 parent bd23a0b
commit 4c0a2be
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/test_language_model_for_completion.py b/tests/test_language_model_for_completion.py
@@ -1,5 +1,5 @@
 from unittest.mock import Mock
-from typing import Any
+from typing import Any, Dict
 
 from transformers import (
     AutoTokenizer,
@@ -47,7 +47,7 @@ def prepare_inputs_for_generation(
         self,
         input_ids: torch.Tensor,
         **kwargs,
-    ) -> dict[str, Any]:
+    ) -> Dict[str, Any]:
         model_inputs = {"input_ids": input_ids}
         return model_inputs