diff --git a/docs/mri_advanced.md b/docs/mri_advanced.md index d52483f4c..df68e6696 100644 --- a/docs/mri_advanced.md +++ b/docs/mri_advanced.md @@ -50,26 +50,38 @@ benchmark = client.mri.create_new_benchmark( Tags provide metadata for filtering and organizing benchmark results without showing them to evaluators. These tags can also be set and used in the frontend. To view the frontend, you can use the `view` method of the benchmark or leaderboard. +Each tag is either a plain string or a `Tag` with a `value` and an optional `category` — you can mix both in the same list. Bare strings are wrapped into `Tag(value, category=None)` for you. You can also record where each prompt came from with an `Origin`. + ```python -# Tags for filtering leaderboard results +from rapidata import Tag, Origin + +# Tags per prompt — plain strings and Tag objects can be mixed tags = [ - ["landscape", "outdoor", "beach"], - ["landscape", "outdoor", "mountain"], - ["outdoor", "city"], - ["indoor", "vehicle"] + [Tag("beach", category="scene"), "outdoor"], + [Tag("mountain", category="scene"), "outdoor"], + [Tag("city", category="scene")], + [Tag("vehicle", category="object"), "indoor"], ] benchmark = client.mri.create_new_benchmark( name="Tagged Benchmark", identifiers=["scene_1", "scene_2", "scene_3", "scene_4"], prompts=["A sunny beach", "A mountain landscape", "A city skyline", "A car in a garage"], - tags=tags + tags=tags, + origins=[Origin("coco"), Origin("coco"), Origin("coco"), Origin("coco")], ) -# Filter leaderboard results by tags -standings = leaderboard.get_standings(tags=["landscape", "outdoor"]) +# Filter leaderboard results by tag value +standings = leaderboard.get_standings(tags=["outdoor"]) ``` +!!! note + Plain-string tags keep working (e.g. `tags=[["landscape", "outdoor"]]`); they are + treated as tags with no category. Read tags back with the values-only + `benchmark.tags` (`list[list[str]]`) for backwards compatibility, or the + structured `benchmark.structured_tags` (`list[list[Tag]]`) for categories, and + prompt origins with `benchmark.origins`. + ### Adding prompts and assets after benchmark creation If you have already created a benchmark and want to add new prompts and assets after the fact. Note however that these will only take effect for new models. @@ -80,7 +92,7 @@ benchmark.add_prompts( identifiers=["new_style"], prompts=["Generate artwork in this new style"], prompt_assets=["https://assets.rapidata.ai/new_style_ref.jpg"], - tags=[["abstract", "modern"]] + tags=[[Tag("abstract", category="style"), "modern"]], ) ``` diff --git a/pyproject.toml b/pyproject.toml index ff8ce0045..3013cca82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "rapidata" -version = "3.17.2" +version = "3.18.0" description = "Rapidata package containing the Rapidata Python Client to interact with the Rapidata Web API in an easy way." authors = [{ name = "Rapidata AG", email = "info@rapidata.ai" }] requires-python = ">=3.10,<4" diff --git a/src/rapidata/__init__.py b/src/rapidata/__init__.py index 0e3909325..2ce63c6f7 100644 --- a/src/rapidata/__init__.py +++ b/src/rapidata/__init__.py @@ -1,4 +1,4 @@ -__version__ = "3.17.2" +__version__ = "3.18.0" from .rapidata_client import ( RapidataClient, @@ -64,6 +64,8 @@ Gender, DeviceFilter, DeviceType, + Tag, + Origin, Datapoint, ContextManager, FailedUploadException, diff --git a/src/rapidata/api_client/models/create_prompt_for_benchmark_endpoint_input.py b/src/rapidata/api_client/models/create_prompt_for_benchmark_endpoint_input.py index a7b8cbbfe..94cf1c2c7 100644 --- a/src/rapidata/api_client/models/create_prompt_for_benchmark_endpoint_input.py +++ b/src/rapidata/api_client/models/create_prompt_for_benchmark_endpoint_input.py @@ -19,6 +19,8 @@ from pydantic import BaseModel, ConfigDict, Field, StrictStr from typing import Any, ClassVar, Dict, List, Optional from rapidata.api_client.models.i_asset_input import IAssetInput +from rapidata.api_client.models.prompt_origin import PromptOrigin +from rapidata.api_client.models.prompt_tagging import PromptTagging from pydantic import ValidationError from rapidata.api_client.lazy_model import LazyValidatedModel from typing import Optional, Set @@ -31,8 +33,9 @@ class CreatePromptForBenchmarkEndpointInput(LazyValidatedModel): identifier: StrictStr = Field(description="An identifier associated with the prompt.") prompt: Optional[StrictStr] = Field(default=None, description="The prompt text.") prompt_asset: Optional[IAssetInput] = Field(default=None, description="The asset associated with the prompt.", alias="promptAsset") - tags: Optional[List[StrictStr]] = None - __properties: ClassVar[List[str]] = ["identifier", "prompt", "promptAsset", "tags"] + tags: Optional[List[PromptTagging]] = Field(default=None, description="The tags associated with the prompt. Each tag has a value and an optional category.") + origin: Optional[PromptOrigin] = Field(default=None, description="The origin of the prompt.") + __properties: ClassVar[List[str]] = ["identifier", "prompt", "promptAsset", "tags", "origin"] # model_config is inherited from LazyValidatedModel @@ -72,6 +75,16 @@ def to_dict(self) -> Dict[str, Any]: # override the default output from pydantic by calling `to_dict()` of prompt_asset if self.prompt_asset: _dict['promptAsset'] = self.prompt_asset.to_dict() + # override the default output from pydantic by calling `to_dict()` of each item in tags (list) + _items = [] + if self.tags: + for _item_tags in self.tags: + if _item_tags: + _items.append(_item_tags.to_dict()) + _dict['tags'] = _items + # override the default output from pydantic by calling `to_dict()` of origin + if self.origin: + _dict['origin'] = self.origin.to_dict() # set to None if prompt (nullable) is None # and model_fields_set contains the field if self.prompt is None and "prompt" in self.model_fields_set: @@ -82,6 +95,11 @@ def to_dict(self) -> Dict[str, Any]: if self.tags is None and "tags" in self.model_fields_set: _dict['tags'] = None + # set to None if origin (nullable) is None + # and model_fields_set contains the field + if self.origin is None and "origin" in self.model_fields_set: + _dict['origin'] = None + return _dict @classmethod @@ -97,12 +115,11 @@ def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]: "identifier": obj.get("identifier"), "prompt": obj.get("prompt"), "promptAsset": IAssetInput.from_dict(obj["promptAsset"]) if obj.get("promptAsset") is not None else None, - "tags": obj.get("tags") + "tags": [PromptTagging.from_dict(_item) for _item in obj["tags"]] if obj.get("tags") is not None else None, + "origin": PromptOrigin.from_dict(obj["origin"]) if obj.get("origin") is not None else None } try: _obj = cls.model_validate(_data) except ValidationError as _val_error: _obj = cls._lazy_construct(_data, _val_error) return _obj - - diff --git a/src/rapidata/api_client/models/get_prompts_by_benchmark_endpoint_output.py b/src/rapidata/api_client/models/get_prompts_by_benchmark_endpoint_output.py index ff7973b2d..e079068d9 100644 --- a/src/rapidata/api_client/models/get_prompts_by_benchmark_endpoint_output.py +++ b/src/rapidata/api_client/models/get_prompts_by_benchmark_endpoint_output.py @@ -20,6 +20,8 @@ from pydantic import BaseModel, ConfigDict, Field, StrictStr from typing import Any, ClassVar, Dict, List, Optional from rapidata.api_client.models.i_asset import IAsset +from rapidata.api_client.models.prompt_origin import PromptOrigin +from rapidata.api_client.models.prompt_tagging import PromptTagging from pydantic import ValidationError from rapidata.api_client.lazy_model import LazyValidatedModel from typing import Optional, Set @@ -35,8 +37,9 @@ class GetPromptsByBenchmarkEndpointOutput(LazyValidatedModel): prompt_asset: Optional[IAsset] = Field(default=None, description="The optional asset associated with the prompt.", alias="promptAsset") identifier: StrictStr = Field(description="The identifier associated with the prompt.") created_at: datetime = Field(description="The timestamp when the prompt was created.", alias="createdAt") - tags: List[StrictStr] - __properties: ClassVar[List[str]] = ["id", "englishPrompt", "originalPrompt", "promptAsset", "identifier", "createdAt", "tags"] + tags: List[PromptTagging] + origin: Optional[PromptOrigin] = None + __properties: ClassVar[List[str]] = ["id", "englishPrompt", "originalPrompt", "promptAsset", "identifier", "createdAt", "tags", "origin"] # model_config is inherited from LazyValidatedModel @@ -76,6 +79,16 @@ def to_dict(self) -> Dict[str, Any]: # override the default output from pydantic by calling `to_dict()` of prompt_asset if self.prompt_asset: _dict['promptAsset'] = self.prompt_asset.to_dict() + # override the default output from pydantic by calling `to_dict()` of each item in tags (list) + _items = [] + if self.tags: + for _item_tags in self.tags: + if _item_tags: + _items.append(_item_tags.to_dict()) + _dict['tags'] = _items + # override the default output from pydantic by calling `to_dict()` of origin + if self.origin: + _dict['origin'] = self.origin.to_dict() # set to None if english_prompt (nullable) is None # and model_fields_set contains the field if self.english_prompt is None and "english_prompt" in self.model_fields_set: @@ -86,6 +99,11 @@ def to_dict(self) -> Dict[str, Any]: if self.original_prompt is None and "original_prompt" in self.model_fields_set: _dict['originalPrompt'] = None + # set to None if origin (nullable) is None + # and model_fields_set contains the field + if self.origin is None and "origin" in self.model_fields_set: + _dict['origin'] = None + return _dict @classmethod @@ -104,7 +122,9 @@ def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]: "promptAsset": IAsset.from_dict(obj["promptAsset"]) if obj.get("promptAsset") is not None else None, "identifier": obj.get("identifier"), "createdAt": obj.get("createdAt"), - "tags": obj.get("tags") + # Wire sends Tag objects; tolerate a bare string (older backend) by wrapping it as an uncategorized tag. + "tags": [PromptTagging.from_dict(_item) if isinstance(_item, dict) else PromptTagging(value=_item) for _item in obj["tags"]] if obj.get("tags") is not None else None, + "origin": PromptOrigin.from_dict(obj["origin"]) if obj.get("origin") is not None else None } try: _obj = cls.model_validate(_data) diff --git a/src/rapidata/api_client/models/prompt_origin.py b/src/rapidata/api_client/models/prompt_origin.py new file mode 100644 index 000000000..d3adfb6ec --- /dev/null +++ b/src/rapidata/api_client/models/prompt_origin.py @@ -0,0 +1,87 @@ +# coding: utf-8 + +""" + Rapidata Asset API + + The API for the Rapidata Asset service + + The version of the OpenAPI document: v1 + Generated by OpenAPI Generator (https://openapi-generator.tech) + + Do not edit the class manually. +""" # noqa: E501 + + +from __future__ import annotations +import pprint +import re # noqa: F401 +import json + +from pydantic import BaseModel, ConfigDict, Field, StrictStr +from typing import Any, ClassVar, Dict, List, Optional +from pydantic import ValidationError +from rapidata.api_client.lazy_model import LazyValidatedModel +from typing import Optional, Set +from typing_extensions import Self + +class PromptOrigin(LazyValidatedModel): + """ + PromptOrigin + """ # noqa: E501 + value: StrictStr = Field(description="The origin value.") + __properties: ClassVar[List[str]] = ["value"] + + # model_config is inherited from LazyValidatedModel + + + def to_str(self) -> str: + """Returns the string representation of the model using alias""" + return pprint.pformat(self.model_dump(by_alias=True)) + + def to_json(self) -> str: + """Returns the JSON representation of the model using alias""" + # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead + return json.dumps(self.to_dict()) + + @classmethod + def from_json(cls, json_str: str) -> Optional[Self]: + """Create an instance of PromptOrigin from a JSON string""" + return cls.from_dict(json.loads(json_str)) + + def to_dict(self) -> Dict[str, Any]: + """Return the dictionary representation of the model using alias. + + This has the following differences from calling pydantic's + `self.model_dump(by_alias=True)`: + + * `None` is only added to the output dict for nullable fields that + were set at model initialization. Other fields with value `None` + are ignored. + """ + excluded_fields: Set[str] = set([ + ]) + + _dict = self.model_dump( + by_alias=True, + exclude=excluded_fields, + exclude_none=True, + ) + return _dict + + @classmethod + def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]: + """Create an instance of PromptOrigin from a dict""" + if obj is None: + return None + + if not isinstance(obj, dict): + return cls.model_validate(obj) + + _data = { + "value": obj.get("value") + } + try: + _obj = cls.model_validate(_data) + except ValidationError as _val_error: + _obj = cls._lazy_construct(_data, _val_error) + return _obj diff --git a/src/rapidata/api_client/models/prompt_tagging.py b/src/rapidata/api_client/models/prompt_tagging.py new file mode 100644 index 000000000..500f5f1a8 --- /dev/null +++ b/src/rapidata/api_client/models/prompt_tagging.py @@ -0,0 +1,94 @@ +# coding: utf-8 + +""" + Rapidata Asset API + + The API for the Rapidata Asset service + + The version of the OpenAPI document: v1 + Generated by OpenAPI Generator (https://openapi-generator.tech) + + Do not edit the class manually. +""" # noqa: E501 + + +from __future__ import annotations +import pprint +import re # noqa: F401 +import json + +from pydantic import BaseModel, ConfigDict, Field, StrictStr +from typing import Any, ClassVar, Dict, List, Optional +from pydantic import ValidationError +from rapidata.api_client.lazy_model import LazyValidatedModel +from typing import Optional, Set +from typing_extensions import Self + +class PromptTagging(LazyValidatedModel): + """ + PromptTagging + """ # noqa: E501 + value: StrictStr = Field(description="The tag value.") + category: Optional[StrictStr] = Field(default=None, description="The optional category the tag belongs to.") + __properties: ClassVar[List[str]] = ["value", "category"] + + # model_config is inherited from LazyValidatedModel + + + def to_str(self) -> str: + """Returns the string representation of the model using alias""" + return pprint.pformat(self.model_dump(by_alias=True)) + + def to_json(self) -> str: + """Returns the JSON representation of the model using alias""" + # TODO: pydantic v2: use .model_dump_json(by_alias=True, exclude_unset=True) instead + return json.dumps(self.to_dict()) + + @classmethod + def from_json(cls, json_str: str) -> Optional[Self]: + """Create an instance of PromptTagging from a JSON string""" + return cls.from_dict(json.loads(json_str)) + + def to_dict(self) -> Dict[str, Any]: + """Return the dictionary representation of the model using alias. + + This has the following differences from calling pydantic's + `self.model_dump(by_alias=True)`: + + * `None` is only added to the output dict for nullable fields that + were set at model initialization. Other fields with value `None` + are ignored. + """ + excluded_fields: Set[str] = set([ + ]) + + _dict = self.model_dump( + by_alias=True, + exclude=excluded_fields, + exclude_none=True, + ) + # set to None if category (nullable) is None + # and model_fields_set contains the field + if self.category is None and "category" in self.model_fields_set: + _dict['category'] = None + + return _dict + + @classmethod + def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]: + """Create an instance of PromptTagging from a dict""" + if obj is None: + return None + + if not isinstance(obj, dict): + return cls.model_validate(obj) + + _data = { + "value": obj.get("value"), + "category": obj.get("category") + } + try: + _obj = cls.model_validate(_data) + except ValidationError as _val_error: + _obj = cls._lazy_construct(_data, _val_error) + return _obj diff --git a/src/rapidata/rapidata_client/__init__.py b/src/rapidata/rapidata_client/__init__.py index 0ebd59f92..99984027f 100644 --- a/src/rapidata/rapidata_client/__init__.py +++ b/src/rapidata/rapidata_client/__init__.py @@ -27,6 +27,7 @@ RapidataRetrievalMode, EffortSelection, ) +from .benchmark.prompt_metadata import Origin, Tag from .datapoints import Datapoint from .context import ContextManager from .datapoints.metadata import ( diff --git a/src/rapidata/rapidata_client/benchmark/_prompt_uploader.py b/src/rapidata/rapidata_client/benchmark/_prompt_uploader.py index 400fda603..d083c3512 100644 --- a/src/rapidata/rapidata_client/benchmark/_prompt_uploader.py +++ b/src/rapidata/rapidata_client/benchmark/_prompt_uploader.py @@ -8,6 +8,7 @@ from tqdm.auto import tqdm from rapidata.rapidata_client.config import logger, rapidata_config, tracer +from rapidata.rapidata_client.benchmark.prompt_metadata import Origin, Tag from rapidata.rapidata_client.datapoints._asset_uploader import AssetUploader if TYPE_CHECKING: @@ -16,12 +17,18 @@ @dataclass class BenchmarkPrompt: - """A single prompt to register on a benchmark.""" + """A single prompt to register on a benchmark. + + ``tags`` are always structured :class:`Tag` objects here; the high-level + API wraps any bare strings into ``Tag(value, category=None)`` before + constructing a ``BenchmarkPrompt``. + """ identifier: str prompt: str | None = None prompt_asset: str | None = None - tags: list[str] = field(default_factory=list) + tags: list[Tag] = field(default_factory=list) + origin: Origin | None = None class BenchmarkPromptUploader: @@ -45,6 +52,8 @@ def upload(self, prompt: BenchmarkPrompt) -> None: IAssetInputExistingAssetInput, ) from rapidata.api_client.models.i_asset_input import IAssetInput + from rapidata.api_client.models.prompt_origin import PromptOrigin + from rapidata.api_client.models.prompt_tagging import PromptTagging self._openapi_service.leaderboard.benchmark_api.benchmark_benchmark_id_prompt_post( benchmark_id=self._benchmark_id, @@ -61,7 +70,15 @@ def upload(self, prompt: BenchmarkPrompt) -> None: if prompt.prompt_asset is not None else None ), - tags=prompt.tags, + tags=[ + PromptTagging(value=tag.value, category=tag.category) + for tag in prompt.tags + ], + origin=( + PromptOrigin(value=prompt.origin.value) + if prompt.origin is not None + else None + ), ), ) diff --git a/src/rapidata/rapidata_client/benchmark/prompt_metadata.py b/src/rapidata/rapidata_client/benchmark/prompt_metadata.py new file mode 100644 index 000000000..030230aa9 --- /dev/null +++ b/src/rapidata/rapidata_client/benchmark/prompt_metadata.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class Tag: + """A structured tag attached to a benchmark prompt. + + Tags are metadata used to filter and organize benchmark results; they are + never shown to the annotators. ``category`` optionally groups related tags + (e.g. ``Tag("landscape", category="scene")``); leave it ``None`` for a bare + tag. + """ + + value: str + category: str | None = None + + +@dataclass +class Origin: + """Where a benchmark prompt originated from (e.g. a source dataset).""" + + value: str diff --git a/src/rapidata/rapidata_client/benchmark/rapidata_benchmark.py b/src/rapidata/rapidata_client/benchmark/rapidata_benchmark.py index e1e8729c8..fcb5cb1a8 100644 --- a/src/rapidata/rapidata_client/benchmark/rapidata_benchmark.py +++ b/src/rapidata/rapidata_client/benchmark/rapidata_benchmark.py @@ -11,6 +11,7 @@ BenchmarkPrompt, BenchmarkPromptUploader, ) +from rapidata.rapidata_client.benchmark.prompt_metadata import Origin, Tag from rapidata.api_client.models.audience_audience_id_jobs_get_job_id_parameter import ( AudienceAudienceIdJobsGetJobIdParameter, ) @@ -50,6 +51,8 @@ def __init__(self, name: str, id: str, openapi_service: OpenAPIService): self.__leaderboards: list["RapidataLeaderboard"] = [] self.__identifiers: list[str] = [] self.__tags: list[list[str]] = [] + self.__structured_tags: list[list[Tag]] = [] + self.__origins: list[Origin | None] = [] self.__participants: list[BenchmarkParticipant] = [] self.__benchmark_page: str = ( f"https://app.{self._openapi_service.environment}/mri/benchmarks/{self.id}" @@ -74,6 +77,8 @@ def __instantiate_prompts(self) -> None: self.__identifiers = [] self.__prompt_assets = [] self.__tags = [] + self.__structured_tags = [] + self.__origins = [] current_page = 1 total_pages = None @@ -118,7 +123,17 @@ def __instantiate_prompts(self) -> None: else: self.__prompt_assets.append(None) - self.__tags.append(prompt.tags) + prompt_tags = [ + Tag(value=tag.value, category=tag.category) + for tag in prompt.tags + ] + self.__structured_tags.append(prompt_tags) + self.__tags.append([tag.value for tag in prompt_tags]) + self.__origins.append( + Origin(value=prompt.origin.value) + if prompt.origin is not None + else None + ) if current_page >= total_pages: break @@ -186,13 +201,43 @@ def prompt_assets(self) -> list[str | None]: @property def tags(self) -> list[list[str]]: """ - Returns the tags that are registered for the benchmark. + Returns the tag values registered for the benchmark, aligned by index with `prompts`. + + This is the flat, values-only view of the tags. It is kept for + backwards compatibility — prefer `structured_tags` for the + (value + category) representation. """ if not self.__tags: self.__instantiate_prompts() return self.__tags + @property + def structured_tags(self) -> list[list[Tag]]: + """ + Returns the structured tags registered for the benchmark, aligned by index with `prompts`. + + Each :class:`Tag` carries a ``value`` and an optional ``category``. Tags + are used to filter and organize leaderboard results and are NOT shown to + the annotators. + """ + if not self.__structured_tags: + self.__instantiate_prompts() + + return self.__structured_tags + + @property + def origins(self) -> list[Origin | None]: + """ + Returns the origin of each prompt, aligned by index with `prompts`. + + A prompt without an origin is represented as ``None``. + """ + if not self.__origins: + self.__instantiate_prompts() + + return self.__origins + @property def leaderboards(self) -> list[RapidataLeaderboard]: """ @@ -277,7 +322,8 @@ def add_prompts( identifiers: Optional[list[str]] = None, prompts: Optional[list[str | None] | list[str]] = None, prompt_assets: Optional[list[str | None] | list[str]] = None, - tags: Optional[list[list[str] | None] | list[list[str]]] = None, + tags: Optional[list[list[str | Tag] | None]] = None, + origins: Optional[list[Origin | str | None]] = None, ) -> None: """ Adds one or more prompts to the benchmark. Everything is matched up by the @@ -293,15 +339,19 @@ def add_prompts( identifiers: The identifiers of the prompts/assets/tags that will be used to match up the media. If not provided, it will use the prompts as the identifiers. prompts: The prompts that will be registered for the benchmark. prompt_assets: The prompt assets that will be registered for the benchmark. - tags: The tags that will be associated with the prompts to use for filtering the leaderboard results. They will NOT be shown to the users. + tags: The tags per prompt, used to filter and organize the leaderboard results. They are NOT shown to the users. Each entry is a list whose items are either a plain string (mapped to `Tag(value, category=None)`) or a :class:`Tag` with a `value` and optional `category`, or None for no tags. + origins: The origin of each prompt (e.g. a source dataset). Each entry is an :class:`Origin`, a plain string (mapped to `Origin(value)`), or None. Example: ```python + from rapidata import Tag, Origin + benchmark.add_prompts( identifiers=["id1", "id2"], prompts=["prompt 1", "prompt 2"], prompt_assets=["https://assets.rapidata.ai/prompt_1.jpg", "https://assets.rapidata.ai/prompt_2.jpg"], - tags=[["tag1", "tag2"], ["tag2"]], + tags=[[Tag("landscape", category="scene"), "outdoor"], [Tag("portrait")]], + origins=[Origin("coco"), Origin("coco")], ) ``` """ @@ -333,17 +383,24 @@ def add_prompts( if tags is not None: if not isinstance(tags, list): - raise ValueError("Tags must be a list of lists of strings or None.") + raise ValueError("Tags must be a list of lists of str/Tag or None.") for tag in tags: if tag is not None and ( not isinstance(tag, list) - or not all(isinstance(item, str) for item in tag) + or not all(isinstance(item, (str, Tag)) for item in tag) ): raise ValueError( - "Tags must be a list of lists of strings or None." + "Tags must be a list of lists of str/Tag or None." ) + if origins is not None: + if not isinstance(origins, list) or not all( + origin is None or isinstance(origin, (Origin, str)) + for origin in origins + ): + raise ValueError("Origins must be a list of Origin/str or None.") + if not identifiers and not prompts: raise ValueError( "At least one of identifiers or prompts must be provided." @@ -378,11 +435,20 @@ def add_prompts( prompt_assets = cast(list[str | None], [None] * expected_length) if not tags: - tags = cast(list[list[str] | None], [None] * expected_length) + tags = cast(list[list[str | Tag] | None], [None] * expected_length) - if not (expected_length == len(prompts) == len(prompt_assets) == len(tags)): + if not origins: + origins = cast(list[Origin | str | None], [None] * expected_length) + + if not ( + expected_length + == len(prompts) + == len(prompt_assets) + == len(tags) + == len(origins) + ): raise ValueError( - "Identifiers, prompts, media assets, and tags must have the same length or set to None." + "Identifiers, prompts, media assets, tags, and origins must have the same length or set to None." ) # Snapshot once: `self.identifiers` is a property whose getter re-fetches @@ -400,12 +466,27 @@ def add_prompts( f"Identifiers already exist in the benchmark: {already_registered}" ) + def build_tags(tag: list[str | Tag] | None) -> list[Tag]: + # Bare strings are wrapped into an uncategorized Tag; Tag objects + # pass through untouched. + return [ + Tag(value=item) if isinstance(item, str) else item + for item in (tag or []) + ] + + def build_origin(origin: Origin | str | None) -> Origin | None: + return Origin(value=origin) if isinstance(origin, str) else origin + to_upload = [ BenchmarkPrompt( - identifier, prompt, asset, tag if tag is not None else [] + identifier, + prompt, + asset, + build_tags(tag), + build_origin(origin), ) - for identifier, prompt, asset, tag in zip( - identifiers, prompts, prompt_assets, tags + for identifier, prompt, asset, tag, origin in zip( + identifiers, prompts, prompt_assets, tags, origins ) ] @@ -415,7 +496,9 @@ def add_prompts( self.__prompt_assets.append( self.__normalize_cached_asset(uploaded.prompt_asset) ) - self.__tags.append(uploaded.tags) + self.__structured_tags.append(uploaded.tags) + self.__tags.append([tag.value for tag in uploaded.tags]) + self.__origins.append(uploaded.origin) # The English translation is produced server-side and is unknown for # the just-added prompts. Clear it so the next access lazily re-fetches diff --git a/src/rapidata/rapidata_client/benchmark/rapidata_benchmark_manager.py b/src/rapidata/rapidata_client/benchmark/rapidata_benchmark_manager.py index 750bd5760..c698c7925 100644 --- a/src/rapidata/rapidata_client/benchmark/rapidata_benchmark_manager.py +++ b/src/rapidata/rapidata_client/benchmark/rapidata_benchmark_manager.py @@ -1,5 +1,6 @@ from typing import Optional from rapidata.rapidata_client.benchmark.rapidata_benchmark import RapidataBenchmark +from rapidata.rapidata_client.benchmark.prompt_metadata import Origin, Tag from rapidata.api_client.models.create_benchmark_endpoint_input import ( CreateBenchmarkEndpointInput, ) @@ -28,7 +29,8 @@ def create_new_benchmark( identifiers: Optional[list[str]] = None, prompts: Optional[list[str | None] | list[str]] = None, prompt_assets: Optional[list[str | None] | list[str]] = None, - tags: Optional[list[list[str] | None] | list[list[str]]] = None, + tags: Optional[list[list[str | Tag] | None]] = None, + origins: Optional[list[Origin | str | None]] = None, ) -> RapidataBenchmark: """ Creates a new benchmark with the given name, identifiers, prompts, and media assets. @@ -41,17 +43,21 @@ def create_new_benchmark( identifiers: The identifiers of the prompts/assets/tags that will be used to match up the media. If not provided, it will use the prompts as the identifiers. prompts: The prompts that will be registered for the benchmark. prompt_assets: The prompt assets that will be registered for the benchmark. - tags: The tags that will be associated with the prompts to use for filtering the leaderboard results. They will NOT be shown to the users. + tags: The tags per prompt, used to filter and organize the leaderboard results. They are NOT shown to the users. Each entry is a list whose items are either a plain string (mapped to `Tag(value, category=None)`) or a :class:`Tag` with a `value` and optional `category`, or None for no tags. + origins: The origin of each prompt (e.g. a source dataset). Each entry is an :class:`Origin`, a plain string (mapped to `Origin(value)`), or None. Example: ```python + from rapidata import Tag, Origin + name = "Example Benchmark" identifiers = ["id1", "id2", "id3"] prompts = ["prompt 1", "prompt 2", "prompt 3"] prompt_assets = ["https://assets.rapidata.ai/prompt_1.jpg", "https://assets.rapidata.ai/prompt_2.jpg", "https://assets.rapidata.ai/prompt_3.jpg"] - tags = [["tag1", "tag2"], ["tag2"], ["tag2", "tag3"]] + tags = [[Tag("tag1", category="group"), "tag2"], [Tag("tag2")], ["tag2", "tag3"]] + origins = [Origin("coco"), Origin("coco"), Origin("coco")] - benchmark = create_new_benchmark(name=name, identifiers=identifiers, prompts=prompts, prompt_assets=prompt_assets, tags=tags) + benchmark = create_new_benchmark(name=name, identifiers=identifiers, prompts=prompts, prompt_assets=prompt_assets, tags=tags, origins=origins) ``` """ with tracer.start_as_current_span( @@ -76,7 +82,7 @@ def create_new_benchmark( name, benchmark_result.id, self.__openapi_service ) - benchmark.add_prompts(identifiers, prompts, prompt_assets, tags) + benchmark.add_prompts(identifiers, prompts, prompt_assets, tags, origins) return benchmark diff --git a/src/rapidata/types/__init__.py b/src/rapidata/types/__init__.py index bbfd1946b..5f3bca029 100644 --- a/src/rapidata/types/__init__.py +++ b/src/rapidata/types/__init__.py @@ -18,6 +18,7 @@ from rapidata.rapidata_client.benchmark.participant.participant import ( BenchmarkParticipant, ) +from rapidata.rapidata_client.benchmark.prompt_metadata import Origin, Tag # Selection Types from rapidata.rapidata_client.selection.ab_test_selection import AbTestSelection @@ -108,6 +109,8 @@ "RapidataBenchmarkManager", "RapidataLeaderboard", "BenchmarkParticipant", + "Tag", + "Origin", # Selection Types "AbTestSelection", "CappedSelection", diff --git a/uv.lock b/uv.lock index 405824b17..9c18e523b 100644 --- a/uv.lock +++ b/uv.lock @@ -2610,7 +2610,7 @@ wheels = [ [[package]] name = "rapidata" -version = "3.17.1" +version = "3.18.0" source = { editable = "." } dependencies = [ { name = "authlib" },