data-designer 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

data_designer/cli/forms/model_builder.py

@@ -6,8 +6,13 @@ from typing import Any
 from data_designer.cli.forms.builder import FormBuilder
 from data_designer.cli.forms.field import NumericField, SelectField, TextField
 from data_designer.cli.forms.form import Form
-from data_designer.config.models import ModelConfig
-from data_designer.config.utils.constants import MAX_TEMPERATURE, MAX_TOP_P, MIN_TEMPERATURE, MIN_TOP_P
+from data_designer.cli.ui import confirm_action, print_error, print_text
+from data_designer.config.models import (
+    ChatCompletionInferenceParams,
+    EmbeddingInferenceParams,
+    GenerationType,
+    ModelConfig,
+)
 
 
 class ModelFormBuilder(FormBuilder[ModelConfig]):
@@ -19,7 +24,7 @@ class ModelFormBuilder(FormBuilder[ModelConfig]):
         self.available_providers = available_providers or []
 
     def create_form(self, initial_data: dict[str, Any] | None = None) -> Form:
-        """Create the model configuration form."""
+        """Create the model configuration form with basic fields."""
         fields = []
 
         # Model alias
@@ -29,7 +34,7 @@ class ModelFormBuilder(FormBuilder[ModelConfig]):
                 "Model alias (used in your configs)",
                 default=initial_data.get("alias") if initial_data else None,
                 required=True,
-                validator=self._validate_alias,
+                validator=self.validate_alias,
             )
         )
 
@@ -37,10 +42,10 @@ class ModelFormBuilder(FormBuilder[ModelConfig]):
         fields.append(
             TextField(
                 "model",
-                "Model ID",
+                "Model",
                 default=initial_data.get("model") if initial_data else None,
                 required=True,
-                validator=lambda x: (False, "Model ID is required") if not x else (True, None),
+                validator=lambda x: (False, "Model is required") if not x else (True, None),
             )
         )
 
@@ -61,46 +66,222 @@ class ModelFormBuilder(FormBuilder[ModelConfig]):
             # Single provider - will be set automatically
             pass
 
-        # Inference parameters
-        fields.extend(
-            [
+        # Generation type
+        # Extract from inference_parameters if present (for existing models)
+        default_gen_type = GenerationType.CHAT_COMPLETION
+        if initial_data:
+            inference_params = initial_data.get("inference_parameters", {})
+            default_gen_type = inference_params.get("generation_type", default_gen_type)
+
+        fields.append(
+            SelectField(
+                "generation_type",
+                "Generation type",
+                options={
+                    GenerationType.CHAT_COMPLETION: "Chat completion",
+                    GenerationType.EMBEDDING: "Embedding",
+                },
+                default=default_gen_type,
+            )
+        )
+
+        return Form(self.title, fields)
+
+    def create_inference_params_form(
+        self, generation_type: GenerationType, initial_params: dict[str, Any] | None = None
+    ) -> Form:
+        """Create generation-type-specific inference parameters form."""
+        initial_params = initial_params or {}
+        fields = []
+
+        if generation_type == GenerationType.CHAT_COMPLETION:
+            # Temperature
+            fields.append(
                 NumericField(
                     "temperature",
-                    f"Temperature ({MIN_TEMPERATURE}-{MAX_TEMPERATURE})",
-                    default=initial_data.get("inference_parameters", {}).get("temperature", 0.7)
-                    if initial_data
-                    else 0.7,
-                    min_value=MIN_TEMPERATURE,
-                    max_value=MAX_TEMPERATURE,
-                ),
+                    "Temperature <dim>(0.0-2.0)</dim>",
+                    default=initial_params.get("temperature"),
+                    min_value=0.0,
+                    max_value=2.0,
+                    required=False,
+                    help_text="Higher values make output more random, lower values more deterministic",
+                )
+            )
+
+            # Top P
+            fields.append(
                 NumericField(
                     "top_p",
-                    f"Top P ({MIN_TOP_P}-{MAX_TOP_P})",
-                    default=initial_data.get("inference_parameters", {}).get("top_p", 0.9) if initial_data else 0.9,
-                    min_value=MIN_TOP_P,
-                    max_value=MAX_TOP_P,
-                ),
+                    "Top P <dim>(0.0-1.0)</dim>",
+                    default=initial_params.get("top_p"),
+                    min_value=0.0,
+                    max_value=1.0,
+                    required=False,
+                    help_text="Controls diversity via nucleus sampling",
+                )
+            )
+
+            # Max tokens
+            fields.append(
                 NumericField(
                     "max_tokens",
-                    "Max tokens",
-                    default=initial_data.get("inference_parameters", {}).get("max_tokens", 2048)
-                    if initial_data
-                    else 2048,
-                    min_value=1,
-                    max_value=100000,
-                ),
-            ]
-        )
+                    "Max tokens <dim>(maximum total tokens including input and output)</dim>",
+                    default=initial_params.get("max_tokens"),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Maximum number of tokens including both input prompt and generated response",
+                )
+            )
 
