nemo-platform-plugin 0.1.0__py3-none-any.whl

nemo_platform/beta/evaluator/values/scores.py

@@ -0,0 +1,300 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Score configuration types for metric definitions."""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from abc import ABC, abstractmethod
+from typing import Annotated, Any, Literal, Self
+
+import jsonschema
+from jsonschema import validators
+from pydantic import BaseModel, ConfigDict, Discriminator, Field, Tag, model_validator
+
+from nemo_platform.beta.evaluator.values.results import MetricScore, RubricScoreStat, ScoreStats
+
+_logger = logging.getLogger(__name__)
+
+
+class JSONScoreParser(BaseModel):
+    """Parse a score from JSON structured content."""
+
+    type: Literal["json"] = "json"
+    json_path: str = Field(
+        description="The JSON path to parse the score from the judge response when using structured output."
+    )
+
+
+class RegexScoreParser(BaseModel):
+    """Parse a score from content in any format using regular expression."""
+
+    type: Literal["regex"] = "regex"
+    pattern: str = Field(description="The regular expression to parse the score from the judge response.")
+    method: Literal["search", "match"] = Field(
+        default="match",
+        description="The regex method to use: 'search' looks anywhere in the string, 'match' starts from the beginning.",
+    )
+
+    @model_validator(mode="after")
+    def valid_regex(self) -> Self:
+        if self.type != "regex":
+            # model_validator will run when resolving parser JSONScoreParser | RegexScoreParser
+            # if it's RegexScoreParser, continue with validation.
+            return self
+
+        try:
+            re.compile(self.pattern)
+        except re.error as e:
+            raise ValueError(f"invalid regex for score parser: {e}")
+        return self
+
+
+class _Score(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    name: str = Field(
+        pattern=r"^[a-z0-9_]+$",
+        description="The name of the score. Only lowercase letters, numbers, and underscores allowed.",
+    )
+    description: str | None = Field(default=None, description="Human-readable description of the score.")
+    parser: JSONScoreParser | RegexScoreParser = Field(
+        default_factory=lambda data: JSONScoreParser(json_path=data["name"]),
+        description="The method to parse the score. When used with llm-judge metric, and no parser is set, JSONScoreParser is the default parser inferred from the score parameters.",
+    )
+
+
+class Rubric(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    label: str = Field(
+        description='The label to use for the level of the rubric grading criteria. (e.g., "helpful", "not_helpful", "positive")'
+    )
+    description: str | None = Field(
+        default=None,
+        description="Describe the semantic meaning of each criteria for the given rubric. If no judge template is set, the input description for labels are included in the generated judge prompt.",
+    )
+    value: float | int = Field(
+        description="The score value to assign for the criteria used for aggregation and ranking."
+    )
+
+
+class RubricScore(_Score):
+    """Score definition for a rubric with optional parser. If no parser is set, JSONScoreParser is the default parser inferred from the score parameters"""
+
+    model_config = ConfigDict(extra="forbid")
+    rubric: list[Rubric] = Field(min_length=2, description="The rubric for the score.")
+
+
+class RangeScore(_Score):
+    """Score definition for a range of values with optional parser. If no parser is set, JSONScoreParser is the default parser inferred from the score parameters"""
+
+    model_config = ConfigDict(extra="forbid")
+    minimum: float | int = Field(description="Minimum value for the score range. Must be less than maximum.")
+    maximum: float | int = Field(description="Maximum value for the score range. Must be greater than minimum.")
+
+    @model_validator(mode="after")
+    def valid_range(self) -> Self:
+        if self.minimum >= self.maximum:
+            raise ValueError(f"minimum must be less than maximum: {self.minimum} < {self.maximum}")
+        return self
+
+
+class RemoteScore(_Score):
+    """Score configuration for remote metrics.
+
+    Unlike RangeScore, minimum and maximum are optional (default to None = no bounds).
+    This avoids JSON serialization issues with infinity values.
+    """
+
+    minimum: float | int | None = Field(
+        default=None,
+        description="Minimum value for the score range. Defaults to None (no lower bound).",
+    )
+    maximum: float | int | None = Field(
+        default=None,
+        description="Maximum value for the score range. Defaults to None (no upper bound).",
+    )
+    parser: JSONScoreParser = Field(
+        default_factory=lambda data: JSONScoreParser(json_path=data["name"]),
+        description="The method to parse the score. Only JSON parsing is supported for remote metrics.",
+    )
+
+    @model_validator(mode="after")
+    def valid_range(self) -> Self:
+        """Validate that minimum < maximum when both are configured."""
+        if self.minimum is not None and self.maximum is not None:
+            if self.minimum >= self.maximum:
+                raise ValueError(f"minimum must be less than maximum: {self.minimum} >= {self.maximum}")
+        return self
+
+
+def score_discriminator(data: dict[str, Any] | RubricScore | RangeScore) -> Literal["rubric", "range"]:
+    if "rubric" in data or isinstance(data, RubricScore):
+        return "rubric"
+    return "range"
+
+
+Score = Annotated[
+    (Annotated[RubricScore, Tag("rubric")] | Annotated[RangeScore, Tag("range")]), Discriminator(score_discriminator)
+]
+
+
+class ScoreParser(ABC):
+    """Parse model output text into a normalized ``MetricScore``."""
+
+    class Params(BaseModel):
+        pass
+
+    rubric_mapping: dict[str, Rubric] | None = None
+
+    def __init__(self, score: Score):
+        self.score = score
+        if isinstance(score, RubricScore):
+            self.rubric_mapping = {rubric.label: rubric for rubric in score.rubric}
+
+    @abstractmethod
+    def parse(self, text: str | None) -> MetricScore:
+        """Parse the provided text and extract a single score."""
+        raise NotImplementedError
+
+    def _get_rubric_score(self, label: str) -> MetricScore:
+        """Map rubric label to score value and build rubric distribution stats."""
+        assert isinstance(self.score, RubricScore)
+        assert self.rubric_mapping is not None
+        rubric_distribution = [
+            RubricScoreStat(
+                label=rubric.label, description=rubric.description, value=rubric.value, count=int(label == rubric.label)
+            )
+            for rubric in self.score.rubric
+        ]
+        rubric = self.rubric_mapping.get(label)
+        score_value = float("nan") if rubric is None else rubric.value
+        return MetricScore(
+            name=self.score.name, value=score_value, stats=ScoreStats(rubric_distribution=rubric_distribution)
+        )
+
+
+class ScoreParserRegex(ScoreParser):
+    """Parse score values from free-form text using a regex capture group."""
+
+    parser_type: Literal["regex"] = "regex"
+    pattern: re.Pattern
+    method: Literal["search", "match"]
+
+    def __init__(self, score: Score):
+        super().__init__(score)
+        if not isinstance(score.parser, RegexScoreParser):
+            raise ValueError(f"incompatible score parser to initialize ScoreParserRegex: {type(score.parser)}")
+        try:
+            self.pattern = re.compile(score.parser.pattern)
+        except re.error as e:
+            raise ValueError(f"invalid regex pattern for score parser with LLM-as-a-Judge: {score.parser.pattern} {e}")
+        self.method = score.parser.method
+
+    def parse(self, text: str | None) -> MetricScore:
+        if not text:
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        match = self.pattern.search(text) if self.method == "search" else self.pattern.match(text)
+        if match is None:
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        groups = match.groups()
+        if len(groups) == 0:
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        try:
+            if self.rubric_mapping:
+                return self._get_rubric_score(groups[0])
+            return MetricScore(name=self.score.name, value=float(groups[0]))
+        except ValueError:
+            # This is expected when models drift from requested output format.
+            _logger.info("Failed to parse score from text: %s.", text)
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+
+class ScoreParserJSON(ScoreParser):
+    """Parse score values from JSON output using ``json_path`` key lookup."""
+
+    parser_type: Literal["json"] = "json"
+    json_path: str
+
+    # Optional when structured output is defined.
+    structured_output: dict | None = None
+    json_schema: dict | None = None
+    json_validator: jsonschema.Validator | None = None
+
+    def __init__(self, score: Score, structured_output: dict | None = None):
+        super().__init__(score)
+        if not isinstance(score.parser, JSONScoreParser):
+            raise ValueError(f"incompatible score parser to initialize ScoreParserJSON: {type(score.parser)}")
+
+        self.json_path = score.parser.json_path
+
+        if structured_output:
+            self._validate_structured_output(score, structured_output)
+            self.structured_output = structured_output
+            self.json_schema = structured_output["schema"]
+
+    def _validate_structured_output(self, score: Score, structured_output: dict) -> None:
+        json_schema = structured_output.get("schema")
+        if not json_schema:
+            raise ValueError("missing schema for structured output")
+
+        # Validate schema itself and verify parser expectations for json_path.
+        validator = validators.validator_for(json_schema)
+        validator.check_schema(json_schema)
+
+        schema_type = json_schema.get("type", "")
+        if schema_type != "object":
+            raise ValueError(f"schema must be type 'object' for JSON score parser: {schema_type}")
+
+        assert score.parser is not None
+        assert isinstance(score.parser, JSONScoreParser)
+        assert score.parser.json_path is not None
+
+        schema_defined_json_path = json_schema.get("properties", {}).get(score.parser.json_path)
+        if not schema_defined_json_path:
+            raise ValueError(
+                f"schema must have {score.parser.json_path} defined as an object property for JSON score parser: {json_schema}"
+            )
+        json_type = schema_defined_json_path.get("type")
+        if isinstance(score, RubricScore):
+            if json_type != "string":
+                raise ValueError(
+                    f"expected string type in schema for property '{score.parser.json_path}' when used with score rubric and JSON score parser: {json_type}"
+                )
+        elif json_type not in ["number", "integer", "boolean"]:
+            raise ValueError(
+                f"schema property {score.parser.json_path} must be type number, integer, or boolean for JSON score parser: {schema_defined_json_path}"
+            )
+
+    def parse(self, text: str | None) -> MetricScore:
+        if not isinstance(text, str):
+            return MetricScore(name=self.score.name, value=float("nan"))
+        try:
+            obj = json.loads(text)
+        except json.JSONDecodeError:
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        if not isinstance(obj, dict):
+            # handles if model returns just a number, for example
+            _logger.warning("Expected JSON object, got %s. Returning NaN score", type(obj))
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        score = obj.get(self.json_path)
+        if score is None:
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        if self.rubric_mapping:
+            return self._get_rubric_score(score)
+
+        if isinstance(score, str):
+            return MetricScore(name=self.score.name, value=float("nan"))
+
+        if isinstance(score, bool):
+            score = 1.0 if score else 0.0
+
+        return MetricScore(name=self.score.name, value=score)

nemo_platform/beta/safe_synthesizer/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+version = "0.0.1"

nemo_platform/beta/safe_synthesizer/config.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+try:
+    from nemo_safe_synthesizer.config import (
+        DataParameters,
+        DifferentialPrivacyHyperparams,
+        EvaluationParameters,
+        GenerateParameters,
+        PiiReplacerConfig,
+        SafeSynthesizerJobConfig,
+        SafeSynthesizerParameters,
+        TimeSeriesParameters,
+        TrainingHyperparams,
+    )
+
+    __all__ = [
+        "DataParameters",
+        "DifferentialPrivacyHyperparams",
+        "EvaluationParameters",
+        "GenerateParameters",
+        "PiiReplacerConfig",
+        "SafeSynthesizerJobConfig",
+        "SafeSynthesizerParameters",
+        "TimeSeriesParameters",
+        "TrainingHyperparams",
+    ]
+except ImportError:
+    __all__ = []

nemo_platform/beta/safe_synthesizer/job.py

@@ -0,0 +1,309 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import time
+from base64 import b64encode
+from io import BytesIO
+from pathlib import Path
+from typing import Iterator
+
+import pandas as pd
+from nemo_platform import NeMoPlatform
+from nemo_platform.types import PlatformJobLog, PlatformJobStatusResponse
+from nemo_platform.types.safe_synthesizer.jobs.safe_synthesizer_summary import (
+    SafeSynthesizerSummary,
+)
+from typing_extensions import Self
+
+logger = logging.getLogger(__name__)
+
+
+class ReportHtml(object):
+    """Class to hold the HTML report string."""
+
+    def __init__(self, html: str):
+        self.raw_html = html
+        self.as_data_uri = f"data:text/html;base64,{b64encode(self.raw_html.encode()).decode()}"
+
+    def save(self, path: str | Path) -> None:
+        """Save the evaluation report to a file.
+
+        Args:
+            path: The path to save the report to.
+        """
+        with open(path, "w") as f:
+            f.write(self.raw_html)
+
+    def display_report_in_notebook(self, width="100%", height=1000) -> None:
+        """Display the evaluation report in a jupyter notebook.
+
+        Requires the IPython library to be installed.
+
+        Args:
+            width: The width of the iframe to display the report in.
+            height: The height of the iframe to display the report in.
+        """
+        try:
+            from IPython.display import IFrame, display
+        except ImportError:
+            logger.warning("IPython is required to display reports in notebooks. Report will not be displayed.")
+            return
+
+        display(IFrame(self.as_data_uri, width=width, height=height))
+
+    @classmethod
+    def read(cls, path: str | Path) -> Self:
+        """Read the evaluation report from a file.
+
+        Args:
+            path: The path to read the report from.
+        """
+        with open(path, "r") as f:
+            raw_html = f.read()
+        return cls(raw_html)
+
+
+class SafeSynthesizerJob:
+    """Interface for convenient interaction with Safe Synthesizer jobs.
+
+    This class provides a wrapper around the Safe Synthesizer job SDK to make common operations easier.
+
+    An instance is returned from the `create_job` method of the `SafeSynthesizerJobBuilder` class.
+    Or create an instance with a job id and a NeMoPlatform client.
+
+    Examples:
+        >>> from nemo_platform import NeMoPlatform
+        >>> from nemo_platform.beta.safe_synthesizer.job import SafeSynthesizerJob
+        >>> client = NeMoPlatform(base_url=..., inference_base_url=...)
+        >>> job = SafeSynthesizerJob(job_id=..., client=client, workspace="default")
+        >>> job.fetch_status()
+        >>> job.wait_for_completion()
+        >>> job.fetch_summary()
+        >>> df = job.fetch_data()
+        >>> job.save_report("./evaluation_report.html")
+
+        And in a jupyter notebook to display the report inline:
+
+        >>> job.display_report_in_notebook()
+    """
+
+    def __init__(self, job_name: str, client: NeMoPlatform, workspace: str = "default"):
+        """Initialize a SafeSynthesizerJob instance.
+
+        Args:
+            job_name: The name of the job to interact with (note: despite the parameter name, this should be the job name, not entity ID).
+            client: The NeMoPlatform client to use to interact with the job.
+            workspace: The workspace to use to interact with the job.
+        """
+        self.job_name = job_name
+        self._client = client
+        self._workspace = workspace
+
+    def fetch_status(self) -> str:
+        """Fetch the status of the job.
+
+        Returns:
+            The status of the job.
+        """
+        return self._client.safe_synthesizer.jobs.get_status(self.job_name, workspace=self._workspace).status
+
+    def fetch_status_info(self) -> PlatformJobStatusResponse:
+        """Fetch the status information of the job.
+
+        Returns:
+            The status information of the job.
+        """
+        return self._client.safe_synthesizer.jobs.get_status(self.job_name, workspace=self._workspace)
+
+    def wait_for_completion(
+        self, poll_interval: int = 10, verbose: bool = True, log_timeout: float | None = None
+    ) -> None:
+        """Block until the job is completed.
+
+        Prints the logs by default. Uses incremental log fetching via page cursor
+        to only fetch new logs on each poll.
+
+        Args:
+            poll_interval: The interval in seconds to poll the job status.
+            verbose: Gets logs and prints them at this interval. Default: True
+            log_timeout: Timeout in seconds for log requests. Defaults to 300s (5 minutes).
+        """
+        # Track cursor to continue from last position, and seen timestamps for dedup
+        last_page_cursor: str | None = None
+        seen_log_keys: set[str] = set()
+        previous_status_info = None
+        current_status_info = self.fetch_status_info()
+        while current_status_info.status not in ["completed", "error", "cancelled"]:
+            if verbose:
+                logging_level = None
+                try:
+                    httpx_logger = logging.getLogger("httpx")
+                    logging_level = httpx_logger.level
+                    httpx_logger.setLevel("ERROR")
+                    # Fetch logs starting from last cursor position
+                    new_logs, last_page_cursor = self._fetch_logs_incremental(
+                        page_cursor=last_page_cursor, timeout=log_timeout
+                    )
+                    for new_log in new_logs:
+                        # Deduplicate using timestamp + message hash (cursor may re-fetch last page)
+                        log_key = f"{new_log.timestamp}:{hash(new_log.message)}"
+                        if log_key not in seen_log_keys:
+                            print(new_log.message.strip())
+                            seen_log_keys.add(log_key)
+                except Exception:
+                    logger.exception("Error fetching logs")
+                finally:
+                    if logging_level:
+                        logging.getLogger("httpx").setLevel(logging_level)
+            current_status_info = self.fetch_status_info()
+            if current_status_info != previous_status_info:
+                print(
+                    f"Job status changed to status: '{current_status_info.status}',",
+                    f"status_details: {current_status_info.status_details},",
+                    f"error_details: {current_status_info.error_details}",
+                )
+                previous_status_info = current_status_info
+            time.sleep(poll_interval)
+        if current_status_info.status in ["error", "cancelled"]:
+            raise RuntimeError(
+                f"Job '{self.job_name}' ended with status '{current_status_info.status}'. "
+                f"Details: {current_status_info.status_details}. "
+                f"Error: {current_status_info.error_details}. "
+                "Check job logs with job.print_logs() for more details."
+            )
+
+    def fetch_summary(self) -> SafeSynthesizerSummary:
+        """Fetch the summary of the job.
+
+        Returns:
+            A summary of machine-readable metrics for a completed job. Raises a 404 error if the job is not finished.
+        """
+        return self._client.safe_synthesizer.jobs.results.download_summary(self.job_name, workspace=self._workspace)
+
+    def fetch_report(self) -> ReportHtml:
+        """Fetch the evaluation report of the job as a string of html.
+
+        Recommended to use save_report or display_report_in_notebook for most use cases.
+
+        Returns:
+            A string containing the html representation of the evaluation report.
+        """
+        response = self._client.safe_synthesizer.jobs.results.download_evaluation_report(
+            self.job_name, workspace=self._workspace
+        )
+        return ReportHtml(html=response.read().decode("utf-8"))
+
+    def display_report_in_notebook(self, width="100%", height=1000) -> None:
+        """Display the evaluation report in a jupyter notebook.
+
+        Requires the IPython library to be installed.
+
+        Args:
+            width: The width of the iframe to display the report in.
+            height: The height of the iframe to display the report in.
+        """
+
+        report = self.fetch_report()
+        # Create a data URI from the report HTML
+        report.display_report_in_notebook()
+
+    def save_report(self, path: str | Path) -> None:
+        """Save the evaluation report to a file.
+
+        Args:
+            path: The path to save the report to.
+        """
+        report = self.fetch_report()
+        report.save(path)
+
+    def fetch_data(self) -> pd.DataFrame:
+        """Fetch the synthetic data of the job as a pandas DataFrame.
+
+        Returns:
+            A pandas DataFrame containing the synthetic data.
+        """
+        response = self._client.safe_synthesizer.jobs.results.download_synthetic_data(
+            self.job_name, workspace=self._workspace
+        )
+        return pd.read_csv(BytesIO(response.read()))
+
+    def _fetch_logs_incremental(
+        self, page_cursor: str | None = None, timeout: float | None = None
+    ) -> tuple[list[PlatformJobLog], str | None]:
+        """Fetch logs incrementally starting from a page cursor.
+
+        Fetches all pages from the given cursor position to the end,
+        returning the collected logs and a cursor for the next poll.
+
+        The returned cursor points to the last page fetched, so subsequent
+        polls will re-fetch that page (to catch any newly appended logs)
+        plus any new pages. Use deduplication on the caller side to avoid
+        printing duplicate logs.
+
+        Args:
+            page_cursor: Cursor to start from (None = start from beginning).
+            timeout: Timeout in seconds. Defaults to 300s (5 minutes).
+
+        Returns:
+            Tuple of (list of log objects, cursor to use for next poll).
+        """
+        if timeout is None:
+            timeout = 300.0
+
+        all_logs: list[PlatformJobLog] = []
+        current_cursor = page_cursor
+        last_cursor_with_data: str | None = page_cursor
+
+        while True:
+            response = self._client.with_options(timeout=timeout).safe_synthesizer.jobs.get_logs(
+                self.job_name, page_cursor=current_cursor, workspace=self._workspace
+            )
+
+            if response.data:
+                all_logs.extend(response.data)
+                # Remember this cursor - it had data
+                last_cursor_with_data = current_cursor
+
+            if response.next_page is None:
+                # No more pages - return the last cursor that had data
+                # so next poll re-fetches that page (in case new logs appended)
+                return all_logs, last_cursor_with_data
+            current_cursor = response.next_page
+
+    def fetch_logs(self, timeout: float | None = None) -> Iterator[PlatformJobLog]:
+        """Fetch the logs of the job as an iterator over log objects.
+
+        Recommended to use print_logs for human-readable output. This method returns an iterator
+        over log objects and is useful for programmatic access.
+
+        Args:
+            timeout: Timeout in seconds for each log request. Defaults to 300 seconds (5 minutes)
+                to handle jobs with many log files. Set to None to use the SDK default (60s).
+
+        Returns:
+            A generator for the log objects.
+        """
+        # Use a longer timeout for log queries since they can be slow when
+        # there are many parquet log files to scan
+        if timeout is None:
+            timeout = 300.0  # 5 minutes default for log queries
+
+        page_cursor = None
+        while True:
+            response = self._client.with_options(timeout=timeout).safe_synthesizer.jobs.get_logs(
+                self.job_name, page_cursor=page_cursor, workspace=self._workspace
+            )
+            yield from response.data
+            if response.next_page is None:
+                break
+            page_cursor = response.next_page
+
+    def print_logs(self, timeout: float | None = None) -> None:
+        """Print the logs of the job to stdout.
+
+        Args:
+            timeout: Timeout in seconds for each log request. Defaults to 300s (5 minutes).
+        """
+        for log in self.fetch_logs(timeout=timeout):
+            print(log.message.strip())