mlrun 1.10.0rc18__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/datastore/model_provider/huggingface_provider.py

@@ -11,17 +11,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
-from typing import TYPE_CHECKING, Optional, TypeVar, Union
+import threading
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import mlrun
-from mlrun.datastore.model_provider.model_provider import ModelProvider
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
 
 if TYPE_CHECKING:
     from transformers.pipelines.base import Pipeline
-
-T = TypeVar("T")
-ChatType = list[dict[str, str]]  # according to transformers.pipelines.text_generation
+    from transformers.pipelines.text_generation import ChatType
 
 
 class HuggingFaceProvider(ModelProvider):
@@ -34,8 +36,14 @@ class HuggingFaceProvider(ModelProvider):
     This class extends the ModelProvider base class and implements Hugging Face-specific
     functionality, including pipeline initialization, default text generation operations,
     and custom operations tailored to the Hugging Face Transformers pipeline API.
+
+    Note: The pipeline object will download the model (if not already cached) and load it
+    into memory for inference. Ensure you have the required CPU/GPU and memory to use this operation.
     """
 
+    #  locks for threading use cases
+    _client_lock = threading.Lock()
+
     def __init__(
         self,
         parent,
@@ -60,18 +68,20 @@ class HuggingFaceProvider(ModelProvider):
         )
         self.options = self.get_client_options()
         self._expected_operation_type = None
-        self.load_client()
+        self._download_model()
 
     @staticmethod
-    def _extract_string_output(result) -> str:
+    def _extract_string_output(response: list[dict]) -> str:
         """
-        Extracts the first generated string from Hugging Face pipeline output,
-        regardless of whether it's plain text-generation or chat-style output.
+        Extracts the first generated string from Hugging Face pipeline output
         """
-        if not isinstance(result, list) or len(result) == 0:
+        if not isinstance(response, list) or len(response) == 0:
             raise ValueError("Empty or invalid pipeline output")
-
-        return result[0].get("generated_text")
+        if len(response) != 1:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "HuggingFaceProvider: extracting string from response is only supported for single-response outputs"
+            )
+        return response[0].get("generated_text")
 
     @classmethod
     def parse_endpoint_and_path(cls, endpoint, subpath) -> (str, str):
@@ -81,6 +91,120 @@ class HuggingFaceProvider(ModelProvider):
             subpath = ""
         return endpoint, subpath
 
+    @property
+    def client(self) -> Any:
+        """
+        Lazily return the HuggingFace-pipeline client.
+
+        If the client has not been initialized yet, it will be created
+        by calling `load_client`.
+        """
+        self.load_client()
+        return self._client
+
+    def _download_model(self):
+        """
+        Pre-downloads model files locally to prevent race conditions in multiprocessing.
+
+        Uses snapshot_download with local_dir_use_symlinks=False to ensure proper
+        file copying for safe concurrent access across multiple processes.
+
+        :raises:
+            ImportError: If huggingface_hub package is not installed.
+        """
+        try:
+            from huggingface_hub import snapshot_download
+
+            # Download the model and tokenizer files directly to the cache.
+            snapshot_download(
+                repo_id=self.model,
+                local_dir_use_symlinks=False,
+                token=self._get_secret_or_env("HF_TOKEN") or None,
+            )
+        except ImportError as exc:
+            raise ImportError("huggingface_hub package is not installed") from exc
+
+    def _response_handler(
+        self,
+        response: Union[str, list],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        messages: Union[str, list[str], "ChatType", list["ChatType"]] = None,
+        **kwargs,
+    ) -> Union[str, list, dict[str, Any]]:
+        """
+        Processes and formats the raw response from the HuggingFace pipeline according to the specified format.
+
+        The response should exclude the user’s input (no repetition in the output).
+        This can be accomplished by invoking the pipeline with `return_full_text=False`.
+
+        :param response:                The raw response from the HuggingFace pipeline, typically a list of dictionaries
+                                        containing generated text sequences.
+        :param invoke_response_format:  Determines how the response should be processed and returned. Options:
+
+                                       - STRING: Return only the main generated content as a string,
+                                                 for single-answer responses.
+                                       - USAGE: Return a dictionary combining the string response with
+                                                token usage statistics:
+
+                                       .. code-block:: json
+
+                                       {
+                                           "answer": "<generated_text>",
+                                           "usage": {
+                                               "prompt_tokens": <int>,
+                                               "completion_tokens": <int>,
+                                               "total_tokens": <int>
+                                           }
+                                       }
+
+                                       Note: Token counts are estimated after answer generation and
+                                       may differ from the actual tokens generated by the model due to
+                                       internal decoding behavior and implementation details.
+
+                                       - FULL: Return the full raw response object.
+
+        :param messages:               The original input messages used for token count estimation in USAGE mode.
+                                       Can be a string, list of strings, or chat format messages.
+        :param kwargs:                 Additional parameters for response processing.
+
+        :return:                       The processed response in the format specified by `invoke_response_format`.
+                                       Can be a string, dictionary, or the original response object.
+
+        :raises MLRunInvalidArgumentError: If extracting the string response fails.
+        :raises MLRunRuntimeError: If applying the chat template to the model fails during token usage calculation.
+        """
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
+            str_response = self._extract_string_output(response)
+            if invoke_response_format == InvokeResponseFormat.STRING:
+                return str_response
+            if invoke_response_format == InvokeResponseFormat.USAGE:
+                tokenizer = self.client.tokenizer
+                if not isinstance(messages, str):
+                    try:
+                        messages = tokenizer.apply_chat_template(
+                            messages, tokenize=False, add_generation_prompt=True
+                        )
+                    except Exception as e:
+                        raise mlrun.errors.MLRunRuntimeError(
+                            f"Failed to apply chat template using the tokenizer for model '{self.model}'. "
+                            "This may indicate that the tokenizer does not support chat formatting, "
+                            "or that the input format is invalid. "
+                            f"Original error: {e}"
+                        )
+                prompt_tokens = len(tokenizer.encode(messages))
+                completion_tokens = len(tokenizer.encode(str_response))
+                total_tokens = prompt_tokens + completion_tokens
+                usage = {
+                    "prompt_tokens": prompt_tokens,
+                    "completion_tokens": completion_tokens,
+                    "total_tokens": total_tokens,
+                }
+                response = {
+                    UsageResponseKeys.ANSWER: str_response,
+                    UsageResponseKeys.USAGE: usage,
+                }
+        return response
+
     def load_client(self) -> None:
         """
         Initializes the Hugging Face pipeline using the provided options.
