seekrai 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

seekrai/resources/finetune.py

@@ -3,7 +3,6 @@ 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,
@@ -11,12 +10,11 @@ from seekrai.types import (
     FinetuneListEvents,
     FinetuneRequest,
     FinetuneResponse,
+    InfrastructureConfig,
     SeekrFlowClient,
     SeekrFlowRequest,
-TrainingConfig,
-InfrastructureConfig
+    TrainingConfig,
 )
-from seekrai.utils import normalize_key
 
 
 class FineTuning:
@@ -27,25 +25,13 @@ class FineTuning:
         self,
         *,
         training_config: TrainingConfig,
-        infrastructure_config: InfrastructureConfig
+        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.
@@ -56,14 +42,13 @@ class FineTuning:
         )
 
         parameter_payload = FinetuneRequest(
-            training_config=training_config,
-            infrastructure_config=infrastructure_config
+            training_config=training_config, infrastructure_config=infrastructure_config
         ).model_dump()
 
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url="fine-tune",
+                url="flow/fine-tune",
                 params=parameter_payload,
             ),
             stream=False,
@@ -88,7 +73,7 @@ class FineTuning:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url="fine-tunes",
+                url="flow/fine-tunes",
             ),
             stream=False,
         )
@@ -115,7 +100,7 @@ class FineTuning:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"fine-tunes/{id}",
+                url=f"flow/fine-tunes/{id}",
             ),
             stream=False,
         )
@@ -142,7 +127,7 @@ class FineTuning:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url=f"fine-tunes/{id}/cancel",
+                url=f"flow/fine-tunes/{id}/cancel",
             ),
             stream=False,
         )
@@ -169,7 +154,7 @@ class FineTuning:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"fine-tunes/{id}/events",
+                url=f"flow/fine-tunes/{id}/events",
             ),
             stream=False,
         )
@@ -196,33 +181,50 @@ class FineTuning:
         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
+        raise NotImplementedError("Function not yet implemented")
+        # 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) -> FinetuneListEvents:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
         )
 
-        return FinetuneDownloadResult(
-            object="local",
-            id=id,
-            checkpoint_step=checkpoint_step,
-            filename=downloaded_filename,
-            size=file_size,
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/fine-tunes/{id}/promote-model",
+                params={"fine_tune_id": id},
+            ),
+            stream=False,
         )
 
+        assert isinstance(response, SeekrFlowResponse)
+
+        return FinetuneListEvents(**response.data)
 
-    def promote(self, id: str):
+    def demote(self, id: str) -> FinetuneListEvents:
         requestor = api_requestor.APIRequestor(
             client=self._client,
         )
@@ -230,13 +232,12 @@ class FineTuning:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"fine-tunes/{id}/promote-model",
+                url=f"flow/fine-tunes/{id}/demote-model",
                 params={"fine_tune_id": id},
             ),
             stream=False,
         )
 
-
         assert isinstance(response, SeekrFlowResponse)
 
         return FinetuneListEvents(**response.data)
@@ -249,32 +250,13 @@ class AsyncFineTuning:
     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,
+        training_config: TrainingConfig,
+        infrastructure_config: InfrastructureConfig,
     ) -> 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.
         """
@@ -284,20 +266,13 @@ class AsyncFineTuning:
         )
 
         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,
+            training_config=training_config, infrastructure_config=infrastructure_config
         ).model_dump()
 
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url="fine-tunes",
+                url="flow/fine-tunes",
                 params=parameter_payload,
             ),
             stream=False,
@@ -322,7 +297,7 @@ class AsyncFineTuning:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url="fine-tunes",
+                url="flow/fine-tunes",
             ),
             stream=False,
         )
@@ -349,7 +324,7 @@ class AsyncFineTuning:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"fine-tunes/{id}",
+                url=f"flow/fine-tunes/{id}",
             ),
             stream=False,
         )
@@ -376,7 +351,7 @@ class AsyncFineTuning:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url=f"fine-tunes/{id}/cancel",
+                url=f"flow/fine-tunes/{id}/cancel",
             ),
             stream=False,
         )
@@ -403,7 +378,7 @@ class AsyncFineTuning:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url=f"fine-tunes/{id}/events",
+                url=f"flow/fine-tunes/{id}/events",
             ),
             stream=False,
         )

seekrai/resources/images.py

@@ -72,7 +72,7 @@ class Images:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="POST",
-                url="images/generations",
+                url="inference/images/generations",
                 params=parameter_payload,
             ),
             stream=False,
@@ -145,7 +145,7 @@ class AsyncImages:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="POST",
-                url="images/generations",
+                url="inference/images/generations",
                 params=parameter_payload,
             ),
             stream=False,

seekrai/resources/models.py

@@ -1,28 +1,62 @@
 from __future__ import annotations
 
-from typing import List
+import os
+from pathlib import Path
+from typing import Any, List
+
+from tqdm import tqdm
 
 from seekrai.abstract import api_requestor
+from seekrai.constants import DISABLE_TQDM
 from seekrai.seekrflow_response import SeekrFlowResponse
-from seekrai.types import (
-    ModelObject,
-    SeekrFlowClient,
-    SeekrFlowRequest,
-)
+from seekrai.types import ModelList, ModelResponse, SeekrFlowClient, SeekrFlowRequest
+from seekrai.types.models import ModelType
 
 
 class Models:
     def __init__(self, client: SeekrFlowClient) -> None:
         self._client = client
 
+    def upload(
+        self,
+        file: Path | str,
+        *,
+        model_type: ModelType | str = ModelType.OBJECT_DETECTION,
+    ) -> ModelResponse:
+        if isinstance(file, str):
+            file = Path(file)
+
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+        file_size = os.stat(file.as_posix()).st_size
+
+        with tqdm(
+            total=file_size,
+            unit="B",
+            unit_scale=True,
+            desc=f"Uploading file {file.name}",
+            disable=bool(DISABLE_TQDM),
+        ):
+            with file.open("rb") as f:
+                response, _, _ = requestor.request(
+                    options=SeekrFlowRequest(
+                        method="PUT",
+                        url="flow/pt-models",
+                        files={"files": f, "filename": file.name},
+                        params={"purpose": model_type},
+                    ),
+                )
+        return ModelResponse(**response.data)
+
     def list(
         self,
-    ) -> List[ModelObject]:
+    ) -> ModelList:
         """
         Method to return list of models on the API
 
         Returns:
