snowflake-ml-python 1.7.0__py3-none-any.whl → 1.7.1__py3-none-any.whl

snowflake/cortex/__init__.py

@@ -3,6 +3,7 @@ from snowflake.cortex._complete import Complete, CompleteOptions
 from snowflake.cortex._embed_text_768 import EmbedText768
 from snowflake.cortex._embed_text_1024 import EmbedText1024
 from snowflake.cortex._extract_answer import ExtractAnswer
+from snowflake.cortex._finetune import Finetune, FinetuneJob, FinetuneStatus
 from snowflake.cortex._sentiment import Sentiment
 from snowflake.cortex._summarize import Summarize
 from snowflake.cortex._translate import Translate
@@ -14,6 +15,9 @@ __all__ = [
     "EmbedText768",
     "EmbedText1024",
     "ExtractAnswer",
+    "Finetune",
+    "FinetuneJob",
+    "FinetuneStatus",
     "Sentiment",
     "Summarize",
     "Translate",

snowflake/cortex/_complete.py

@@ -1,7 +1,8 @@
 import json
 import logging
 import time
-from typing import Any, Callable, Iterator, List, Optional, TypedDict, Union, cast
+from io import BytesIO
+from typing import Any, Callable, Dict, Iterator, List, Optional, TypedDict, Union, cast
 from urllib.parse import urlunparse
 
 import requests
@@ -16,8 +17,10 @@ from snowflake.cortex._util import (
 )
 from snowflake.ml._internal import telemetry
 from snowflake.snowpark import context, functions
+from snowflake.snowpark._internal.utils import is_in_stored_procedure
 
 logger = logging.getLogger(__name__)
+_REST_COMPLETE_URL = "/api/v2/cortex/inference:complete"
 
 
 class ConversationMessage(TypedDict):
@@ -84,6 +87,76 @@ def retry(func: Callable[..., requests.Response]) -> Callable[..., requests.Resp
     return inner
 
 
+def _make_common_request_headers() -> Dict[str, str]:
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json, text/event-stream",
+    }
+    return headers
+
+
+def _make_request_body(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+) -> Dict[str, Any]:
+    data = {
+        "model": model,
+        "stream": True,
+    }
+    if isinstance(prompt, List):
+        data["messages"] = prompt
+    else:
+        data["messages"] = [{"content": prompt}]
+
+    if options:
+        if "max_tokens" in options:
+            data["max_tokens"] = options["max_tokens"]
+            data["max_output_tokens"] = options["max_tokens"]
+        if "temperature" in options:
+            data["temperature"] = options["temperature"]
+        if "top_p" in options:
+            data["top_p"] = options["top_p"]
+    return data
+
+
+# XP endpoint returns a dict response which needs to be converted to a format which can
+# be consumed by the SSEClient. This method does that.
+def _xp_dict_to_response(raw_resp: Dict[str, Any]) -> requests.Response:
+    response = requests.Response()
+    response.status_code = int(raw_resp["status"])
+    response.headers = raw_resp["headers"]
+
+    data = raw_resp["content"]
+    data = json.loads(data)
+    # Convert the dictionary to a string format that resembles the SSE event format
+    # For example, if the dict is {'event': 'message', 'data': 'your data'}, it should be formatted like this:
+    sse_format_data = ""
+    for event in data:
+        event_type = event.get("event", "message")
+        event_data = event.get("data", "")
+        event_data = json.dumps(event_data)
+        sse_format_data += f"event: {event_type}\ndata: {event_data}\n\n"  # Add each event with new lines
+
+    response.raw = BytesIO(sse_format_data.encode("utf-8"))
+    return response
+
+
+@retry
+def _call_complete_xp(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+    deadline: Optional[float] = None,
+) -> requests.Response:
+    headers = _make_common_request_headers()
+    body = _make_request_body(model, prompt, options)
+    import _snowflake
+
+    raw_resp = _snowflake.send_snow_api_request("POST", _REST_COMPLETE_URL, {}, headers, body, {}, deadline)
+    return _xp_dict_to_response(raw_resp)
+
+
 @retry
 def _call_complete_rest(
     model: str,
@@ -110,36 +183,16 @@ def _call_complete_rest(
     scheme = "https"
     if hasattr(session.connection, "scheme"):
         scheme = session.connection.scheme
-    url = urlunparse((scheme, session.connection.host, "api/v2/cortex/inference:complete", "", "", ""))
+    url = urlunparse((scheme, session.connection.host, _REST_COMPLETE_URL, "", "", ""))
 
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f'Snowflake Token="{session.connection.rest.token}"',
-        "Accept": "application/json, text/event-stream",
-    }
-
-    data = {
-        "model": model,
-        "stream": True,
-    }
-    if isinstance(prompt, List):
-        data["messages"] = prompt
-    else:
-        data["messages"] = [{"content": prompt}]
-
-    if options:
-        if "max_tokens" in options:
-            data["max_tokens"] = options["max_tokens"]
-            data["max_output_tokens"] = options["max_tokens"]
-        if "temperature" in options:
-            data["temperature"] = options["temperature"]
-        if "top_p" in options:
-            data["top_p"] = options["top_p"]
+    headers = _make_common_request_headers()
+    headers["Authorization"] = f'Snowflake Token="{session.connection.rest.token}"'
 
+    body = _make_request_body(model, prompt, options)
     logger.debug(f"making POST request to {url} (model={model})")
     return requests.post(
         url,
-        json=data,
+        json=body,
         headers=headers,
         stream=True,
     )
@@ -164,49 +217,24 @@ def _complete_call_sql_function_snowpark(
     return cast(snowpark.Column, functions.builtin(function)(*args))
 
 
-def _complete_call_sql_function_immediate(
-    function: str,
+def _complete_non_streaming_immediate(
     model: str,
     prompt: Union[str, List[ConversationMessage]],
     options: Optional[CompleteOptions],
-    session: Optional[snowpark.Session],
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
 ) -> str:
-    session = session or context.get_active_session()
-    if session is None:
-        raise SnowflakeAuthenticationException(
-            """Session required. Provide the session through a session=... argument or ensure an active session is
-            available in your environment."""
-        )
+    response = _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    return "".join(response)
 
-    # https://docs.snowflake.com/en/sql-reference/functions/complete-snowflake-cortex
-    if options is not None or not isinstance(prompt, str):
-        if isinstance(prompt, List):
-            prompt_arg = prompt
-        else:
-            prompt_arg = [{"role": "user", "content": prompt}]
-        options = options or {}
-        lit_args = [
-            functions.lit(model),
-            functions.lit(prompt_arg),
-            functions.lit(options),
-        ]
-    else:
-        lit_args = [
-            functions.lit(model),
-            functions.lit(prompt),
-        ]
-
-    empty_df = session.create_dataframe([snowpark.Row()])
-    df = empty_df.select(functions.builtin(function)(*lit_args))
-    return cast(str, df.collect()[0][0])
 
-
-def _complete_sql_impl(
+def _complete_non_streaming_impl(
     function: str,
     model: Union[str, snowpark.Column],
     prompt: Union[str, List[ConversationMessage], snowpark.Column],
     options: Optional[Union[CompleteOptions, snowpark.Column]],
-    session: Optional[snowpark.Session],
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
 ) -> Union[str, snowpark.Column]:
     if isinstance(prompt, snowpark.Column):
         if options is not None:
@@ -217,7 +245,24 @@ def _complete_sql_impl(
         raise ValueError("'model' cannot be a snowpark.Column when 'prompt' is a string.")
     if isinstance(options, snowpark.Column):
         raise ValueError("'options' cannot be a snowpark.Column when 'prompt' is a string.")
-    return _complete_call_sql_function_immediate(function, model, prompt, options, session)
+    return _complete_non_streaming_immediate(
+        model=model, prompt=prompt, options=options, session=session, deadline=deadline
+    )
+
+
+def _complete_rest(
+    model: str,
+    prompt: Union[str, List[ConversationMessage]],
+    options: Optional[CompleteOptions] = None,
+    session: Optional[snowpark.Session] = None,
+    deadline: Optional[float] = None,
+) -> Iterator[str]:
+    if is_in_stored_procedure():  # type: ignore[no-untyped-call]
+        response = _call_complete_xp(model=model, prompt=prompt, options=options, deadline=deadline)
+    else:
+        response = _call_complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    assert response.status_code >= 200 and response.status_code < 300
+    return _return_stream_response(response, deadline)
 
 
 def _complete_impl(
@@ -239,10 +284,8 @@ def _complete_impl(
             raise ValueError("in REST mode, 'model' must be a string")
         if not isinstance(prompt, str) and not isinstance(prompt, List):
             raise ValueError("in REST mode, 'prompt' must be a string or a list of ConversationMessage")
-        response = _call_complete_rest(model, prompt, options, session=session, deadline=deadline)
-        assert response.status_code >= 200 and response.status_code < 300
-        return _return_stream_response(response, deadline)
-    return _complete_sql_impl(function, model, prompt, options, session)
+        return _complete_rest(model=model, prompt=prompt, options=options, session=session, deadline=deadline)
+    return _complete_non_streaming_impl(function, model, prompt, options, session, deadline)
 
 
 @telemetry.send_api_usage_telemetry(

snowflake/cortex/_finetune.py

@@ -0,0 +1,273 @@
+import json
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union, cast
+
+from snowflake import snowpark
+from snowflake.cortex._util import (
+    CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+    call_sql_function_literals,
+)
+from snowflake.ml._internal import telemetry
+from snowflake.ml._internal.utils import snowpark_dataframe_utils
+
+_CORTEX_FINETUNE_SYSTEM_FUNCTION_NAME = "SNOWFLAKE.CORTEX.FINETUNE"
+CORTEX_FINETUNE_TELEMETRY_SUBPROJECT = "FINETUNE"
+CORTEX_FINETUNE_FIRST_VERSION = "1.7.0"
+CORTEX_FINETUNE_DOCUMENTATION_URL = "https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-finetuning"
+
+
+class FinetuneError(Exception):
+    def __init__(self, message: str, original_exception: Optional[Exception] = None) -> None:
+        """Finetuning Exception Class.
+
+        Args:
+            message: Error message to be reported.
+            original_exception: Original exception. This is the exception raised to users by telemetry.
+
+        Attributes:
+            original_exception: Original exception with an error code in its message.
+        """
+        self.original_exception = original_exception
+        self._pretty_msg = message + repr(self.original_exception) if self.original_exception is not None else ""
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._pretty_msg!r})"
+
+    def __str__(self) -> str:
+        return self._pretty_msg
+
+
+@dataclass
+class FinetuneStatus:
+    """Fine-tuning job status."""
+
+    id: Optional[str] = None
+    """Workflow ID for the fine-tuning run."""
+
+    status: Optional[str] = None
+    """Status string, e.g. PENDING, RUNNING, SUCCESS, ERROR, CANCELLED."""
+
+    base_model: Optional[str] = None
+    """Name of the base model that is being fine-tuned."""
+
+    created_on: Optional[int] = None
+    """Creation timestamp of the Fine-tuning job in milliseconds."""
+
+    error: Optional[Dict[str, Any]] = None
+    """Error message propagated from the job."""
+
+    finished_on: Optional[int] = None
+    """Completion timestamp of the Fine-tuning job in milliseconds."""
+
+    progress: Optional[float] = None
+    """Progress made as a fraction of total [0.0,1.0]."""
+
+    training_result: Optional[List[Dict[str, Any]]] = None
+    """Detailed metrics report for a completed training."""
+
+    trained_tokens: Optional[int] = None
+    """Number of tokens trained on. If multiple epochs are run, this can be larger than number of tokens in the
+    training data."""
+
+    training_data: Optional[str] = None
+    """Training data query."""
+
+    validation_data: Optional[str] = None
+    """Validation data query."""
+
+    model: Optional[str] = None
+    """Location of the fine-tuned model."""
+
+
+class FinetuneJob:
+    def __init__(self, session: Optional[snowpark.Session], status: FinetuneStatus) -> None:
+        """Fine-tuning Job.
+
+        Args:
+            session: Snowpark session to use to communicate with Snowflake.
+            status: FinetuneStatus for this job.
+        """
+        self._session = session
+        self.status = status
+
+    def __repr__(self) -> str:
+        return self.status.__repr__()
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, FinetuneJob):
+            raise NotImplementedError(
+                f"Equality comparison of FinetuneJob with objects of type {type(other)} is not implemented."
+            )
+        return self.status == other.status
+
+    @snowpark._internal.utils.experimental(version=CORTEX_FINETUNE_FIRST_VERSION)
+    @telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+        subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
+    )
+    def cancel(self) -> bool:
+        """Cancel a fine-tuning run.
+
+        No confirmation will be required.
+
+        [Documentation](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-finetuning)
+
+        Args:
+
+        Returns:
+            True if the cancellation was successful, False otherwise.
+        """
+        result = _finetune_impl(operation="CANCEL", session=self._session, function_args=[self.status.id])
+        return result is not None and isinstance(result, str) and result.startswith("Canceled Cortex Fine-tuning job: ")
+
+    @snowpark._internal.utils.experimental(version=CORTEX_FINETUNE_FIRST_VERSION)
+    @telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+        subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
+    )
+    def describe(self) -> FinetuneStatus:
+        """Describe a fine-tuning run.
+
+        Args:
+
+        Returns:
+            FinetuneStatus containing of attributes of the fine-tuning run.
+        """
+        result_string = _finetune_impl(operation="DESCRIBE", session=self._session, function_args=[self.status.id])
+
+        result = FinetuneStatus(**cast(Dict[str, Any], _try_load_json(result_string)))
+        return result
+
+
+class Finetune:
+    @snowpark._internal.utils.experimental(version=CORTEX_FINETUNE_FIRST_VERSION)
+    @telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+        subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
+    )
+    def __init__(self, session: Optional[snowpark.Session] = None) -> None:
+        """Cortex Fine-Tuning API.
+
+        [Documentation](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-finetuning)
+
+        Args:
+            session: Snowpark session to be used. If none is given, we will attempt to
+                use the currently active session.
+        """
+        self._session = session
+
+    @snowpark._internal.utils.experimental(version=CORTEX_FINETUNE_FIRST_VERSION)
+    @telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+        subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
+    )
+    def create(
+        self,
+        name: str,
+        base_model: str,
+        training_data: Union[str, snowpark.DataFrame],
+        validation_data: Optional[Union[str, snowpark.DataFrame]] = None,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> FinetuneJob:
+        """Create a new fine-tuning runs.
+
+        The expected format of training and validation data is two fields or columns:
+            "prompt": the input to the model and
+            "completion": the output that the model is expected to generate.
+
+        Both data parameters "training_data" and "validation_data" expect to be one of
+            (1) stage path to JSONL-formatted data,
+            (2) select-query string resulting in a table,
+            (3) Snowpark DataFrame containing the data
+
+        [Documentation](https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-finetuning)
+
+        Args:
+            name: Name of the resulting fine-tuned model.
+            base_model: The name of the base model to start fine-tuning from.
+            training_data: Data used for fine-tuning the model.
+            validation_data: Data used for validating the fine-tuned model (not used in training)
+            options: Dictionary of additional options to be passed to the training procedure.
+                Please refer to the official documentation for a list of available options.
+
+        Returns:
+            The identifier of the fine-tuning run.
+
+        Raises:
+            ValueError: If the Snowpark DataFrame used is incompatible with this API.
+                This can happen if the DataFrame contains multiple queries.
+        """
+
+        # Handle data provided as snowpark dataframes
+        if isinstance(training_data, snowpark.DataFrame):
+            if snowpark_dataframe_utils.is_single_query_snowpark_dataframe(training_data):
+                training_string = str(training_data.queries["queries"][0])
+            else:
+                raise ValueError(
+                    "Snowpark DataFrame given in 'training_data' contains "
+                    + f'{training_data.queries["queries"]} queries and '
+                    + f'{training_data.queries["post_actions"]} post_actions. It needs '
+                    "to contain exactly one query and no post_actions."
+                )
+        else:
+            training_string = training_data
+
+        validation_string: Optional[str] = None
+        if isinstance(validation_data, snowpark.DataFrame):
+            if snowpark_dataframe_utils.is_single_query_snowpark_dataframe(validation_data):
+                validation_string = str(validation_data.queries["queries"][0])
+            else:
+                raise ValueError(
+                    "Snowpark DataFrame given in 'validation_data' contains "
+                    + f'{validation_data.queries["queries"]} queries and '
+                    + f'{validation_data.queries["post_actions"]} post_actions. It needs '
+                    "to contain exactly one query and no post_actions."
+                )
+        else:
+            validation_string = validation_data
+
+        result = _finetune_impl(
+            operation="CREATE",
+            session=self._session,
+            function_args=[name, base_model, training_string, validation_string, options],
+        )
+        finetune_status = FinetuneStatus(id=result)
+        finetune_run = FinetuneJob(self._session, finetune_status)
+        return finetune_run
+
+    @snowpark._internal.utils.experimental(version=CORTEX_FINETUNE_FIRST_VERSION)
+    @telemetry.send_api_usage_telemetry(
+        project=CORTEX_FUNCTIONS_TELEMETRY_PROJECT,
+        subproject=CORTEX_FINETUNE_TELEMETRY_SUBPROJECT,
+    )
+    def list_jobs(self) -> List["FinetuneJob"]:
+        """Show current and past fine-tuning runs.
+
+        Returns:
+            List of dictionaries of attributes of the fine-tuning runs. Please refer to the official documentation for a
+                list of expected fields.
+        """
+        result_string = _finetune_impl(operation="SHOW", session=self._session, function_args=[])
+        result = _try_load_json(result_string)
+
+        return [FinetuneJob(session=self._session, status=FinetuneStatus(**run_status)) for run_status in result]
+
+
+def _try_load_json(json_string: str) -> Union[Dict[Any, Any], List[Any]]:
+    try:
+        result = json.loads(str(json_string))
+    except json.JSONDecodeError as e:
+        message = f"""Unable to parse JSON from: "{json_string}". """
+        raise FinetuneError(message=message, original_exception=e)
+    except Exception as e:
+        message = f"""Unable to parse JSON from: "{json_string}". """
+        raise FinetuneError(message=message, original_exception=e)
+    else:
+        if not isinstance(result, dict) and not isinstance(result, list):
+            message = f"""Unable to parse JSON from: "{json_string}". Result was not a dictionary."""
+            raise FinetuneError(message=message)
+    return result
+
+
+def _finetune_impl(operation: str, session: Optional[snowpark.Session], function_args: List[Any]) -> str:
+    return call_sql_function_literals(_CORTEX_FINETUNE_SYSTEM_FUNCTION_NAME, session, operation, *function_args)