Support for list_upload_schemas & batch task types (#86)

Author: vhaldemar

examples/async/datasets/upload_schemas.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import asyncio
+import pprint
+
+from yandex_cloud_ml_sdk import AsyncYCloudML
+
+
+async def main() -> None:
+    sdk = AsyncYCloudML(folder_id='b1ghsjum2v37c2un8h64')
+    sdk.setup_default_logging()
+
+    for task_type in (
+        'TextToTextGeneration',
+        'TextToTextGenerationRequest',
+        'ImageTextToTextGenerationRequest',
+        'TextEmbeddingsPair',
+        'TextEmbeddingsTriplet',
+        'TextClassificationMultilabel',
+        'TextClassificationMulticlass',
+    ):
+        schemas = await sdk.datasets.list_upload_schemas(task_type)
+        print(f'Schemas for {task_type=}:')
+        pprint.pprint([schema.json for schema in schemas])
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

examples/sync/datasets/upload_schemas.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import pprint
+
+from yandex_cloud_ml_sdk import YCloudML
+
+
+def main() -> None:
+    sdk = YCloudML(folder_id='b1ghsjum2v37c2un8h64')
+    sdk.setup_default_logging()
+
+    for task_type in (
+        'TextToTextGeneration',
+        'TextToTextGenerationRequest',
+        'ImageTextToTextGenerationRequest',
+        'TextEmbeddingsPair',
+        'TextEmbeddingsTriplet',
+        'TextClassificationMultilabel',
+        'TextClassificationMulticlass',
+    ):
+        schemas = sdk.datasets.list_upload_schemas(task_type)
+        print(f'Schemas for {task_type=}:')
+        pprint.pprint([schema.json for schema in schemas])
+
+
+if __name__ == '__main__':
+    main()

src/yandex_cloud_ml_sdk/_datasets/domain.py

@@ -1,12 +1,14 @@
 # pylint: disable=protected-access,no-name-in-module
 from __future__ import annotations
 
+import warnings
 from typing import AsyncIterator, Generic, Iterable, Iterator, Union
 
 from typing_extensions import TypeAlias
 from yandex.cloud.ai.dataset.v1.dataset_service_pb2 import (
     CreateDatasetRequest, CreateDatasetResponse, DescribeDatasetRequest, DescribeDatasetResponse, ListDatasetsRequest,
-    ListDatasetsResponse, ListUploadFormatsRequest, ListUploadFormatsResponse
+    ListDatasetsResponse, ListUploadFormatsRequest, ListUploadFormatsResponse, ListUploadSchemasRequest,
+    ListUploadSchemasResponse
 )
 from yandex.cloud.ai.dataset.v1.dataset_service_pb2_grpc import DatasetServiceStub
 
@@ -17,6 +19,7 @@
 
 from .dataset import AsyncDataset, Dataset, DatasetTypeT
 from .draft import AsyncDatasetDraft, DatasetDraft, DatasetDraftT
+from .schema import DatasetUploadSchema
 from .status import DatasetStatus
 from .task_types import KnownTaskType, TaskTypeProxy
 
@@ -172,6 +175,8 @@ async def _list_upload_formats(
         *,
         timeout: float = 60,
     ) -> tuple[str, ...]:
+        warnings.warn("dataset.list_upload_formats is deprecated", category=DeprecationWarning)
+
         logger.debug('Fetching available dataset upload formats for task_type=%s', task_type)
         request = ListUploadFormatsRequest(
             task_type=task_type
@@ -191,6 +196,35 @@ async def _list_upload_formats(
         )
         return tuple(response.formats)
 
+    async def _list_upload_schemas(
+        self,
+        task_type: str,
+        *,
+        timeout: float = 60,
+    ) -> tuple[DatasetUploadSchema, ...]:
+        logger.debug('Fetching available dataset upload schemas for task_type=%s', task_type)
+        request = ListUploadSchemasRequest(
+            task_type=task_type,
+            folder_id=self._folder_id,
+        )
+
+        async with self._client.get_service_stub(DatasetServiceStub, timeout=timeout) as stub:
+            response = await self._client.call_service(
+                stub.ListUploadSchemas,
+                request,
+                timeout=timeout,
+                expected_type=ListUploadSchemasResponse,
+            )
+
+        logger.info(
+            '%d dataset upload schemas successfully fetched for a task_type=%s',
+            len(response.schemas), task_type,
+        )
+        return tuple(
+            DatasetUploadSchema._from_proto(proto=schema, sdk=self._sdk)
+            for schema in response.schemas
+        )
+
 
 class AsyncDatasets(BaseDatasets[AsyncDataset, AsyncDatasetDraft]):
     _dataset_impl = AsyncDataset
@@ -231,6 +265,14 @@ async def list_upload_formats(
     ) -> tuple[str, ...]:
         return await self._list_upload_formats(task_type=task_type, timeout=timeout)
 
+    async def list_upload_schemas(
+        self,
+        task_type: str,
+        *,
+        timeout: float = 60,
+    ) -> tuple[DatasetUploadSchema, ...]:
+        return await self._list_upload_schemas(task_type=task_type, timeout=timeout)
+
 
 class Datasets(BaseDatasets[Dataset, DatasetDraft]):
     _dataset_impl = Dataset
@@ -239,6 +281,7 @@ class Datasets(BaseDatasets[Dataset, DatasetDraft]):
     __get = run_sync(BaseDatasets._get)
     __list = run_sync_generator(BaseDatasets._list)
     __list_upload_formats = run_sync(BaseDatasets._list_upload_formats)
+    __list_upload_schemas = run_sync(BaseDatasets._list_upload_schemas)
 
     def get(
         self,
@@ -273,3 +316,11 @@ def list_upload_formats(
         timeout: float = 60,
     ) -> tuple[str, ...]:
         return self.__list_upload_formats(task_type=task_type, timeout=timeout)
+
+    def list_upload_schemas(
+        self,
+        task_type: str,
+        *,
+        timeout: float = 60,
+    ) -> tuple[DatasetUploadSchema, ...]:
+        return self.__list_upload_schemas(task_type=task_type, timeout=timeout)

src/yandex_cloud_ml_sdk/_datasets/schema.py

@@ -0,0 +1,30 @@
+# pylint: disable=no-name-in-module
+from __future__ import annotations
+
+import ast
+from dataclasses import dataclass
+
+from typing_extensions import Self
+from yandex.cloud.ai.dataset.v1.dataset_pb2 import DatasetUploadSchema as ProtoDatasetUploadSchema
+
+from yandex_cloud_ml_sdk._types.proto import ProtoBased, SDKType
+from yandex_cloud_ml_sdk._types.schemas import JsonSchemaType
+
+
+@dataclass(frozen=True)
+class DatasetUploadSchema(ProtoBased[ProtoDatasetUploadSchema]):
+    task_type: str
+    upload_format: str
+    raw_schema: str
+
+    @classmethod
+    def _from_proto(cls, *, proto: ProtoDatasetUploadSchema, sdk: SDKType) -> Self:
+        return cls(
+            task_type=proto.task_type,
+            upload_format=proto.upload_format,
+            raw_schema=proto.schema
+        )
+
+    @property
+    def json(self) -> JsonSchemaType:
+        return ast.literal_eval(self.raw_schema)

src/yandex_cloud_ml_sdk/_datasets/task_types.py

@@ -19,6 +19,9 @@ class KnownTaskType(str, Enum):
     TextEmbeddingsPair = 'TextEmbeddingsPair'
     TextEmbeddingsTriplet = 'TextEmbeddingsTriplet'
 
+    TextToTextGenerationRequest = 'TextToTextGenerationRequest'
+    ImageTextToTextGenerationRequest = 'ImageTextToTextGenerationRequest'
+
 
 class TaskTypeProxy:
     def __init__(self, task_type: KnownTaskType):
@@ -66,6 +69,10 @@ def draft_from_path(self):
     def list_upload_formats(self):
         return partial(self._domain.list_upload_formats, task_type=self._task_type)
 
+    @property
+    def list_upload_schemas(self):
+        return partial(self._domain.list_upload_schemas, task_type=self._task_type)
+
     @property
     def list(self):
         return partial(self._domain.list, task_type=self._task_type)