-            List[ModelObject]: List of model objects
+            List[ModelResponse]: List of model objects
         """
 
         requestor = api_requestor.APIRequestor(
@@ -32,15 +66,81 @@ class Models:
         response, _, _ = requestor.request(
             options=SeekrFlowRequest(
                 method="GET",
-                url="models",
+                url="flow/pt-models",
             ),
             stream=False,
         )
 
         assert isinstance(response, SeekrFlowResponse)
-        assert isinstance(response.data, list)
+        return ModelList(**response.data)
+
+    def promote(self, id: str) -> ModelResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/pt-models/{id}/promote-model",
+                params={"model_id": id},
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return ModelResponse(**response.data)
+
+    def demote(self, id: str) -> ModelResponse:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        response, _, _ = requestor.request(
+            options=SeekrFlowRequest(
+                method="GET",
+                url=f"flow/pt-models/{id}/demote-model",
+                params={"model_id": id},
+            ),
+            stream=False,
+        )
+
+        assert isinstance(response, SeekrFlowResponse)
+
+        return ModelResponse(**response.data)
+
+    def predict(self, id: str, file: Path | str) -> Any:
+        requestor = api_requestor.APIRequestor(
+            client=self._client,
+        )
+
+        if isinstance(file, str):
+            file = Path(file)
+
+        file_size = os.stat(file.as_posix()).st_size
+
+        with tqdm(
+            total=file_size,
+            unit="B",
+            unit_scale=True,
+            desc=f"Uploading file {file.name}",
+            disable=bool(DISABLE_TQDM),
+        ):
+            with file.open("rb") as f:
+                response, _, _ = requestor.request(
+                    options=SeekrFlowRequest(
+                        method="POST",
+                        url="flow/pt-models/predict",
+                        files={"files": f, "filename": file.name},
+                        params={"model_id": id},
+                    ),
+                    stream=False,
+                )
+
+        assert isinstance(response, SeekrFlowResponse)
 
-        return [ModelObject(**model) for model in response.data]
+        return response.data
 
 
 class AsyncModels:
@@ -49,12 +149,12 @@ class AsyncModels:
 
     async def list(
         self,
-    ) -> List[ModelObject]:
+    ) -> List[ModelResponse]:
         """
         Async method to return list of models on API
 
         Returns:
