clarifai 11.8.1__py3-none-any.whl → 11.8.3__py3-none-any.whl

clarifai/__init__.py

@@ -1 +1 @@
-__version__ = "11.8.1"
+__version__ = "11.8.3"

clarifai/cli/model.py

@@ -7,13 +7,18 @@ import click
 
 from clarifai.cli.base import cli, pat_display
 from clarifai.utils.cli import (
+    check_lmstudio_installed,
     check_ollama_installed,
     check_requirements_installed,
+    customize_huggingface_model,
+    customize_lmstudio_model,
     customize_ollama_model,
     parse_requirements,
     validate_context,
 )
 from clarifai.utils.constants import (
+    DEFAULT_HF_MODEL_REPO_BRANCH,
+    DEFAULT_LMSTUDIO_MODEL_REPO_BRANCH,
     DEFAULT_LOCAL_RUNNER_APP_ID,
     DEFAULT_LOCAL_RUNNER_COMPUTE_CLUSTER_CONFIG,
     DEFAULT_LOCAL_RUNNER_COMPUTE_CLUSTER_ID,
@@ -22,8 +27,9 @@ from clarifai.utils.constants import (
     DEFAULT_LOCAL_RUNNER_MODEL_TYPE,
     DEFAULT_LOCAL_RUNNER_NODEPOOL_CONFIG,
     DEFAULT_LOCAL_RUNNER_NODEPOOL_ID,
-    DEFAULT_OLLAMA_MODEL_REPO,
     DEFAULT_OLLAMA_MODEL_REPO_BRANCH,
+    DEFAULT_TOOLKIT_MODEL_REPO,
+    DEFAULT_VLLM_MODEL_REPO_BRANCH,
 )
 from clarifai.utils.logging import logger
 from clarifai.utils.misc import (
@@ -68,14 +74,14 @@ def model():
 )
 @click.option(
     '--toolkit',
-    type=click.Choice(['ollama'], case_sensitive=False),
+    type=click.Choice(['ollama', 'huggingface', 'lmstudio', 'vllm'], case_sensitive=False),
     required=False,
-    help='Toolkit to use for model initialization. Currently supports "ollama".',
+    help='Toolkit to use for model initialization. Currently supports "ollama", "huggingface", "lmstudio" and "vllm".',
 )
 @click.option(
     '--model-name',
     required=False,
-    help='Model name to configure when using --toolkit. For ollama toolkit, this sets the Ollama model to use (e.g., "llama3.1", "mistral", etc.).',
+    help='Model name to configure when using --toolkit. For ollama toolkit, this sets the Ollama model to use (e.g., "llama3.1", "mistral", etc.). For vllm & huggingface toolkit, this sets the Hugging Face model repo_id (e.g., "unsloth/Llama-3.2-1B-Instruct").\n For lmstudio toolkit, this sets the LM Studio model name (e.g., "qwen/qwen3-4b-thinking-2507").\n',
 )
 @click.option(
     '--port',
@@ -112,14 +118,16 @@ def init(
     when cloning private repositories. The --branch option can be used to specify a specific
     branch to clone from.
 
-    MODEL_PATH: Path where to create the model directory structure. If not specified, the current directory is used by default.
-    MODEL_TYPE_ID: Type of model to create. If not specified, defaults to "text-to-text" for text models.
-    GITHUB_PAT: GitHub Personal Access Token for authentication when cloning private repositories.
-    GITHUB_URL: GitHub repository URL or "repo" format to clone a repository from. If provided, the entire repository contents will be copied to the target directory instead of using default templates.
-    TOOLKIT: Toolkit to use for model initialization. Currently supports "ollama".
-    MODEL_NAME: Model name to configure when using --toolkit. For ollama toolkit, this sets the Ollama model to use (e.g., "llama3.1", "mistral", etc.).
-    PORT: Port to run the Ollama server on. Defaults to 23333.
-    CONTEXT_LENGTH: Context length for the Ollama model. Defaults to 8192.
+    MODEL_PATH: Path where to create the model directory structure. If not specified, the current directory is used by default.\n
+
+    OPTIONS:\n
+    MODEL_TYPE_ID: Type of model to create. If not specified, defaults to "text-to-text" for text models.\n
+    GITHUB_PAT: GitHub Personal Access Token for authentication when cloning private repositories.\n
+    GITHUB_URL: GitHub repository URL or "repo" format to clone a repository from. If provided, the entire repository contents will be copied to the target directory instead of using default templates.\n
+    TOOLKIT: Toolkit to use for model initialization. Currently supports "ollama", "huggingface", "lmstudio" and "vllm".\n
+    MODEL_NAME: Model name to configure when using --toolkit. For ollama toolkit, this sets the Ollama model to use (e.g., "llama3.1", "mistral", etc.). For vllm & huggingface toolkit, this sets the Hugging Face model repo_id (e.g., "Qwen/Qwen3-4B-Instruct-2507"). For lmstudio toolkit, this sets the LM Studio model name (e.g., "qwen/qwen3-4b-thinking-2507").\n
+    PORT: Port to run the (Ollama/lmstudio) server on. Defaults to 23333.\n
+    CONTEXT_LENGTH: Context length for the (Ollama/lmstudio) model. Defaults to 8192.\n
     """
     # Resolve the absolute path
     model_path = os.path.abspath(model_path)
@@ -152,8 +160,22 @@ def init(
                 "Ollama is not installed. Please install it from `https://ollama.com/` to use the Ollama toolkit."
             )
             raise click.Abort()
-        github_url = DEFAULT_OLLAMA_MODEL_REPO
+        github_url = DEFAULT_TOOLKIT_MODEL_REPO
         branch = DEFAULT_OLLAMA_MODEL_REPO_BRANCH
+    elif toolkit == 'huggingface':
+        github_url = DEFAULT_TOOLKIT_MODEL_REPO
+        branch = DEFAULT_HF_MODEL_REPO_BRANCH
+    elif toolkit == 'lmstudio':
+        if not check_lmstudio_installed():
+            logger.error(
+                "LM Studio is not installed. Please install it from `https://lmstudio.com/` to use the LM Studio toolkit."
+            )
+            raise click.Abort()
+        github_url = DEFAULT_TOOLKIT_MODEL_REPO
+        branch = DEFAULT_LMSTUDIO_MODEL_REPO_BRANCH
+    elif toolkit == 'vllm':
+        github_url = DEFAULT_TOOLKIT_MODEL_REPO
+        branch = DEFAULT_VLLM_MODEL_REPO_BRANCH
 
     if github_url:
         downloader = GitHubDownloader(
@@ -209,6 +231,44 @@ def init(
                     repo_url = format_github_repo_url(github_url)
                     repo_url = f"https://github.com/{owner}/{repo}"
 
+                try:
+                    # Create a temporary directory for cloning
+                    with tempfile.TemporaryDirectory(prefix="clarifai_model_") as clone_dir:
+                        # Clone the repository with explicit branch parameter
+                        if not clone_github_repo(repo_url, clone_dir, github_pat, branch):
+                            logger.error(f"Failed to clone repository from {repo_url}")
+                            github_url = None  # Fall back to template mode
+
+                        else:
+                            # Copy the entire repository content to target directory (excluding .git)
+                            for item in os.listdir(clone_dir):
+                                if item == '.git':
+                                    continue
+
+                                source_path = os.path.join(clone_dir, item)
+                                target_path = os.path.join(model_path, item)
+
+                                if os.path.isdir(source_path):
+                                    shutil.copytree(source_path, target_path, dirs_exist_ok=True)
+                                else:
+                                    shutil.copy2(source_path, target_path)
+
+                            logger.info(f"Successfully cloned repository to {model_path}")
+                            logger.info(
+                                "Model initialization complete with GitHub repository clone"
+                            )
+                            logger.info("Next steps:")
+                            logger.info("1. Review the model configuration")
+                            logger.info("2. Install any required dependencies manually")
+                            logger.info(
+                                "3. Test the model locally using 'clarifai model local-test'"
+                            )
+                            return
+
+                except Exception as e:
+                    logger.error(f"Failed to clone GitHub repository: {e}")
+                    github_url = None  # Fall back to template mode
+
     if toolkit:
         logger.info(f"Initializing model from GitHub repository: {github_url}")
 
@@ -218,35 +278,42 @@ def init(
         else:
             repo_url = format_github_repo_url(github_url)
 
-    try:
-        # Create a temporary directory for cloning
-        with tempfile.TemporaryDirectory(prefix="clarifai_model_") as clone_dir:
-            # Clone the repository with explicit branch parameter
-            if not clone_github_repo(repo_url, clone_dir, github_pat, branch):
-                logger.error(f"Failed to clone repository from {repo_url}")
-                github_url = None  # Fall back to template mode
-
-            else:
-                # Copy the entire repository content to target directory (excluding .git)
-                for item in os.listdir(clone_dir):
-                    if item == '.git':
-                        continue
-
-                    source_path = os.path.join(clone_dir, item)
-                    target_path = os.path.join(model_path, item)
-
-                    if os.path.isdir(source_path):
-                        shutil.copytree(source_path, target_path, dirs_exist_ok=True)
-                    else:
-                        shutil.copy2(source_path, target_path)
+        try:
+            # Create a temporary directory for cloning
+            with tempfile.TemporaryDirectory(prefix="clarifai_model_") as clone_dir:
+                # Clone the repository with explicit branch parameter
+                if not clone_github_repo(repo_url, clone_dir, github_pat, branch):
+                    logger.error(f"Failed to clone repository from {repo_url}")
+                    github_url = None  # Fall back to template mode
 
-    except Exception as e:
-        logger.error(f"Failed to clone GitHub repository: {e}")
-        github_url = None
+                else:
+                    # Copy the entire repository content to target directory (excluding .git)
+                    for item in os.listdir(clone_dir):
+                        if item == '.git':
+                            continue
+
+                        source_path = os.path.join(clone_dir, item)
+                        target_path = os.path.join(model_path, item)
+
+                        if os.path.isdir(source_path):
+                            shutil.copytree(source_path, target_path, dirs_exist_ok=True)
+                        else:
+                            shutil.copy2(source_path, target_path)
+
+        except Exception as e:
+            logger.error(f"Failed to clone GitHub repository: {e}")
+            github_url = None
 
     if (model_name or port or context_length) and (toolkit == 'ollama'):
         customize_ollama_model(model_path, model_name, port, context_length)
 
+    if (model_name or port or context_length) and (toolkit == 'lmstudio'):
+        customize_lmstudio_model(model_path, model_name, port, context_length)
+
+    if model_name and (toolkit == 'huggingface' or toolkit == 'vllm'):
+        # Update the config.yaml file with the provided model name
+        customize_huggingface_model(model_path, model_name)
+
     if github_url:
         logger.info("Model initialization complete with GitHub repository")
         logger.info("Next steps:")
@@ -294,7 +361,7 @@ def init(
         if os.path.exists(config_path):
             logger.warning(f"File {config_path} already exists, skipping...")
         else:
-            config_model_type_id = "text-to-text"  # default
+            config_model_type_id = DEFAULT_LOCAL_RUNNER_MODEL_TYPE  # default
 
             config_template = get_config_template(config_model_type_id)
             with open(config_path, 'w') as f:

clarifai/cli/pipeline.py

@@ -26,14 +26,19 @@ def pipeline():
 
 @pipeline.command()
 @click.argument("path", type=click.Path(exists=True), required=False, default=".")
-def upload(path):
+@click.option(
+    '--no-lockfile',
+    is_flag=True,
+    help='Skip creating config-lock.yaml file.',
+)
+def upload(path, no_lockfile):
     """Upload a pipeline with associated pipeline steps to Clarifai.
 
     PATH: Path to the pipeline configuration file or directory containing config.yaml. If not specified, the current directory is used by default.
     """
     from clarifai.runners.pipelines.pipeline_builder import upload_pipeline
 
-    upload_pipeline(path)
+    upload_pipeline(path, no_lockfile=no_lockfile)
 
 
 @pipeline.command()
@@ -106,15 +111,32 @@ def run(
 
     validate_context(ctx)
 
+    # Try to load from config-lock.yaml first if no config is specified
+    lockfile_path = os.path.join(os.getcwd(), "config-lock.yaml")
+    if not config and os.path.exists(lockfile_path):
+        logger.info("Found config-lock.yaml, using it as default config source")
+        config = lockfile_path
+
     if config:
         config_data = from_yaml(config)
-        pipeline_id = config_data.get('pipeline_id', pipeline_id)
-        pipeline_version_id = config_data.get('pipeline_version_id', pipeline_version_id)
+
+        # Handle both regular config format and lockfile format
+        if 'pipeline' in config_data and isinstance(config_data['pipeline'], dict):
+            pipeline_config = config_data['pipeline']
+            pipeline_id = pipeline_config.get('id', pipeline_id)
+            pipeline_version_id = pipeline_config.get('version_id', pipeline_version_id)
+            user_id = pipeline_config.get('user_id', user_id)
+            app_id = pipeline_config.get('app_id', app_id)
+        else:
+            # Fallback to flat config structure
+            pipeline_id = config_data.get('pipeline_id', pipeline_id)
+            pipeline_version_id = config_data.get('pipeline_version_id', pipeline_version_id)
+            user_id = config_data.get('user_id', user_id)
+            app_id = config_data.get('app_id', app_id)
+
         pipeline_version_run_id = config_data.get(
             'pipeline_version_run_id', pipeline_version_run_id
         )
-        user_id = config_data.get('user_id', user_id)
-        app_id = config_data.get('app_id', app_id)
         nodepool_id = config_data.get('nodepool_id', nodepool_id)
         compute_cluster_id = config_data.get('compute_cluster_id', compute_cluster_id)
         pipeline_url = config_data.get('pipeline_url', pipeline_url)
@@ -319,6 +341,62 @@ def init(pipeline_path):
     logger.info("3. Run 'clarifai pipeline upload config.yaml' to upload your pipeline")
 
 
+@pipeline.command()
+@click.argument(
+    "lockfile_path", type=click.Path(exists=True), required=False, default="config-lock.yaml"
+)
+def validate_lock(lockfile_path):
+    """Validate a config-lock.yaml file for schema and reference consistency.
+
+    LOCKFILE_PATH: Path to the config-lock.yaml file. If not specified, looks for config-lock.yaml in current directory.
+    """
+    from clarifai.runners.utils.pipeline_validation import PipelineConfigValidator
+    from clarifai.utils.cli import from_yaml
+
+    try:
+        # Load the lockfile
+        lockfile_data = from_yaml(lockfile_path)
+
+        # Validate required fields
+        if "pipeline" not in lockfile_data:
+            raise ValueError("'pipeline' section not found in lockfile")
+
+        pipeline = lockfile_data["pipeline"]
+        required_fields = ["id", "user_id", "app_id", "version_id"]
+
+        for field in required_fields:
+            if field not in pipeline:
+                raise ValueError(f"Required field '{field}' not found in pipeline section")
+            if not pipeline[field]:
+                raise ValueError(f"Required field '{field}' cannot be empty")
+
+        # Validate orchestration spec if present
+        if "orchestration_spec" in pipeline:
+            # Create a temporary config structure for validation
+            temp_config = {
+                "pipeline": {
+                    "id": pipeline["id"],
+                    "user_id": pipeline["user_id"],
+                    "app_id": pipeline["app_id"],
+                    "orchestration_spec": pipeline["orchestration_spec"],
+                }
+            }
+
+            # Use existing validator to check orchestration spec
+            validator = PipelineConfigValidator()
+            validator._validate_orchestration_spec(temp_config)
+
+        logger.info(f"✅ Lockfile {lockfile_path} is valid")
+        logger.info(f"Pipeline: {pipeline['id']}")
+        logger.info(f"User: {pipeline['user_id']}")
+        logger.info(f"App: {pipeline['app_id']}")
+        logger.info(f"Version: {pipeline['version_id']}")
+
+    except Exception as e:
+        logger.error(f"❌ Lockfile validation failed: {e}")
+        raise click.Abort()
+
+
 @pipeline.command(['ls'])
 @click.option('--page_no', required=False, help='Page number to list.', default=1)
 @click.option('--per_page', required=False, help='Number of items per page.', default=16)

clarifai/cli/templates/model_templates.py

@@ -178,7 +178,7 @@ class MyModel(OpenAIModelClass):
 '''
 
 
-def get_config_template(model_type_id: str = "text-to-text") -> str:
+def get_config_template(model_type_id: str = "any-to-any") -> str:
     """Return the template for config.yaml."""
     return f'''# Configuration file for your Clarifai model
 

clarifai/client/base.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any, Callable
+from typing import Any, Callable, Dict, Optional
 
 from clarifai_grpc.grpc.api import resources_pb2
 from google.protobuf import struct_pb2
@@ -88,20 +88,46 @@ class BaseClient:
         self.root_certificates_path = self.auth_helper._root_certificates_path
 
     @property
-    def async_stub(self):
+    def async_stub(self) -> Any:
         """Returns the asynchronous gRPC stub for the API interaction.
-        Lazy initialization of async stub"""
+
+        Returns:
+            Any: The async gRPC stub object for making asynchronous API calls.
+
+        Note:
+            Uses lazy initialization - stub is created on first access.
+        """
         if self._async_stub is None:
             self._async_stub = create_stub(self.auth_helper, is_async=True)
         return self._async_stub
 
     @classmethod
-    def from_env(cls, validate: bool = False):
+    def from_env(cls, validate: bool = False) -> 'BaseClient':
+        """Creates a BaseClient instance from environment variables.
+
+        Args:
+            validate (bool): Whether to validate the authentication credentials.
+                           Defaults to False.
+
+        Returns:
+            BaseClient: A new BaseClient instance configured from environment variables.
+        """
         auth = ClarifaiAuthHelper.from_env(validate=validate)
         return cls.from_auth_helper(auth)
 
     @classmethod
-    def from_auth_helper(cls, auth: ClarifaiAuthHelper, **kwargs):
+    def from_auth_helper(cls, auth: ClarifaiAuthHelper, **kwargs) -> 'BaseClient':
+        """Creates a BaseClient instance from a ClarifaiAuthHelper.
+
+        Args:
+            auth (ClarifaiAuthHelper): The authentication helper containing credentials.
+            **kwargs: Additional keyword arguments to override auth helper values.
+                     Supported keys: user_id, app_id, pat, token, root_certificates_path,
+                     base, ui, url.
+
+        Returns:
+            BaseClient: A new BaseClient instance configured from the auth helper.
+        """
         default_kwargs = {
             "user_id": kwargs.get("user_id", None) or auth.user_id,
             "app_id": kwargs.get("app_id", None) or auth.app_id,
@@ -132,15 +158,18 @@ class BaseClient:
 
         return cls(**kwargs)
 
-    def _grpc_request(self, method: Callable, argument: Any):
+    def _grpc_request(self, method: Callable[..., Any], argument: Any) -> Any:
         """Makes a gRPC request to the API.
 
         Args:
-            method (Callable): The gRPC method to call.
-            argument (Any): The argument to pass to the gRPC method.
+            method (Callable[..., Any]): The gRPC stub method to call.
+            argument (Any): The protobuf request object to pass to the gRPC method.
 
         Returns:
-            res (Any): The result of the gRPC method call.
+            Any: The protobuf response object from the gRPC method call.
+
+        Raises:
+            Exception: If the API request fails.
         """
 
         try:
@@ -150,14 +179,16 @@ class BaseClient:
         except ApiError:
             raise Exception("ApiError")
 
-    def convert_string_to_timestamp(self, date_str) -> Timestamp:
+    def convert_string_to_timestamp(self, date_str: str) -> Timestamp:
         """Converts a string to a Timestamp object.
 
         Args:
-            date_str (str): The string to convert.
+            date_str (str): The date string to convert. Accepts formats like
+                          '%Y-%m-%dT%H:%M:%S.%fZ' or '%Y-%m-%dT%H:%M:%SZ'.
 
         Returns:
-            Timestamp: The converted Timestamp object.
+            Timestamp: The converted protobuf Timestamp object. Returns empty Timestamp
+                      if the date string format is invalid.
         """
         # Parse the string into a Python datetime object
         try:
@@ -174,15 +205,22 @@ class BaseClient:
 
         return timestamp_obj
 
-    def process_response_keys(self, old_dict, listing_resource=None):
+    def process_response_keys(
+        self, old_dict: Dict[str, Any], listing_resource: Optional[str] = None
+    ) -> Dict[str, Any]:
         """Converts keys in a response dictionary to resource proto format.
 
+        This method processes dictionary keys to match protobuf resource naming conventions,
+        particularly for converting 'id' to '{resource}_id' format.
+
         Args:
-            old_dict (dict): The dictionary to convert.
-            listing_resource (str, optional): The resource type for which the keys are being processed.
+            old_dict (Dict[str, Any]): The dictionary to convert with keys to process.
+            listing_resource (Optional[str]): The resource type for which the keys are being
+                                            processed (e.g., 'model', 'workflow'). If provided,
+                                            renames 'id' to '{listing_resource}_id'.
 
         Returns:
-            new_dict (dict): The dictionary with processed keys.
+            Dict[str, Any]: The dictionary with processed keys and converted values.
         """
         if listing_resource:
             old_dict[f'{listing_resource}_id'] = old_dict['id']

clarifai/client/dataset.py

@@ -4,7 +4,16 @@ import uuid
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
 from multiprocessing import cpu_count
-from typing import Dict, Generator, List, Optional, Tuple, Type, TypeVar, Union
+from typing import (
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+)
 
 import requests
 from clarifai_grpc.grpc.api import resources_pb2, service_pb2
@@ -122,10 +131,13 @@ class Dataset(Lister, BaseClient):
         Args:
             **kwargs: Additional keyword arguments to be passed to Dataset Version.
               - description (str): The description of the dataset version.
-              - metadata (dict): The metadata of the dataset version.
+              - metadata (Dict[str, Any]): The metadata dictionary for the dataset version.
 
         Returns:
-            Dataset: A Dataset object for the specified dataset ID.
+            Dataset: A Dataset object for the newly created dataset version.
+
+        Raises:
+            Exception: If the dataset version creation fails.
 
         Example:
             >>> from clarifai.client.dataset import Dataset
@@ -172,13 +184,13 @@ class Dataset(Lister, BaseClient):
         self.logger.info("\nDataset Version Deleted\n%s", response.status)
 
     def list_versions(
-        self, page_no: int = None, per_page: int = None
+        self, page_no: Optional[int] = None, per_page: Optional[int] = None
     ) -> Generator['Dataset', None, None]:
         """Lists all the versions for the dataset.
 
         Args:
-            page_no (int): The page number to list.
-            per_page (int): The number of items per page.
+            page_no (Optional[int]): The page number to list. If None, lists all pages.
+            per_page (Optional[int]): The number of items per page. If None, uses default.
 
         Yields:
             Dataset: Dataset objects for the versions of the dataset.

clarifai/client/model.py

@@ -2,7 +2,7 @@ import itertools
 import json
 import os
 import time
-from typing import Any, Dict, Generator, Iterable, Iterator, List, Tuple, Union
+from typing import Any, Dict, Generator, Iterable, Iterator, List, Optional, Tuple, Union
 
 import numpy as np
 import requests
@@ -178,20 +178,26 @@ class Model(Lister, BaseClient):
 
         return templates
 
-    def get_params(self, template: str = None, save_to: str = 'params.yaml') -> Dict[str, Any]:
-        """Returns the model params for the model type and yaml file.
+    def get_params(
+        self, template: Optional[str] = None, save_to: str = 'params.yaml'
+    ) -> Dict[str, Any]:
+        """Returns the model params for the model type and saves them to a yaml file.
 
         Args:
-            template (str): The template to use for the model type.
-            yaml_file (str): The yaml file to save the model params.
+            template (Optional[str]): The template to use for the model type. Required for most
+                                    model types except 'clusterer' and 'embedding-classifier'.
+            save_to (str): The yaml file path to save the model params. Defaults to 'params.yaml'.
 
         Returns:
-            params (Dict): Dictionary of model params for the model type.
+            Dict[str, Any]: Dictionary of model params for the model type.
+
+        Raises:
+            UserError: If the model type is not trainable, or if template is required but not provided.
 
         Example:
             >>> from clarifai.client.model import Model
             >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
-            >>> model_params = model.get_params(template='template', yaml_file='model_params.yaml')
+            >>> model_params = model.get_params(template='template', save_to='model_params.yaml')
         """
         if not self.model_info.model_type_id:
             self.load_info()
@@ -260,19 +266,22 @@ class Model(Lister, BaseClient):
             find_and_replace_key(self.training_params, key, value)
 
     def get_param_info(self, param: str) -> Dict[str, Any]:
-        """Returns the param info for the param.
+        """Returns the parameter info for the specified parameter.
 
         Args:
-            param (str): The param to get the info for.
+            param (str): The parameter name to get information for.
 
         Returns:
-            param_info (Dict): Dictionary of model param info for the param.
+            Dict[str, Any]: Dictionary containing model parameter info for the specified param.
+
+        Raises:
+            UserError: If the model type is not trainable or if training params are not loaded.
 
         Example:
             >>> from clarifai.client.model import Model
             >>> model = Model(model_id='model_id', user_id='user_id', app_id='app_id')
-            >>> model_params = model.get_params(template='template', yaml_file='model_params.yaml')
-            >>> model.get_param_info('param')
+            >>> model_params = model.get_params(template='template', save_to='model_params.yaml')
+            >>> param_info = model.get_param_info('learning_rate')
         """
         if self.model_info.model_type_id not in TRAINABLE_MODEL_TYPES:
             raise UserError(f"Model type {self.model_info.model_type_id} is not trainable")
@@ -513,8 +522,9 @@ class Model(Lister, BaseClient):
                 model=self.model_info,
                 runner_selector=self._runner_selector,
             )
+            # Pass in None for async stub will create it.
             self._client = ModelClient(
-                stub=self.STUB, async_stub=self.async_stub, request_template=request_template
+                stub=self.STUB, async_stub=None, request_template=request_template
             )
         return self._client
 

clarifai/client/model_client.py

@@ -314,6 +314,7 @@ class ModelClient:
         self,
         base_url: str = None,
         use_ctx: bool = False,
+        colorize: bool = False,
     ) -> str:
         """Generate a client script for this model.
 
@@ -335,6 +336,7 @@ class ModelClient:
             compute_cluster_id=self.request_template.runner_selector.nodepool.compute_cluster.id,
             nodepool_id=self.request_template.runner_selector.nodepool.id,
             use_ctx=use_ctx,
+            colorize=colorize,
         )
 
     def _define_compatability_functions(self):