-        return Form(self.title, fields)
+            # Max parallel requests
+            fields.append(
+                NumericField(
+                    "max_parallel_requests",
+                    "Max parallel requests <dim>(default: 4)</dim>",
+                    default=initial_params.get("max_parallel_requests", 4),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Maximum number of parallel API requests",
+                )
+            )
 
-    def _validate_alias(self, alias: str) -> tuple[bool, str | None]:
-        """Validate model alias."""
-        if not alias:
-            return False, "Model alias is required"
-        if alias in self.existing_aliases:
-            return False, f"Model alias '{alias}' already exists"
-        return True, None
+            # Timeout
+            fields.append(
+                NumericField(
+                    "timeout",
+                    "Timeout in seconds <dim>(optional)</dim>",
+                    default=initial_params.get("timeout"),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Timeout for each API request in seconds",
+                )
+            )
+
+        else:  # EMBEDDING
+            # Encoding format
+            fields.append(
+                TextField(
+                    "encoding_format",
+                    "Encoding format <dim>(float or base64)</dim>",
+                    default=initial_params.get("encoding_format"),
+                    required=False,
+                    validator=self.validate_encoding_format,
+                )
+            )
+
+            # Dimensions
+            fields.append(
+                NumericField(
+                    "dimensions",
+                    "Dimensions <dim>(number of dimensions for embeddings)</dim>",
+                    default=initial_params.get("dimensions"),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Model-specific dimension size (e.g., 1024, 768)",
+                )
+            )
+
+            # Max parallel requests (common field)
+            fields.append(
+                NumericField(
+                    "max_parallel_requests",
+                    "Max parallel requests <dim>(default: 4)</dim>",
+                    default=initial_params.get("max_parallel_requests", 4),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Maximum number of parallel API requests",
+                )
+            )
+
+            # Timeout (common field)
+            fields.append(
+                NumericField(
+                    "timeout",
+                    "Timeout in seconds <dim>(optional)</dim>",
+                    default=initial_params.get("timeout"),
+                    min_value=1.0,
+                    required=False,
+                    help_text="Timeout for each API request in seconds",
+                )
+            )
+
+        return Form(f"{self.title} - Inference Parameters", fields)
+
+    def build_inference_params(self, generation_type: GenerationType, params_data: dict[str, Any]) -> dict[str, Any]:
+        """Build inference parameters dictionary from form data with proper type conversions."""
+        inference_params = {}
+
+        if generation_type == GenerationType.CHAT_COMPLETION:
+            if params_data.get("temperature") is not None:
+                inference_params["temperature"] = params_data["temperature"]
+            if params_data.get("top_p") is not None:
+                inference_params["top_p"] = params_data["top_p"]
+            if params_data.get("max_tokens") is not None:
+                inference_params["max_tokens"] = int(params_data["max_tokens"])
+
+        else:  # EMBEDDING
+            # Only include fields with actual values; Pydantic will use defaults for missing fields
+            if params_data.get("encoding_format"):
+                inference_params["encoding_format"] = params_data["encoding_format"]
+            if params_data.get("dimensions"):
+                inference_params["dimensions"] = int(params_data["dimensions"])
+
+        # Common fields for both generation types
+        if params_data.get("max_parallel_requests") is not None:
+            inference_params["max_parallel_requests"] = int(params_data["max_parallel_requests"])
+        if params_data.get("timeout") is not None:
+            inference_params["timeout"] = int(params_data["timeout"])
+
+        return inference_params
+
+    def run(self, initial_data: dict[str, Any] | None = None) -> ModelConfig | None:
+        """Run the interactive form with two-step process for generation-type-specific parameters."""
+        # Step 1: Collect basic model configuration
+        basic_form = self.create_form(initial_data)
+
+        if initial_data:
+            basic_form.set_values(initial_data)
+
+        while True:
+            basic_result = basic_form.prompt_all(allow_back=True)
+
+            if basic_result is None:
+                if confirm_action("Cancel configuration?", default=False):
+                    return None
+                continue
+
+            # Step 2: Collect generation-type-specific inference parameters
+            generation_type = basic_result.get("generation_type", GenerationType.CHAT_COMPLETION)
+            initial_params = initial_data.get("inference_parameters") if initial_data else None
+
+            # Print message to indicate we're now configuring inference parameters
+            gen_type_name = "chat completion" if generation_type == GenerationType.CHAT_COMPLETION else "embedding"
+            print_text(
+                f"⚙️  Configuring {gen_type_name} inference parameters [dim](Press Enter to keep current value or skip)[/dim]\n"
+            )
+
+            params_form = self.create_inference_params_form(generation_type, initial_params)
+
+            params_result = params_form.prompt_all(allow_back=True)
+
+            if params_result is None:
+                if confirm_action("Cancel configuration?", default=False):
+                    return None
+                continue
+
+            # Build inference_parameters dict from individual fields
+            inference_params = self.build_inference_params(generation_type, params_result)
+
+            # Merge results
+            full_data = {**basic_result, "inference_parameters": inference_params}
+
+            try:
+                config = self.build_config(full_data)
+                return config
+            except Exception as e:
+                print_error(f"Configuration error: {e}")
+                if not confirm_action("Try again?", default=True):
+                    return None
 
     def build_config(self, form_data: dict[str, Any]) -> ModelConfig:
         """Build ModelConfig from form data."""