-            List[ModelObject]: List of model objects
+            List[ModelResponse]: List of model objects
         """
 
         requestor = api_requestor.APIRequestor(
@@ -64,7 +164,7 @@ class AsyncModels:
         response, _, _ = await requestor.arequest(
             options=SeekrFlowRequest(
                 method="GET",
-                url="models",
+                url="flow/models",
             ),
             stream=False,
         )
@@ -72,4 +172,4 @@ class AsyncModels:
         assert isinstance(response, SeekrFlowResponse)
         assert isinstance(response.data, list)
 
-        return [ModelObject(**model) for model in response.data]
+        return [ModelResponse(**model) for model in response.data]

seekrai/types/__init__.py

@@ -26,14 +26,14 @@ from seekrai.types.finetune import (
     FinetuneListEvents,
     FinetuneRequest,
     FinetuneResponse,
-InfrastructureConfig,
-TrainingConfig,
+    InfrastructureConfig,
+    TrainingConfig,
 )
 from seekrai.types.images import (
     ImageRequest,
     ImageResponse,
 )
-from seekrai.types.models import ModelObject
+from seekrai.types.models import ModelList, ModelResponse
 
 
 __all__ = [
@@ -63,5 +63,6 @@ __all__ = [
     "FileType",
     "ImageRequest",
     "ImageResponse",
-    "ModelObject",
+    "ModelResponse",
+    "ModelList",
 ]

seekrai/types/common.py

@@ -4,7 +4,6 @@ from enum import Enum
 from typing import Any, Dict, List
 
 from pydantic import ConfigDict
-from tqdm.utils import CallbackIOWrapper
 
 from seekrai.types.abstract import BaseModel
 
@@ -58,7 +57,7 @@ class SeekrFlowRequest(BaseModel):
     method: str
     url: str
     headers: Dict[str, str] | None = None
-    params: Dict[str, Any] | CallbackIOWrapper | None = None
+    params: Dict[str, Any] | None = None
     files: Dict[str, Any] | None = None
     allow_redirects: bool = True
     override_headers: bool = False

seekrai/types/files.py

@@ -1,23 +1,24 @@
 from __future__ import annotations
 
+from datetime import datetime
 from enum import Enum
 from typing import List, Literal
 
-from pydantic import Field
-
 from seekrai.types.abstract import BaseModel
 from seekrai.types.common import (
     ObjectType,
 )
-from datetime import datetime
+
 
 class FilePurpose(str, Enum):
     FineTune = "fine-tune"
+    PreTrain = "pre-train"
 
 
 class FileType(str, Enum):
     jsonl = "jsonl"
     parquet = "parquet"
+    pytorch = "pt"
 
 
 class FileRequest(BaseModel):
@@ -25,22 +26,25 @@ class FileRequest(BaseModel):
     Files request type
     """
 
-    # training file ID
-    training_file: str
-    # base model string
-    model: str
-    # number of epochs to train for
-    n_epochs: int
-    # training learning rate
-    learning_rate: float
-    # number of checkpoints to save
-    n_checkpoints: int | None = None
-    # training batch size
-    batch_size: int | None = None
-    # up to 40 character suffix for output model name
-    suffix: str | None = None
-    # weights & biases api key
-    wandb_api_key: str | None = None
+    # # training file ID
+    # training_file: str
+    # # base model string
+    # model: str
+    # # number of epochs to train for
+    # n_epochs: int
+    # # training learning rate
+    # learning_rate: float
+    # # number of checkpoints to save
+    # n_checkpoints: int | None = None
+    # # training batch size
+    # batch_size: int | None = None
+    # # up to 40 character suffix for output model name
+    # suffix: str | None = None
+    # # weights & biases api key
+    # wandb_api_key: str | None = None
+    purpose: FilePurpose
+    filetype: FileType
+    filename: str
 
 
 class FileResponse(BaseModel):

seekrai/types/finetune.py

@@ -1,16 +1,13 @@
 from __future__ import annotations
 
+from datetime import datetime
 from enum import Enum
 from typing import List, Literal
 
-from pydantic import Field
-
 from seekrai.types.abstract import BaseModel
 from seekrai.types.common import (
     ObjectType,
 )
-from datetime import datetime
-
 
 
 class FinetuneJobStatus(str, Enum):
@@ -83,22 +80,10 @@ class FinetuneEvent(BaseModel):
     # object type
     object: Literal[ObjectType.FinetuneEvent]
     # created at datetime stamp
-    created_at: str | None = None
-    # event log level
-    level: FinetuneEventLevels | None = None
-    # event message string
-    message: str | None = None
-    # event type
-    type: FinetuneEventType | None = None
-    # optional: model parameter count
-    param_count: int | None = None
-    # optional: dataset token count
-    token_count: int | None = None
-    # optional: weights & biases url
-    wandb_url: str | None = None
-    # event hash
-    hash: str | None = None
-
+    created_at: datetime | None = None
+    # metrics that we expose
+    loss: float | None = None
+    epoch: float | None = None
 
 
 class TrainingConfig(BaseModel):
@@ -118,22 +103,31 @@ class TrainingConfig(BaseModel):
     experiment_name: str | None = None
     # # weights & biases api key
     # wandb_key: str | None = None
+    # IFT by default
+    pre_train: bool = False
+
+
+class AcceleratorType(str, Enum):
+    GAUDI2 = "GAUDI2"
+    GAUDI3 = "GAUDI3"
+    A100 = "A100"
+    H100 = "H100"
+
 
 class InfrastructureConfig(BaseModel):
-    n_cpu: int
-    n_gpu: int
+    accel_type: AcceleratorType
+    n_accel: int
+
 
 class FinetuneRequest(BaseModel):
     """
     Fine-tune request type
     """
+
     training_config: TrainingConfig
     infrastructure_config: InfrastructureConfig
 
 
-
-
-
 class FinetuneResponse(BaseModel):
     """
     Fine-tune API response type
@@ -170,7 +164,7 @@ class FinetuneResponse(BaseModel):
 
     # list of fine-tune events
     events: List[FinetuneEvent] | None = None
-    inference_available: bool  = False
+    inference_available: bool = False
     # dataset token count
     # TODO
     # token_count: int | None = None