From 5cf11e5ab9591652ee025069658f9af5a98e455e Mon Sep 17 00:00:00 2001 From: Matt Date: Thu, 12 Dec 2024 13:59:24 +0000 Subject: [PATCH] Fix type hints for apply_chat_template (#35216) --- src/transformers/tokenization_utils_base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/transformers/tokenization_utils_base.py b/src/transformers/tokenization_utils_base.py index f4e5b9b3aaf314..de0bc87b26b676 100644 --- a/src/transformers/tokenization_utils_base.py +++ b/src/transformers/tokenization_utils_base.py @@ -28,7 +28,7 @@ from contextlib import contextmanager from dataclasses import dataclass from inspect import isfunction -from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union +from typing import TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Tuple, Union import numpy as np from packaging import version @@ -1527,7 +1527,7 @@ def get_vocab(self) -> Dict[str, int]: def apply_chat_template( self, conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]]], - tools: Optional[List[Dict]] = None, + tools: Optional[List[Union[Dict, Callable]]] = None, documents: Optional[List[Dict[str, str]]] = None, chat_template: Optional[str] = None, add_generation_prompt: bool = False,