@@ -112,14 +293,40 @@ class ModelFormBuilder(FormBuilder[ModelConfig]):
         else:
             provider = None
 
+        # Get generation type (from form data, used to determine which inference params to create)
+        generation_type = form_data.get("generation_type", GenerationType.CHAT_COMPLETION)
+
+        # Get inference parameters dict
+        inference_params_dict = form_data.get("inference_parameters", {})
+
+        # Create the appropriate inference parameters type based on generation_type
+        # The generation_type will be set automatically by the inference params class
+        if generation_type == GenerationType.EMBEDDING:
+            inference_params = EmbeddingInferenceParams(**inference_params_dict)
+        else:
+            inference_params = ChatCompletionInferenceParams(**inference_params_dict)
+
         return ModelConfig(
             alias=form_data["alias"],
             model=form_data["model"],
             provider=provider,
-            inference_parameters={
-                "temperature": form_data["temperature"],
-                "top_p": form_data["top_p"],
-                "max_tokens": int(form_data["max_tokens"]),
-                "max_parallel_requests": 4,
-            },
+            inference_parameters=inference_params,
         )
+
+    def validate_alias(self, alias: str) -> tuple[bool, str | None]:
+        """Validate model alias."""
+        if not alias:
+            return False, "Model alias is required"
+        if alias in self.existing_aliases:
+            return False, f"Model alias '{alias}' already exists"
+        return True, None
+
+    def validate_encoding_format(self, value: str) -> tuple[bool, str | None]:
+        """Validate encoding format for embedding models."""
+        if not value:
+            return True, None  # Optional field
+        if value.lower() in ("clear", "none", "default"):
+            return True, None  # Allow clearing keywords
+        if value not in ("float", "base64"):
+            return False, "Must be either 'float' or 'base64'"
+        return True, None

