seekrai 0.0.1__py3-none-any.whl

seekrai/resources/finetune.py

@@ -0,0 +1,425 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from seekrai.abstract import api_requestor
+from seekrai.filemanager import DownloadManager
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    FinetuneDownloadResult,
+    FinetuneList,
+    FinetuneListEvents,
+    FinetuneRequest,
+    FinetuneResponse,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+TrainingConfig,
+InfrastructureConfig
+)
+from seekrai.utils import normalize_key
+
+
+class FineTuning:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def create(
+        self,
+        *,
+        training_config: TrainingConfig,
+        infrastructure_config: InfrastructureConfig
+        # wandb_api_key: str | None = None,
+    ) -> FinetuneResponse:
+        """
+        Method to initiate a fine-tuning job
+
+        Args:
+            training_file (str): File-ID of a file uploaded to the SeekrFlow API
+            model (str): Name of the base model to run fine-tune job on
+            n_epochs (int, optional): Number of epochs for fine-tuning. Defaults to 1.
+            n_checkpoints (int, optional): Number of checkpoints to save during fine-tuning.
+                Defaults to 1.
+            batch_size (int, optional): Batch size for fine-tuning. Defaults to 32.
+            learning_rate (float, optional): Learning rate multiplier to use for training
+                Defaults to 0.00001.
+            suffix (str, optional): Up to 40 character suffix that will be added to your fine-tuned model name.
+                Defaults to None.
+            wandb_api_key (str, optional): API key for Weights & Biases integration.
+                Defaults to None.
+
+        Returns:
+            FinetuneResponse: Object containing information about fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = FinetuneRequest(
+            training_config=training_config,
+            infrastructure_config=infrastructure_config
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="fine-tune",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    def list(self) -> FinetuneList:
+        """
+        Lists fine-tune job history
+
+        Returns:
+            FinetuneList: Object containing a list of fine-tune jobs
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="fine-tunes",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneList(**response.data)
+
+    def retrieve(self, id: str) -> FinetuneResponse:
+        """
+        Retrieves fine-tune job details
+
+        Args:
+            id (str): Fine-tune ID to retrieve. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneResponse: Object containing information about fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"fine-tunes/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    def cancel(self, id: str) -> FinetuneResponse:
+        """
+        Method to cancel a running fine-tuning job
+
+        Args:
+            id (str): Fine-tune ID to cancel. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneResponse: Object containing information about cancelled fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"fine-tunes/{id}/cancel",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    def list_events(self, id: str) -> FinetuneListEvents:
+        """
+        Lists events of a fine-tune job
+
+        Args:
+            id (str): Fine-tune ID to list events for. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneListEvents: Object containing list of fine-tune events
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"fine-tunes/{id}/events",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneListEvents(**response.data)
+
+    def download(
+        self, id: str, *, output: Path | str | None = None, checkpoint_step: int = -1
+    ) -> FinetuneDownloadResult:
+        """
+        Downloads compressed fine-tuned model or checkpoint to local disk.
+
+        Defaults file location to `$PWD/{model_name}.{extension}`
+
+        Args:
+            id (str): Fine-tune ID to download. A string that starts with `ft-`.
+            output (pathlib.Path | str, optional): Specifies output file name for downloaded model.
+                Defaults to None.
+            checkpoint_step (int, optional): Specifies step number for checkpoint to download.
+                Defaults to -1 (download the final model)
+
+        Returns:
+            FinetuneDownloadResult: Object containing downloaded model metadata
+        """
+
+        url = f"finetune/download?ft_id={id}"
+
+        if checkpoint_step > 0:
+            url += f"&checkpoint_step={checkpoint_step}"
+
+        remote_name = self.retrieve(id).output_name
+
+        download_manager = DownloadManager(self._client)
+
+        if isinstance(output, str):
+            output = Path(output)
+
+        downloaded_filename, file_size = download_manager.download(
+            url, output, normalize_key(remote_name or id), fetch_metadata=True
+        )
+
+        return FinetuneDownloadResult(
+            object="local",
+            id=id,
+            checkpoint_step=checkpoint_step,
+            filename=downloaded_filename,
+            size=file_size,
+        )
+
+
+    def promote(self, id: str):
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"fine-tunes/{id}/promote-model",
+                params={"fine_tune_id": id},
+            ),
+            stream=False,
+        )
+
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneListEvents(**response.data)
+
+
+class AsyncFineTuning:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def create(
+        self,
+        *,
+        training_file: str,
+        model: str,
+        n_epochs: int = 1,
+        n_checkpoints: int | None = 1,
+        batch_size: int | None = 32,
+        learning_rate: float = 0.00001,
+        suffix: str | None = None,
+        wandb_api_key: str | None = None,
+    ) -> FinetuneResponse:
+        """
+        Async method to initiate a fine-tuning job
+
+        Args:
+            training_file (str): File-ID of a file uploaded to the SeekrFlow API
+            model (str): Name of the base model to run fine-tune job on
+            n_epochs (int, optional): Number of epochs for fine-tuning. Defaults to 1.
+            n_checkpoints (int, optional): Number of checkpoints to save during fine-tuning.
+                Defaults to 1.
+            batch_size (int, optional): Batch size for fine-tuning. Defaults to 32.
+            learning_rate (float, optional): Learning rate multiplier to use for training
+                Defaults to 0.00001.
+            suffix (str, optional): Up to 40 character suffix that will be added to your fine-tuned model name.
+                Defaults to None.
+            wandb_api_key (str, optional): API key for Weights & Biases integration.
+                Defaults to None.
+
+        Returns:
+            FinetuneResponse: Object containing information about fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = FinetuneRequest(
+            model=model,
+            training_file=training_file,
+            n_epochs=n_epochs,
+            n_checkpoints=n_checkpoints,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            suffix=suffix,
+            wandb_key=wandb_api_key,
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="fine-tunes",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    async def list(self) -> FinetuneList:
+        """
+        Async method to list fine-tune job history
+
+        Returns:
+            FinetuneList: Object containing a list of fine-tune jobs
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="fine-tunes",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneList(**response.data)
+
+    async def retrieve(self, id: str) -> FinetuneResponse:
+        """
+        Async method to retrieve fine-tune job details
+
+        Args:
+            id (str): Fine-tune ID to retrieve. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneResponse: Object containing information about fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"fine-tunes/{id}",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    async def cancel(self, id: str) -> FinetuneResponse:
+        """
+        Async method to cancel a running fine-tuning job
+
+        Args:
+            id (str): Fine-tune ID to cancel. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneResponse: Object containing information about cancelled fine-tuning job.
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url=f"fine-tunes/{id}/cancel",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneResponse(**response.data)
+
+    async def list_events(self, id: str) -> FinetuneListEvents:
+        """
+        Async method to lists events of a fine-tune job
+
+        Args:
+            id (str): Fine-tune ID to list events for. A string that starts with `ft-`.
+
+        Returns:
+            FinetuneListEvents: Object containing list of fine-tune events
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"fine-tunes/{id}/events",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneListEvents(**response.data)
+
+    async def download(
+        self, id: str, *, output: str | None = None, checkpoint_step: int = -1
+    ) -> str:
+        """
+        TODO: Implement async download method
+        """
+
+        raise NotImplementedError(
+            "AsyncFineTuning.download not implemented. "
+            "Please use FineTuning.download function instead."
+        )

seekrai/resources/images.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+from seekrai.abstract import api_requestor
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    ImageRequest,
+    ImageResponse,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+
+
+class Images:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def generate(
+        self,
+        *,
+        prompt: str,
+        model: str,
+        steps: int | None = 20,
+        seed: int | None = None,
+        n: int | None = 1,
+        height: int | None = 1024,
+        width: int | None = 1024,
+        negative_prompt: str | None = None,
+    ) -> ImageResponse:
+        """
+        Method to generate images based on a given prompt using a specified model.
+
+        Args:
+            prompt (str): A description of the desired images. Maximum length varies by model.
+
+            model (str, optional): The model to use for image generation.
+
+            steps (int, optional): Number of generation steps. Defaults to 20
+
+            seed (int, optional): Seed used for generation. Can be used to reproduce image generations.
+                Defaults to None.
+
+            n (int, optional): Number of image results to generate. Defaults to 1.
+
+            height (int, optional): Height of the image to generate in number of pixels. Defaults to 1024
+
+            width (int, optional): Width of the image to generate in number of pixels. Defaults to 1024
+
+            negative_prompt (str, optional): The prompt or prompts not to guide the image generation.
+                Defaults to None
+
+            image_base64: (str, optional): Reference image used for generation. Defaults to None.
+
+        Returns:
+            ImageResponse: Object containing image data
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = ImageRequest(
+            prompt=prompt,
+            model=model,
+            steps=steps,
+            seed=seed,
+            n=n,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+        ).model_dump()
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="images/generations",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return ImageResponse(**response.data)
+
+
+class AsyncImages:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def generate(
+        self,
+        *,
+        prompt: str,
+        model: str,
+        steps: int | None = 20,
+        seed: int | None = None,
+        n: int | None = 1,
+        height: int | None = 1024,
+        width: int | None = 1024,
+        negative_prompt: str | None = None,
+    ) -> ImageResponse:
+        """
+        Async method to generate images based on a given prompt using a specified model.
+
+        Args:
+            prompt (str): A description of the desired images. Maximum length varies by model.
+
+            model (str, optional): The model to use for image generation.
+
+            steps (int, optional): Number of generation steps. Defaults to 20
+
+            seed (int, optional): Seed used for generation. Can be used to reproduce image generations.
+                Defaults to None.
+
+            n (int, optional): Number of image results to generate. Defaults to 1.
+
+            height (int, optional): Height of the image to generate in number of pixels. Defaults to 1024
+
+            width (int, optional): Width of the image to generate in number of pixels. Defaults to 1024
+
+            negative_prompt (str, optional): The prompt or prompts not to guide the image generation.
+                Defaults to None
+
+            image_base64: (str, optional): Reference image used for generation. Defaults to None.
+
+        Returns:
+            ImageResponse: Object containing image data
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        parameter_payload = ImageRequest(
+            prompt=prompt,
+            model=model,
+            steps=steps,
+            seed=seed,
+            n=n,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+        ).model_dump()
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="POST",
+                url="images/generations",
+                params=parameter_payload,
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return ImageResponse(**response.data)

seekrai/resources/models.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import List
+
+from seekrai.abstract import api_requestor
+from seekrai.seekrflow_response import SeekrFlowResponse
+from seekrai.types import (
+    ModelObject,
+    SeekrFlowClient,
+    SeekrFlowRequest,
+)
+
+
+class Models:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    def list(
+        self,
+    ) -> List[ModelObject]:
+        """
+        Method to return list of models on the API
+
+        Returns:
+            List[ModelObject]: List of model objects
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="models",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        assert isinstance(response.data, list)
+
+        return [ModelObject(**model) for model in response.data]
+
+
+class AsyncModels:
+    def __init__(self, client: SeekrFlowClient) -> None:
+        self._client = client
+
+    async def list(
+        self,
+    ) -> List[ModelObject]:
+        """
+        Async method to return list of models on API
+
+        Returns:
+            List[ModelObject]: List of model objects
+        """
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = await requestor.arequest(
+            options=SeekrFlowRequest(
+                method="GET",
+                url="models",
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+        assert isinstance(response.data, list)
+
+        return [ModelObject(**model) for model in response.data]

seekrai/seekrflow_response.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+
+class SeekrFlowResponse:
+    """
+    API Response class. Stores headers and response data.
+    """
+
+    def __init__(self, data: Dict[str, Any], headers: Dict[str, Any]):
+        self._headers = headers
+        self.data = data
+
+    @property
+    def request_id(self) -> str | None:
+        """
+        Fetches request id from headers
+        """
+        if "cf-ray" in self._headers:
+            return str(self._headers["cf-ray"])
+        return None
+
+    @property
+    def requests_remaining(self) -> int | None:
+        """
+        Number of requests remaining at current rate limit
+        """
+        if "x-ratelimit-remaining" in self._headers:
+            return int(self._headers["x-ratelimit-remaining"])
+        return None
+
+    @property
+    def processed_by(self) -> str | None:
+        """
+        Processing host server name
+        """
+        if "x-hostname" in self._headers:
+            return str(self._headers["x-hostname"])
+        return None
+
+    @property
+    def response_ms(self) -> int | None:
+        """
+        Server request completion time
+        """
+        if "x-total-time" in self._headers:
+            h = self._headers["x-total-time"]
+            return None if h is None else round(float(h))
+        return None