huggingface-hub 0.24.7__py3-none-any.whl → 0.25.0rc0__py3-none-any.whl

huggingface_hub/inference/_generated/_async_client.py

@@ -32,6 +32,7 @@ from typing import (
     List,
     Literal,
     Optional,
+    Set,
     Union,
     overload,
 )
@@ -86,9 +87,6 @@ from huggingface_hub.inference._generated.types import (
     ZeroShotClassificationOutputElement,
     ZeroShotImageClassificationOutputElement,
 )
-from huggingface_hub.inference._types import (
-    ConversationalOutput,  # soon to be removed
-)
 from huggingface_hub.utils import (
     build_hf_headers,
 )
@@ -99,6 +97,7 @@ from .._common import _async_yield_from, _import_aiohttp
 
 if TYPE_CHECKING:
     import numpy as np
+    from aiohttp import ClientResponse, ClientSession
     from PIL.Image import Image
 
 logger = logging.getLogger(__name__)
@@ -120,7 +119,9 @@ class AsyncInferenceClient:
             or a URL to a deployed Inference Endpoint. Defaults to None, in which case a recommended model is
             automatically selected for the task.
             Note: for better compatibility with OpenAI's client, `model` has been aliased as `base_url`. Those 2
-            arguments are mutually exclusive and have the exact same behavior.
+            arguments are mutually exclusive. If using `base_url` for chat completion, the `/chat/completions` suffix
+            path will be appended to the base URL (see the [TGI Messages API](https://huggingface.co/docs/text-generation-inference/en/messages_api)
+            documentation for details). When passing a URL as `model`, the client will not append any suffix path to it.
         token (`str` or `bool`, *optional*):
             Hugging Face token. Will default to the locally saved token if not provided.
             Pass `token=False` if you don't want to send your token to the server.
@@ -134,6 +135,10 @@ class AsyncInferenceClient:
             Values in this dictionary will override the default values.
         cookies (`Dict[str, str]`, `optional`):
             Additional cookies to send to the server.
+        trust_env ('bool', 'optional'):
+            Trust environment settings for proxy configuration if the parameter is `True` (`False` by default).
+        proxies (`Any`, `optional`):
+            Proxies to use for the request.
         base_url (`str`, `optional`):
             Base URL to run inference. This is a duplicated argument from `model` to make [`InferenceClient`]
             follow the same pattern as `openai.OpenAI` client. Cannot be used if `model` is set. Defaults to None.
@@ -151,6 +156,7 @@ class AsyncInferenceClient:
         timeout: Optional[float] = None,
         headers: Optional[Dict[str, str]] = None,
         cookies: Optional[Dict[str, str]] = None,
+        trust_env: bool = False,
         proxies: Optional[Any] = None,
         # OpenAI compatibility
         base_url: Optional[str] = None,
@@ -160,7 +166,8 @@ class AsyncInferenceClient:
             raise ValueError(
                 "Received both `model` and `base_url` arguments. Please provide only one of them."
                 " `base_url` is an alias for `model` to make the API compatible with OpenAI's client."
-                " It has the exact same behavior as `model`."
+                " If using `base_url` for chat completion, the `/chat/completions` suffix path will be appended to the base url."
+                " When passing a URL as `model`, the client will not append any suffix path to it."
             )
         if token is not None and api_key is not None:
             raise ValueError(
@@ -176,11 +183,15 @@ class AsyncInferenceClient:
             self.headers.update(headers)
         self.cookies = cookies
         self.timeout = timeout
+        self.trust_env = trust_env
         self.proxies = proxies
 
         # OpenAI compatibility
         self.base_url = base_url
 
+        # Keep track of the sessions to close them properly
+        self._sessions: Dict["ClientSession", Set["ClientResponse"]] = dict()
+
     def __repr__(self):
         return f"<InferenceClient(model='{self.model if self.model else ''}', timeout={self.timeout})>"
 
@@ -265,7 +276,7 @@ class AsyncInferenceClient:
             warnings.warn("Ignoring `json` as `data` is passed as binary.")
 
         # Set Accept header if relevant
-        headers = self.headers.copy()
+        headers = dict()
         if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
             headers["Accept"] = "image/png"
 
@@ -275,12 +286,10 @@ class AsyncInferenceClient:
             with _open_as_binary(data) as data_as_binary:
                 # Do not use context manager as we don't want to close the connection immediately when returning
                 # a stream
-                client = aiohttp.ClientSession(
-                    headers=headers, cookies=self.cookies, timeout=aiohttp.ClientTimeout(self.timeout)
-                )
+                session = self._get_client_session(headers=headers)
 
                 try:
-                    response = await client.post(url, json=json, data=data_as_binary, proxy=self.proxies)
+                    response = await session.post(url, json=json, data=data_as_binary, proxy=self.proxies)
                     response_error_payload = None
                     if response.status != 200:
                         try:
@@ -289,18 +298,18 @@ class AsyncInferenceClient:
                             pass
                     response.raise_for_status()
                     if stream:
-                        return _async_yield_from(client, response)
+                        return _async_yield_from(session, response)
                     else:
                         content = await response.read()
-                        await client.close()
+                        await session.close()
                         return content
                 except asyncio.TimeoutError as error:
-                    await client.close()
+                    await session.close()
                     # Convert any `TimeoutError` to a `InferenceTimeoutError`
                     raise InferenceTimeoutError(f"Inference call timed out: {url}") from error  # type: ignore
                 except aiohttp.ClientResponseError as error:
                     error.response_error_payload = response_error_payload
-                    await client.close()
+                    await session.close()
                     if response.status == 422 and task is not None:
                         error.message += f". Make sure '{task}' task is supported by the model."
                     if response.status == 503:
@@ -322,9 +331,35 @@ class AsyncInferenceClient:
                         continue
                     raise error
                 except Exception:
-                    await client.close()
+                    await session.close()
                     raise
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.close()
+
+    def __del__(self):
+        if len(self._sessions) > 0:
+            warnings.warn(
+                "Deleting 'AsyncInferenceClient' client but some sessions are still open. "
+                "This can happen if you've stopped streaming data from the server before the stream was complete. "
+                "To close the client properly, you must call `await client.close()` "
+                "or use an async context (e.g. `async with AsyncInferenceClient(): ...`."
+            )
+
+    async def close(self):
+        """Close all open sessions.
+
+        By default, 'aiohttp.ClientSession' objects are closed automatically when a call is completed. However, if you
+        are streaming data from the server and you stop before the stream is complete, you must call this method to
+        close the session properly.
+
+        Another possibility is to use an async context (e.g. `async with AsyncInferenceClient(): ...`).
+        """
+        await asyncio.gather(*[session.close() for session in self._sessions.keys()])
+
     async def audio_classification(
         self,
         audio: ContentT,
@@ -815,134 +850,66 @@ class AsyncInferenceClient:
         '{\n\n"activity": "bike ride",\n"animals": ["puppy", "cat", "raccoon"],\n"animals_seen": 3,\n"location": "park"}'
         ```
         """
-        # Determine model
-        # `self.xxx` takes precedence over the method argument only in `chat_completion`
-        # since `chat_completion(..., model=xxx)` is also a payload parameter for the
-        # server, we need to handle it differently
-        model = self.base_url or self.model or model or self.get_recommended_model("text-generation")
-        is_url = model.startswith(("http://", "https://"))
-
-        # First, resolve the model chat completions URL
-        if model == self.base_url:
-            # base_url passed => add server route
-            model_url = model.rstrip("/")
-            if not model_url.endswith("/v1"):
-                model_url += "/v1"
-            model_url += "/chat/completions"
-        elif is_url:
-            # model is a URL => use it directly
-            model_url = model
-        else:
-            # model is a model ID => resolve it + add server route
-            model_url = self._resolve_url(model).rstrip("/") + "/v1/chat/completions"
+        model_url = self._resolve_chat_completion_url(model)
 
         # `model` is sent in the payload. Not used by the server but can be useful for debugging/routing.
         # If it's a ID on the Hub => use it. Otherwise, we use a random string.
-        model_id = model if not is_url and model.count("/") == 1 else "tgi"
-
-        data = await self.post(
-            model=model_url,
-            json=dict(
-                model=model_id,
-                messages=messages,
-                frequency_penalty=frequency_penalty,
-                logit_bias=logit_bias,
-                logprobs=logprobs,
-                max_tokens=max_tokens,
-                n=n,
-                presence_penalty=presence_penalty,
-                response_format=response_format,
-                seed=seed,
-                stop=stop,
-                temperature=temperature,
-                tool_choice=tool_choice,
-                tool_prompt=tool_prompt,
-                tools=tools,
-                top_logprobs=top_logprobs,
-                top_p=top_p,
-                stream=stream,
-            ),
+        model_id = model or self.model or "tgi"
+        if model_id.startswith(("http://", "https://")):
+            model_id = "tgi"  # dummy value
+
+        payload = dict(
+            model=model_id,
+            messages=messages,
+            frequency_penalty=frequency_penalty,
+            logit_bias=logit_bias,
+            logprobs=logprobs,
+            max_tokens=max_tokens,
+            n=n,
+            presence_penalty=presence_penalty,
+            response_format=response_format,
+            seed=seed,
+            stop=stop,
+            temperature=temperature,
+            tool_choice=tool_choice,
+            tool_prompt=tool_prompt,
+            tools=tools,
+            top_logprobs=top_logprobs,
+            top_p=top_p,
             stream=stream,
         )
+        payload = {key: value for key, value in payload.items() if value is not None}
+        data = await self.post(model=model_url, json=payload, stream=stream)
 
         if stream:
             return _async_stream_chat_completion_response(data)  # type: ignore[arg-type]
 
         return ChatCompletionOutput.parse_obj_as_instance(data)  # type: ignore[arg-type]
 
-    async def conversational(
-        self,
-        text: str,
-        generated_responses: Optional[List[str]] = None,
-        past_user_inputs: Optional[List[str]] = None,
-        *,
-        parameters: Optional[Dict[str, Any]] = None,
-        model: Optional[str] = None,
-    ) -> ConversationalOutput:
-        """
-        Generate conversational responses based on the given input text (i.e. chat with the API).
-
-        <Tip warning={true}>
-
-        [`InferenceClient.conversational`] API is deprecated and will be removed in a future release. Please use
-        [`InferenceClient.chat_completion`] instead.
+    def _resolve_chat_completion_url(self, model: Optional[str] = None) -> str:
+        # Since `chat_completion(..., model=xxx)` is also a payload parameter for the server, we need to handle 'model' differently.
+        # `self.base_url` and `self.model` takes precedence over 'model' argument only in `chat_completion`.
+        model_id_or_url = self.base_url or self.model or model or self.get_recommended_model("text-generation")
 
-        </Tip>
+        # Resolve URL if it's a model ID
+        model_url = (
+            model_id_or_url
+            if model_id_or_url.startswith(("http://", "https://"))
+            else self._resolve_url(model_id_or_url, task="text-generation")
+        )
 
-        Args:
-            text (`str`):
-                The last input from the user in the conversation.
-            generated_responses (`List[str]`, *optional*):
-                A list of strings corresponding to the earlier replies from the model. Defaults to None.
-            past_user_inputs (`List[str]`, *optional*):
-                A list of strings corresponding to the earlier replies from the user. Should be the same length as
-                `generated_responses`. Defaults to None.
-            parameters (`Dict[str, Any]`, *optional*):
-                Additional parameters for the conversational task. Defaults to None. For more details about the available
-                parameters, please refer to [this page](https://huggingface.co/docs/api-inference/detailed_parameters#conversational-task)
-            model (`str`, *optional*):
-                The model to use for the conversational task. Can be a model ID hosted on the Hugging Face Hub or a URL to
-                a deployed Inference Endpoint. If not provided, the default recommended conversational model will be used.
-                Defaults to None.
+        # Strip trailing /
+        model_url = model_url.rstrip("/")
 
-        Returns:
-            `Dict`: The generated conversational output.
+        # Append /chat/completions if not already present
+        if model_url.endswith("/v1"):
+            model_url += "/chat/completions"
 
-        Raises:
-            [`InferenceTimeoutError`]:
-                If the model is unavailable or the request times out.
-            `aiohttp.ClientResponseError`:
-                If the request fails with an HTTP error status code other than HTTP 503.
+        # Append /v1/chat/completions if not already present
+        if not model_url.endswith("/chat/completions"):
+            model_url += "/v1/chat/completions"
 
-        Example:
-        ```py
-        # Must be run in an async context
-        >>> from huggingface_hub import AsyncInferenceClient
-        >>> client = AsyncInferenceClient()
-        >>> output = await client.conversational("Hi, who are you?")
-        >>> output
-        {'generated_text': 'I am the one who knocks.', 'conversation': {'generated_responses': ['I am the one who knocks.'], 'past_user_inputs': ['Hi, who are you?']}, 'warnings': ['Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.']}
-        >>> await client.conversational(
-        ...     "Wow, that's scary!",
-        ...     generated_responses=output["conversation"]["generated_responses"],
-        ...     past_user_inputs=output["conversation"]["past_user_inputs"],
-        ... )
-        ```
-        """
-        warnings.warn(
-            "'InferenceClient.conversational' is deprecated and will be removed starting from huggingface_hub>=0.25. "
-            "Please use the more appropriate 'InferenceClient.chat_completion' API instead.",
-            FutureWarning,
-        )
-        payload: Dict[str, Any] = {"inputs": {"text": text}}
-        if generated_responses is not None:
-            payload["inputs"]["generated_responses"] = generated_responses
-        if past_user_inputs is not None:
-            payload["inputs"]["past_user_inputs"] = past_user_inputs
-        if parameters is not None:
-            payload["parameters"] = parameters
-        response = await self.post(json=payload, model=model, task="conversational")
-        return _bytes_to_dict(response)  # type: ignore
+        return model_url
 
     async def document_question_answering(
         self,
@@ -1373,8 +1340,8 @@ class AsyncInferenceClient:
                     models_by_task.setdefault(model["task"], []).append(model["model_id"])
 
         async def _fetch_framework(framework: str) -> None:
-            async with _import_aiohttp().ClientSession(headers=self.headers) as client:
-                response = await client.get(f"{INFERENCE_ENDPOINT}/framework/{framework}")
+            async with self._get_client_session() as client:
+                response = await client.get(f"{INFERENCE_ENDPOINT}/framework/{framework}", proxy=self.proxies)
                 response.raise_for_status()
                 _unpack_response(framework, await response.json())
 
@@ -1757,7 +1724,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1786,7 +1754,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1815,7 +1784,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1844,7 +1814,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1873,7 +1844,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1901,7 +1873,8 @@ class AsyncInferenceClient:
         repetition_penalty: Optional[float] = None,
         return_full_text: Optional[bool] = False,  # Manual default value
         seed: Optional[int] = None,
-        stop_sequences: Optional[List[str]] = None,  # Same as `stop`
+        stop: Optional[List[str]] = None,
+        stop_sequences: Optional[List[str]] = None,  # Deprecated, use `stop` instead
         temperature: Optional[float] = None,
         top_k: Optional[int] = None,
         top_n_tokens: Optional[int] = None,
@@ -1966,8 +1939,10 @@ class AsyncInferenceClient:
                 Whether to prepend the prompt to the generated text
             seed (`int`, *optional*):
                 Random sampling seed
+            stop (`List[str]`, *optional*):
+                Stop generating tokens if a member of `stop` is generated.
             stop_sequences (`List[str]`, *optional*):
-                Stop generating tokens if a member of `stop_sequences` is generated
+                Deprecated argument. Use `stop` instead.
             temperature (`float`, *optional*):
                 The value used to module the logits distribution.
             top_n_tokens (`int`, *optional*):
@@ -2112,6 +2087,15 @@ class AsyncInferenceClient:
             )
             decoder_input_details = False
 
+        if stop_sequences is not None:
+            warnings.warn(
+                "`stop_sequences` is a deprecated argument for `text_generation` task"
+                " and will be removed in version '0.28.0'. Use `stop` instead.",
+                FutureWarning,
+            )
+        if stop is None:
+            stop = stop_sequences  # use deprecated arg if provided
+
         # Build payload
         parameters = {
             "adapter_id": adapter_id,
@@ -2125,7 +2109,7 @@ class AsyncInferenceClient:
             "repetition_penalty": repetition_penalty,
             "return_full_text": return_full_text,
             "seed": seed,
-            "stop": stop_sequences if stop_sequences is not None else [],
+            "stop": stop if stop is not None else [],
             "temperature": temperature,
             "top_k": top_k,
             "top_n_tokens": top_n_tokens,
@@ -2195,7 +2179,7 @@ class AsyncInferenceClient:
                     repetition_penalty=repetition_penalty,
                     return_full_text=return_full_text,
                     seed=seed,
-                    stop_sequences=stop_sequences,
+                    stop=stop,
                     temperature=temperature,
                     top_k=top_k,
                     top_n_tokens=top_n_tokens,
@@ -2655,6 +2639,47 @@ class AsyncInferenceClient:
         )
         return ZeroShotImageClassificationOutputElement.parse_obj_as_list(response)
 
+    def _get_client_session(self, headers: Optional[Dict] = None) -> "ClientSession":
+        aiohttp = _import_aiohttp()
+        client_headers = self.headers.copy()
+        if headers is not None:
+            client_headers.update(headers)
+
+        # Return a new aiohttp ClientSession with correct settings.
+        session = aiohttp.ClientSession(
+            headers=client_headers,
+            cookies=self.cookies,
+            timeout=aiohttp.ClientTimeout(self.timeout),
+            trust_env=self.trust_env,
+        )
+
+        # Keep track of sessions to close them later
+        self._sessions[session] = set()
+
+        # Override the `._request` method to register responses to be closed
+        session._wrapped_request = session._request
+
+        async def _request(method, url, **kwargs):
+            response = await session._wrapped_request(method, url, **kwargs)
+            self._sessions[session].add(response)
+            return response
+
+        session._request = _request
+
+        # Override the 'close' method to
+        # 1. close ongoing responses
+        # 2. deregister the session when closed
+        session._close = session.close
+
+        async def close_session():
+            for response in self._sessions[session]:
+                response.close()
+            await session._close()
+            self._sessions.pop(session, None)
+
+        session.close = close_session
+        return session
+
     def _resolve_url(self, model: Optional[str] = None, task: Optional[str] = None) -> str:
         model = model or self.model or self.base_url
 
@@ -2761,8 +2786,8 @@ class AsyncInferenceClient:
         else:
             url = f"{INFERENCE_ENDPOINT}/models/{model}/info"
 
-        async with _import_aiohttp().ClientSession(headers=self.headers) as client:
-            response = await client.get(url)
+        async with self._get_client_session() as client:
+            response = await client.get(url, proxy=self.proxies)
             response.raise_for_status()
             return await response.json()
 
@@ -2798,8 +2823,8 @@ class AsyncInferenceClient:
             )
         url = model.rstrip("/") + "/health"
 
-        async with _import_aiohttp().ClientSession(headers=self.headers) as client:
-            response = await client.get(url)
+        async with self._get_client_session() as client:
+            response = await client.get(url, proxy=self.proxies)
             return response.status == 200
 
     async def get_model_status(self, model: Optional[str] = None) -> ModelStatus:
@@ -2840,8 +2865,8 @@ class AsyncInferenceClient:
             raise NotImplementedError("Model status is only available for Inference API endpoints.")
         url = f"{INFERENCE_ENDPOINT}/status/{model}"
 
-        async with _import_aiohttp().ClientSession(headers=self.headers) as client:
-            response = await client.get(url)
+        async with self._get_client_session() as client:
+            response = await client.get(url, proxy=self.proxies)
             response.raise_for_status()
             response_data = await response.json()
 

huggingface_hub/inference/_generated/types/base.py

@@ -15,7 +15,6 @@
 
 import inspect
 import json
-import warnings
 from dataclasses import asdict, dataclass
 from typing import Any, Dict, List, Type, TypeVar, Union, get_args
 
@@ -135,14 +134,6 @@ class BaseInferenceType(dict):
             self[__name] = __value
         return
 
-    def __getitem__(self, __key: Any) -> Any:
-        warnings.warn(
-            f"Accessing '{self.__class__.__name__}' values through dict is deprecated and "
-            "will be removed from version '0.25'. Use dataclass attributes instead.",
-            FutureWarning,
-        )
-        return super().__getitem__(__key)
-
 
 def normalize_key(key: str) -> str:
     # e.g "content-type" -> "content_type", "Accept" -> "accept"

huggingface_hub/inference/_templating.py

@@ -1,9 +1,8 @@
 from functools import lru_cache
 from typing import Callable, Dict, List, Optional, Union
 
-from huggingface_hub.errors import TemplateError
-
-from ..utils import HfHubHTTPError, RepositoryNotFoundError, is_minijinja_available
+from ..errors import HfHubHTTPError, RepositoryNotFoundError, TemplateError
+from ..utils import is_minijinja_available
 
 
 def _import_minijinja():

huggingface_hub/inference_api.py

@@ -1,7 +1,7 @@
 import io
 from typing import Any, Dict, List, Optional, Union
 
-from .constants import INFERENCE_ENDPOINT
+from . import constants
 from .hf_api import HfApi
 from .utils import build_hf_headers, get_session, is_pillow_available, logging, validate_hf_hub_args
 from .utils._deprecation import _deprecate_method
@@ -149,7 +149,7 @@ class InferenceApi:
             assert model_info.pipeline_tag is not None, "Pipeline tag cannot be None"
             self.task = model_info.pipeline_tag
 
-        self.api_url = f"{INFERENCE_ENDPOINT}/pipeline/{self.task}/{repo_id}"
+        self.api_url = f"{constants.INFERENCE_ENDPOINT}/pipeline/{self.task}/{repo_id}"
 
     def __repr__(self):
         # Do not add headers to repr to avoid leaking token.

huggingface_hub/keras_mixin.py

@@ -16,7 +16,7 @@ from huggingface_hub.utils import (
     yaml_dump,
 )
 
-from .constants import CONFIG_NAME
+from . import constants
 from .hf_api import HfApi
 from .utils import SoftTemporaryDirectory, logging, validate_hf_hub_args
 from .utils._typing import CallableT
@@ -202,7 +202,7 @@ def save_pretrained_keras(
         if not isinstance(config, dict):
             raise RuntimeError(f"Provided config to save_pretrained_keras should be a dict. Got: '{type(config)}'")
 
-        with (save_directory / CONFIG_NAME).open("w") as f:
+        with (save_directory / constants.CONFIG_NAME).open("w") as f:
             json.dump(config, f)
 
     metadata = {}