data_designer/cli/main.py

@@ -3,8 +3,8 @@
 
 import typer
 
+from data_designer.cli.commands import download, models, providers, reset
 from data_designer.cli.commands import list as list_cmd
-from data_designer.cli.commands import models, providers, reset
 from data_designer.config.default_model_settings import resolve_seed_default_model_settings
 from data_designer.config.utils.misc import can_run_data_designer_locally
 
@@ -32,7 +32,17 @@ config_app.command(name="models", help="Configure models interactively")(models.
 config_app.command(name="list", help="List current configurations")(list_cmd.list_command)
 config_app.command(name="reset", help="Reset configuration files")(reset.reset_command)
 
+# Create download command group
+download_app = typer.Typer(
+    name="download",
+    help="Download assets for Data Designer",
+    no_args_is_help=True,
+)
+download_app.command(name="personas", help="Download Nemotron-Persona datasets")(download.personas_command)
+
+# Add command groups to main app
 app.add_typer(config_app, name="config")
+app.add_typer(download_app, name="download")
 
 
 def main() -> None:

data_designer/cli/repositories/persona_repository.py

@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import BaseModel
+
+from data_designer.config.utils.constants import (
+    NEMOTRON_PERSONAS_DATASET_PREFIX,
+    NEMOTRON_PERSONAS_DATASET_SIZES,
+)
+
+
+class PersonaLocale(BaseModel):
+    """Metadata for a single persona locale."""
+
+    code: str
+    size: str
+    dataset_name: str
+
+
+class PersonaLocaleRegistry(BaseModel):
+    """Registry for available persona locales."""
+
+    locales: list[PersonaLocale]
+    dataset_prefix: str = NEMOTRON_PERSONAS_DATASET_PREFIX
+
+
+class PersonaRepository:
+    """Repository for persona locale metadata.
+
+    This repository provides access to built-in persona locale metadata.
+    Unlike ConfigRepository subclasses, this is read-only reference data
+    about what's available in NGC, not user configuration.
+    """
+
+    def __init__(self) -> None:
+        """Initialize repository with built-in locale metadata."""
+        self._registry = self._initialize_registry()
+
+    def _initialize_registry(self) -> PersonaLocaleRegistry:
+        """Initialize registry from constants."""
+        locales = [
+            PersonaLocale(
+                code=code,
+                size=size,
+                dataset_name=f"{NEMOTRON_PERSONAS_DATASET_PREFIX}{code.lower()}",
+            )
+            for code, size in NEMOTRON_PERSONAS_DATASET_SIZES.items()
+        ]
+        return PersonaLocaleRegistry(locales=locales)
+
+    def list_all(self) -> list[PersonaLocale]:
+        """Get all available persona locales.
+
+        Returns:
+            List of all available persona locales
+        """
+        return list(self._registry.locales)
+
+    def get_by_code(self, code: str) -> PersonaLocale | None:
+        """Get a specific locale by code.
+
+        Args:
+            code: Locale code (e.g., 'en_US', 'ja_JP')
+
+        Returns:
+            PersonaLocale if found, None otherwise
+        """
+        return next((locale for locale in self._registry.locales if locale.code == code), None)
+
+    def get_dataset_name(self, code: str) -> str | None:
+        """Get the NGC dataset name for a locale.
+
+        Args:
+            code: Locale code (e.g., 'en_US', 'ja_JP')
+
+        Returns:
+            Dataset name if locale exists, None otherwise
+        """
+        locale = self.get_by_code(code)
+        return locale.dataset_name if locale else None
+
+    def get_dataset_prefix(self) -> str:
+        """Get the dataset prefix for all persona datasets.
+
+        Returns:
+            Dataset prefix string
+        """
+        return self._registry.dataset_prefix

data_designer/cli/services/__init__.py

@@ -1,7 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from data_designer.cli.services.download_service import DownloadService
 from data_designer.cli.services.model_service import ModelService
 from data_designer.cli.services.provider_service import ProviderService
 
-__all__ = ["ModelService", "ProviderService"]
+__all__ = ["DownloadService", "ModelService", "ProviderService"]

data_designer/cli/services/download_service.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import glob
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+from data_designer.cli.repositories.persona_repository import PersonaRepository
+
+
+class DownloadService:
+    """Business logic for downloading assets via NGC CLI."""
+
+    def __init__(self, config_dir: Path, persona_repository: PersonaRepository):
+        self.config_dir = config_dir
+        self.managed_assets_dir = config_dir / "managed-assets" / "datasets"
+        self.persona_repository = persona_repository
+
+    def get_available_locales(self) -> dict[str, str]:
+        """Get dictionary of available persona locales (locale code -> locale code)."""
+        locales = self.persona_repository.list_all()
+        return {locale.code: locale.code for locale in locales}
+
+    def download_persona_dataset(self, locale: str) -> Path:
+        """Download persona dataset for a specific locale using NGC CLI and move to managed assets.
+
+        Args:
+            locale: Locale code (e.g., 'en_US', 'ja_JP')
+
+        Returns:
+            Path to the managed assets datasets directory
+
+        Raises:
+            ValueError: If locale is invalid
+            subprocess.CalledProcessError: If NGC CLI command fails
+        """
+        locale_obj = self.persona_repository.get_by_code(locale)
+        if not locale_obj:
+            raise ValueError(f"Invalid locale: {locale}")
+
+        self.managed_assets_dir.mkdir(parents=True, exist_ok=True)
+
+        # Use temporary directory for download
+        with tempfile.TemporaryDirectory() as temp_dir:
+            # Run NGC CLI download command (without version to get latest)
+            cmd = [
+                "ngc",
+                "registry",
+                "resource",
+                "download-version",
+                f"nvidia/nemotron-personas/{locale_obj.dataset_name}",
+                "--dest",
+                temp_dir,
+            ]
+
+            subprocess.run(cmd, check=True)
+
+            download_pattern = f"{temp_dir}/{locale_obj.dataset_name}*/*.parquet"
+            parquet_files = glob.glob(download_pattern)
+
+            if not parquet_files:
+                raise FileNotFoundError(f"No parquet files found matching pattern: {download_pattern}")
+
+            # Move each parquet file to managed assets
+            for parquet_file in parquet_files:
+                source = Path(parquet_file)
+                dest = self.managed_assets_dir / source.name
+                shutil.move(str(source), str(dest))
+
+        return self.managed_assets_dir
+
+    def get_managed_assets_directory(self) -> Path:
+        """Get the directory where managed datasets are stored."""
+        return self.managed_assets_dir
+
+    def is_locale_downloaded(self, locale: str) -> bool:
+        """Check if a locale has already been downloaded to managed assets.
+
+        Args:
+            locale: Locale code to check
+
+        Returns:
+            True if the locale dataset exists in managed assets
+        """
+        locale_obj = self.persona_repository.get_by_code(locale)
+        if not locale_obj:
+            return False
+
+        if not self.managed_assets_dir.exists():
+            return False
+
+        # Look for any parquet files that start with the dataset pattern
+        parquet_files = glob.glob(str(self.managed_assets_dir / f"{locale}.parquet"))
+
+        return len(parquet_files) > 0