From 232fec547e92d7ae0dc12b06c78dad63ec9dc5b4 Mon Sep 17 00:00:00 2001 From: Vladimir Lipkin Date: Tue, 12 Nov 2024 16:44:22 +0300 Subject: [PATCH] Add example docs --- src/yandex_cloud_ml_sdk/_models/__init__.py | 16 ++++++-- .../_models/image_generation/function.py | 24 ++++++++++- .../_models/image_generation/model.py | 41 +++++++++++++++++++ .../_models/image_generation/result.py | 3 ++ src/yandex_cloud_ml_sdk/_sdk.py | 11 +++-- 5 files changed, 87 insertions(+), 8 deletions(-) diff --git a/src/yandex_cloud_ml_sdk/_models/__init__.py b/src/yandex_cloud_ml_sdk/_models/__init__.py index 3276b0c..a6f96d8 100644 --- a/src/yandex_cloud_ml_sdk/_models/__init__.py +++ b/src/yandex_cloud_ml_sdk/_models/__init__.py @@ -8,16 +8,22 @@ from yandex_cloud_ml_sdk._types.function import BaseFunction from .completions.function import AsyncCompletions, BaseCompletions, Completions -from .image_generation.function import AsyncImageGeneration, ImageGeneration -from .text_classifiers.function import AsyncTextClassifiers, TextClassifiers -from .text_embeddings.function import AsyncTextEmbeddings, TextEmbeddings +from .image_generation.function import AsyncImageGeneration, BaseImageGeneration, ImageGeneration +from .text_classifiers.function import AsyncTextClassifiers, BaseTextClassifiers, TextClassifiers +from .text_embeddings.function import AsyncTextEmbeddings, BaseTextEmbeddings, TextEmbeddings if TYPE_CHECKING: from yandex_cloud_ml_sdk._sdk import BaseSDK class BaseModels(BaseDomain): + """@private""" + completions: BaseCompletions + text_embeddings: BaseTextEmbeddings + text_classifiers: BaseTextClassifiers + image_generation: BaseImageGeneration + """API for image generation models""" def __init__(self, name: str, sdk: BaseSDK): super().__init__(name=name, sdk=sdk) @@ -32,6 +38,8 @@ def _init_functions(self) -> None: class AsyncModels(BaseModels): + """API for working with models at asynchronous code""" + completions: AsyncCompletions text_embeddings: AsyncTextEmbeddings text_classifiers: AsyncTextClassifiers @@ -39,6 +47,8 @@ class AsyncModels(BaseModels): class Models(BaseModels): + """API for working with models at synchronous code""" + completions: Completions text_embeddings: TextEmbeddings text_classifiers: TextClassifiers diff --git a/src/yandex_cloud_ml_sdk/_models/image_generation/function.py b/src/yandex_cloud_ml_sdk/_models/image_generation/function.py index 1c183e7..0f62aa8 100644 --- a/src/yandex_cloud_ml_sdk/_models/image_generation/function.py +++ b/src/yandex_cloud_ml_sdk/_models/image_generation/function.py @@ -8,13 +8,15 @@ class BaseImageGeneration(BaseModelFunction[ModelTypeT]): + """@private""" + @override def __call__( self, model_name: str, *, model_version: str = 'latest', - ): + ) -> ModelTypeT: if '://' in model_name: uri = model_name else: @@ -28,8 +30,28 @@ def __call__( class ImageGeneration(BaseImageGeneration): + """ + Synchronous API which returns + `yandex_cloud_ml_sdk._models.image_generation.model.ImageGenerationModel` + + >>> import yandex_cloud_ml_sdk + >>> sdk = yandex_cloud_ml_sdk.YCloudML(folder_id="...") + >>> model: ImageGenerationModel = sdk.models.image_generation('yandex-art') + >>> rc_model: ImageGenerationModel = sdk.models.image_generation('yandex-art', model_version="rc") + """ + _model_type = ImageGenerationModel class AsyncImageGeneration(BaseImageGeneration): + """ + Asynchronous API which returns + `yandex_cloud_ml_sdk._models.image_generation.model.AsyncImageGenerationModel` + + >>> import yandex_cloud_ml_sdk + >>> sdk = yandex_cloud_ml_sdk.AsyncYCloudML(folder_id="...") + >>> model: AsyncImageGenerationModel = sdk.models.image_generation('yandex-art') + >>> rc_model: AsyncImageGenerationModel = sdk.models.image_generation('yandex-art', model_version="rc") + """ + _model_type = AsyncImageGenerationModel diff --git a/src/yandex_cloud_ml_sdk/_models/image_generation/model.py b/src/yandex_cloud_ml_sdk/_models/image_generation/model.py index b742bb4..4895ae6 100644 --- a/src/yandex_cloud_ml_sdk/_models/image_generation/model.py +++ b/src/yandex_cloud_ml_sdk/_models/image_generation/model.py @@ -24,6 +24,10 @@ class BaseImageGenerationModel( ModelAsyncMixin[ImageGenerationModelConfig, ImageGenerationModelResult, OperationTypeT], ): + """ + Base class for sync and async ImageGeneration models. + """ # NB: we can't @private this because we need to pdoc configure method + _config_type = ImageGenerationModelConfig _result_type = ImageGenerationModelResult _operation_type: type[OperationTypeT] @@ -37,6 +41,16 @@ def configure( # type: ignore[override] height_ratio: UndefinedOr[int] = UNDEFINED, mime_type: UndefinedOr[str] = UNDEFINED, ) -> Self: + """ + Returns a copy of the model with config fields changed. + + :param seed: Seed of a run + :param width_ratio: Width ratio + :param height_ratio: Height_ratio + :param mime_type: Mime type + :type mime_type: str # NB: pydoc ignores this + """ + return super().configure( seed=seed, width_ratio=width_ratio, @@ -79,6 +93,9 @@ async def _run_deferred( class AsyncImageGenerationModel(BaseImageGenerationModel[AsyncOperation[ImageGenerationModelResult]]): + """ + Async class for working with a image generation model + """ _operation_type = AsyncOperation[ImageGenerationModelResult] async def run_deferred( @@ -87,6 +104,16 @@ async def run_deferred( *, timeout: float = 60, ) -> AsyncOperation[ImageGenerationModelResult]: + """ + Run a model and return an `Operation`. + + >>> import yandex_cloud_ml_sdk + >>> sdk = yandex_cloud_ml_sdk.AsyncYCloudML(folder_id='...') + >>> model = sdk.models.image_generation('yandex-art') + >>> operation = await model.run_deferred('kitten') + >>> resullt: ImageGenerationModelResult = await operation + """ + return await self._run_deferred( messages=messages, timeout=timeout @@ -94,6 +121,10 @@ async def run_deferred( class ImageGenerationModel(BaseImageGenerationModel[Operation[ImageGenerationModelResult]]): + """ + Sync class for working with a image generation model + """ + _operation_type = Operation[ImageGenerationModelResult] __run_deferred = run_sync(BaseImageGenerationModel[Operation[ImageGenerationModelResult]]._run_deferred) @@ -103,6 +134,16 @@ def run_deferred( *, timeout: float = 60, ) -> Operation[ImageGenerationModelResult]: + """ + Run a model and return an `Operation`. + + >>> import yandex_cloud_ml_sdk + >>> sdk = yandex_cloud_ml_sdk.YCloudML(folder_id='...') + >>> model = sdk.models.image_generation('yandex-art') + >>> operation = model.run_deferred('kitten') + >>> resullt: ImageGenerationModelResult = operation.wait() + """ + # Mypy thinks that self.__run_deferred returns OperationTypeT instead of Operation return self.__run_deferred( # type: ignore[return-value] messages=messages, diff --git a/src/yandex_cloud_ml_sdk/_models/image_generation/result.py b/src/yandex_cloud_ml_sdk/_models/image_generation/result.py index 28091e7..1584685 100644 --- a/src/yandex_cloud_ml_sdk/_models/image_generation/result.py +++ b/src/yandex_cloud_ml_sdk/_models/image_generation/result.py @@ -15,8 +15,11 @@ @dataclass(frozen=True, repr=False) class ImageGenerationModelResult(BaseResult[ImageGenerationResponse]): + """Result of a image generation run""" _proto_result_type = ImageGenerationResponse image_bytes: bytes + """Bytes which contains a resulting image""" + model_version: str @classmethod diff --git a/src/yandex_cloud_ml_sdk/_sdk.py b/src/yandex_cloud_ml_sdk/_sdk.py index 590371c..03b6a81 100644 --- a/src/yandex_cloud_ml_sdk/_sdk.py +++ b/src/yandex_cloud_ml_sdk/_sdk.py @@ -27,6 +27,7 @@ class BaseSDK: tools: Tools models: BaseModels + """The domain contains the Python API for working with models""" threads: BaseThreads files: BaseFiles assistants: BaseAssistants @@ -51,14 +52,12 @@ def __init__( :param folder_id: Yandex Cloud folder identifier which will be billed for models usage. - :type folder_id: str :param endpoint: domain:port pair for Yandex Cloud API or any other grpc compatible target. - :type endpoint: str :param auth: string with API Key, IAM token or one of yandex_cloud_ml_sdk.auth objects; - in case of default Undefined value, there will be a mechanism to get token + in case of default `Undefined` value, there will be a mechanism to get token from environment - :type api_key | BaseAuth: str + :type api_key: BaseAuth | str :param service_map: a way to redefine endpoints for one or more cloud subservices with a format of dict {service_name: service_address}. :type service_map: Dict[str, str] @@ -143,6 +142,8 @@ def _get_event_loop() -> asyncio.AbstractEventLoop: class AsyncYCloudML(BaseSDK): + """SDK client intended to be used with asyncio""" + tools: Tools models: AsyncModels files: AsyncFiles @@ -154,6 +155,8 @@ class AsyncYCloudML(BaseSDK): class YCloudML(BaseSDK): + """SDK client intended to be used with an usual synchronous code""" + tools: Tools models: Models files: Files