@@ -91,15 +215,20 @@ class HuggingFaceProvider(ModelProvider):
 
         Note: Hugging Face pipelines are synchronous and do not support async invocation.
 
-        Raises:
+        :raises:
             ImportError: If the `transformers` package is not installed.
         """
+        if self._client:
+            return
         try:
             from transformers import pipeline, AutoModelForCausalLM  # noqa
             from transformers import AutoTokenizer  # noqa
             from transformers.pipelines.base import Pipeline  # noqa
 
-            self._client = pipeline(model=self.model, **self.options)
+            self.options["model_kwargs"] = self.options.get("model_kwargs", {})
+            self.options["model_kwargs"]["local_files_only"] = True
+            with self._client_lock:
+                self._client = pipeline(model=self.model, **self.options)
             self._expected_operation_type = Pipeline
         except ImportError as exc:
             raise ImportError("transformers package is not installed") from exc
@@ -117,25 +246,40 @@ class HuggingFaceProvider(ModelProvider):
 
     def custom_invoke(
         self, operation: Optional["Pipeline"] = None, **invoke_kwargs
-    ) -> Optional[T]:
+    ) -> Union[list, dict, Any]:
         """
-        HuggingFace implementation of `ModelProvider.custom_invoke`.
-        Use the default config in provider client/ user defined client:
+        Invokes a HuggingFace pipeline operation with the given keyword arguments.
+
+        This method provides flexibility to use a custom pipeline object for specific tasks
+        (e.g., image classification, sentiment analysis).
+
+        The operation must be a Pipeline object from the transformers library that accepts keyword arguments.
 
         Example:
-        ```python
+            ```python
+            from transformers import pipeline
+            from PIL import Image
+
+            # Using custom pipeline for image classification
             image = Image.open(image_path)
-            pipeline_object =  pipeline("image-classification", model="microsoft/resnet-50")
+            pipeline_object = pipeline("image-classification", model="microsoft/resnet-50")
             result = hf_provider.custom_invoke(
                 pipeline_object,
                 inputs=image,
             )
-        ```
+            ```
 
+        :param operation:      A Pipeline object from the transformers library.
+                               If not provided, defaults to the provider's configured pipeline.
+        :param invoke_kwargs:  Keyword arguments to pass to the pipeline operation.
+                               These are merged with `default_invoke_kwargs` and may include
+                               parameters such as `inputs`, `max_length`, `temperature`, or task-specific options.
 
-        :param operation:               A pipeline object
-        :param invoke_kwargs:           Keyword arguments to pass to the operation.
-        :return:                        The full response returned by the operation.
+        :return:               The full response returned by the pipeline operation.
+                               Format depends on the pipeline task (list for text generation,
+                               dict for classification, etc.).
+
+        :raises MLRunInvalidArgumentError: If the operation is not a valid Pipeline object.
 
         """
         invoke_kwargs = self.get_invoke_kwargs(invoke_kwargs)
@@ -150,34 +294,74 @@ class HuggingFaceProvider(ModelProvider):
 
     def invoke(
         self,
-        messages: Union[str, list[str], ChatType, list[ChatType]] = None,
-        as_str: bool = False,
+        messages: Union[str, list[str], "ChatType", list["ChatType"]],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
         **invoke_kwargs,
-    ) -> Optional[Union[str, list, T]]:
+    ) -> Union[str, list, dict[str, Any]]:
         """
-        HuggingFace-specific implementation of `ModelProvider.invoke`.
-        Invokes a HuggingFace model operation using the synchronous client.
-        For complete usage details, refer to `ModelProvider.invoke`.
+        HuggingFace-specific implementation of model invocation using the synchronous pipeline client.
+        Invokes a HuggingFace model operation for text generation tasks.
+
+        Note: Ensure your environment has sufficient computational resources (CPU/GPU and memory) to run the model.
+
         :param messages:
-                            Same as ModelProvider.invoke.
+            Input for the text generation model. Can be provided in multiple formats:
+
+            - A single string: Direct text input for generation
+            - A list of strings: Multiple text inputs for batch processing
+            - Chat format: A list of dictionaries with "role" and "content" keys:
+
+            .. code-block:: json
+
+                [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "What is the capital of France?"}
+                ]
+
+        :param invoke_response_format: InvokeResponseFormat
+            Specifies the format of the returned response. Options:
 
-        :param as_str:
-                            If `True`, returns only the main content from a single response
-                            (intended for single-response use cases).
-                            If `False`, returns the full response object, whose type depends on
-                            the client (e.g., `pipeline`).
+            - "string": Returns only the generated text content, extracted from a single response.
+            - "usage":  Combines the generated text with metadata (e.g., token usage), returning a dictionary:
+
+            .. code-block:: json
+                {
+                    "answer": "<generated_text>",
+                    "usage": {
+                        "prompt_tokens": <int>,
+                        "completion_tokens": <int>,
+                        "total_tokens": <int>
+                    }
+                }
+
+            Note: For usage mode, the model tokenizer should support apply_chat_template.
+
+            - "full":   Returns the raw response object from the HuggingFace model,
+                        typically a list of generated sequences (dictionaries).
+                        This format does not include token usage statistics.
 
         :param invoke_kwargs:
-                            Same as ModelProvider.invoke.
-        :return:            Same as ModelProvider.invoke.
+            Additional keyword arguments passed to the HuggingFace pipeline.
+
+        :return:
+            A string, dictionary, or list of model outputs, depending on `invoke_response_format`.
+
+        :raises MLRunInvalidArgumentError:
+            If the pipeline task is not "text-generation" or if the response contains multiple outputs when extracting
+            string content.
+        :raises MLRunRuntimeError:
+            If using "usage" response mode and the model tokenizer does not support chat template formatting.
         """
         if self.client.task != "text-generation":
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "HuggingFaceProvider.invoke supports text-generation task only"
             )
-        if as_str:
+        if InvokeResponseFormat.is_str_response(invoke_response_format.value):
             invoke_kwargs["return_full_text"] = False
         response = self.custom_invoke(text_inputs=messages, **invoke_kwargs)
-        if as_str:
-            return self._extract_string_output(response)
+        response = self._response_handler(
+            messages=messages,
+            response=response,
+            invoke_response_format=invoke_response_format,
+        )
         return response

mlrun/datastore/model_provider/mock_model_provider.py

@@ -0,0 +1,87 @@
+# Copyright 2023 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Optional, Union
+
+import mlrun
+from mlrun.datastore.model_provider.model_provider import (
+    InvokeResponseFormat,
+    ModelProvider,
+    UsageResponseKeys,
+)
+
+
+class MockModelProvider(ModelProvider):
+    support_async = False
+
+    def __init__(
+        self,
+        parent,
+        kind,
+        name,
+        endpoint="",
+        secrets: Optional[dict] = None,
+        default_invoke_kwargs: Optional[dict] = None,
+    ):
+        super().__init__(
+            parent=parent, name=name, kind=kind, endpoint=endpoint, secrets=secrets
+        )
+        self.default_invoke_kwargs = default_invoke_kwargs or {}
+        self._client = None
+        self._async_client = None
+
+    @staticmethod
+    def _extract_string_output(response: Any) -> str:
+        """
+        Extracts string response from response object
+        """
+        pass
+
+    def load_client(self) -> None:
+        """
+        Initializes the SDK client for the model provider with the given keyword arguments
+        and assigns it to an instance attribute (e.g., self._client).
+
+        Subclasses should override this method to:
+        - Create and configure the provider-specific client instance.
+        - Assign the client instance to self._client.
+        """
+
+        pass
+
+    def invoke(
+        self,
+        messages: Union[list[dict], Any],
+        invoke_response_format: InvokeResponseFormat = InvokeResponseFormat.FULL,
+        **invoke_kwargs,
+    ) -> Union[str, dict[str, Any], Any]:
+        if invoke_response_format == InvokeResponseFormat.STRING:
+            return (
+                "You are using a mock model provider, no actual inference is performed."
+            )
+        elif invoke_response_format == InvokeResponseFormat.FULL:
+            return {
+                UsageResponseKeys.USAGE: {"prompt_tokens": 0, "completion_tokens": 0},
+                UsageResponseKeys.ANSWER: "You are using a mock model provider, no actual inference is performed.",
+                "extra": {},
+            }
+        elif invoke_response_format == InvokeResponseFormat.USAGE:
+            return {
+                UsageResponseKeys.ANSWER: "You are using a mock model provider, no actual inference is performed.",
+                UsageResponseKeys.USAGE: {"prompt_tokens": 0, "completion_tokens": 0},
+            }
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"Unsupported invoke response format: {invoke_response_format}"
+            )