diff --git a/docs/training/index.rst b/docs/training/index.rst index e2eabfae54..416560cf20 100644 --- a/docs/training/index.rst +++ b/docs/training/index.rst @@ -299,6 +299,31 @@ Train pre-configured models from the SageMaker JumpStart hub using ``ModelTraine # Complex queries with filters text_gen = search_public_hub_models(query="@task:text-generation") +**Reading the Catalog Metadata of a Hub Model:** + +.. code-block:: python + + from sagemaker.core.jumpstart.hub.catalog import HubModelCatalogEntry + from sagemaker.core.resources import HubContent + + # Every text-generation model of the public hub, as typed catalog entries + entries = [ + HubModelCatalogEntry.from_hub_content(content) + for content in search_public_hub_models(query="@task:text-generation") + ] + small_sft_models = [ + entry.hub_content_name + for entry in entries + if entry.supports_technique("sft") and entry.model_size in ("<1b", "1b-10b") + ] + + # One model, read from DescribeHubContent (the document fills fields the keyword cap dropped) + content = HubContent.get( + hub_name="SageMakerPublicHub", hub_content_type="Model", hub_content_name="meta-textgeneration-llama-3-2-1b" + ) + entry = HubModelCatalogEntry.from_hub_content(content) + print(entry.license, entry.provider, entry.tasks, entry.fine_tuning_techniques()) + Key points: - ``from_jumpstart_config()`` auto-configures training image, instance type, and default hyperparameters @@ -306,6 +331,7 @@ Key points: - JumpStart provides built-in datasets so you can start training immediately - Supports HuggingFace, XGBoost, CatBoost, LightGBM, and many more frameworks - Use ``list_jumpstart_models()`` and ``search_public_hub_models()`` to discover available models +- Use ``HubModelCatalogEntry`` to read a model's task, license, provider, size, context window, languages, modalities, capabilities, and recipes without parsing ``HubContentSearchKeywords`` by hand :doc:`Full example notebook <../v3-examples/training-examples/jumpstart-training-example>` diff --git a/sagemaker-core/src/sagemaker/core/jumpstart/hub/catalog.py b/sagemaker-core/src/sagemaker/core/jumpstart/hub/catalog.py new file mode 100644 index 0000000000..29b4ec3ae6 --- /dev/null +++ b/sagemaker-core/src/sagemaker/core/jumpstart/hub/catalog.py @@ -0,0 +1,325 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Typed catalog view of a hub model. + +A hub model carries its catalog metadata in two places. ``ListHubContents`` returns up to +50 ``HubContentSearchKeywords`` per model, tagged with a prefix per field (``@task:``, +``@license:``, ``@provider:``, ``@model-size:``, ``@context-window:``, ``@framework:``, +``@language:``, ``@data-type:``, ``@input-modality:``, ``@output-modality:``, +``@model-type:``, ``@capability:``, ``@recipe:``). ``DescribeHubContent`` and +``Search`` return the full ``HubContentDocument`` JSON, whose top-level fields hold the +same values in their original casing without the 50-keyword cap. + +``HubModelCatalogEntry`` reads either source into one object, so a caller filters on +``entry.license`` instead of splitting keyword strings by hand. A keyword, when present, +is authoritative. The document fills in only what the keyword cap dropped. +""" + +from __future__ import absolute_import + +import json +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, Iterable, List, Mapping, Optional + +KEYWORD_PREFIXES = { + "tasks": "@task:", + "license": "@license:", + "provider": "@provider:", + "model_size": "@model-size:", + "context_window": "@context-window:", + "framework": "@framework:", + "languages": "@language:", + "data_types": "@data-type:", + "input_modalities": "@input-modality:", + "output_modalities": "@output-modality:", + "model_type": "@model-type:", +} +CAPABILITY_PREFIX = "@capability:" +RECIPE_PREFIX = "@recipe:" +CUSTOMIZATION_CAPABILITY = "customization" +BEDROCK_CAPABILITY = "bedrock_console" + +_DOCUMENT_FIELDS = { + "tasks": "Task", + "license": "License", + "provider": "Provider", + "model_size": "ModelSize", + "context_window": "ContextWindow", + "framework": "Framework", + "languages": "Languages", + "data_types": "DataType", + "input_modalities": "InputModalities", + "output_modalities": "OutputModalities", + "model_type": "ModelTypes", +} +_LIST_FIELDS = ( + "tasks", + "languages", + "data_types", + "input_modalities", + "output_modalities", +) + + +@dataclass(frozen=True) +class HubRecipe: + """One ``@recipe:`` keyword or ``RecipeCollection`` entry of a hub model. + + The hub tags a recipe as ``@recipe:{type}_{technique}`` or + ``@recipe:{type}_{technique}_{strategy}``, all lowercase: ``@recipe:finetuning_sft_lora`` + is a FineTuning recipe for the SFT technique with the LoRA strategy. + """ + + recipe_type: str + technique: str + strategy: Optional[str] = None + + @classmethod + def from_keyword(cls, keyword: str) -> Optional["HubRecipe"]: + """Parse a ``@recipe:`` keyword, or return ``None`` for any other keyword.""" + if not keyword.lower().startswith(RECIPE_PREFIX): + return None + parts = keyword[len(RECIPE_PREFIX) :].lower().split("_", 2) + if len(parts) < 2 or not parts[0] or not parts[1]: + return None + strategy = parts[2] if len(parts) == 3 and parts[2] else None + return cls(recipe_type=parts[0], technique=parts[1], strategy=strategy) + + @classmethod + def from_document(cls, recipe: Mapping[str, Any]) -> Optional["HubRecipe"]: + """Parse one ``RecipeCollection`` entry of a hub model document.""" + recipe_type = recipe.get("Type") or recipe.get("RecipeType") + technique = recipe.get("CustomizationTechnique") or recipe.get("EvaluationType") + if not recipe_type or not technique: + return None + strategy = recipe.get("Strategy") or recipe.get("TrainingType") + return cls( + recipe_type=str(recipe_type).lower(), + technique=str(technique).lower(), + strategy=str(strategy).lower() if strategy else None, + ) + + +@dataclass +class HubModelCatalogEntry: + """The catalog metadata of one hub model, read from keywords and the document. + + Attributes: + hub_content_name: The model's hub content name (the JumpStart model ID). + hub_content_version: The version the entry was read from. + display_name: ``HubContentDisplayName`` when the hub provides one. + description: ``HubContentDescription`` when the hub provides one. + tasks: Every ``@task:`` keyword, extended with document ``Task`` values the keyword + cap dropped. Keyword values keep the hub's spelling. + license: The ``@license:`` keyword, or the document ``License``. + provider: The ``@provider:`` keyword, or the document ``Provider``. + model_size: The raw ``@model-size:`` value (``"7b"``, ``"10b-70b"``, ``"<1b"``), or + the document ``ModelSize``. + context_window: The raw ``@context-window:`` value, or the document ``ContextWindow``. + framework: The ``@framework:`` keyword, or the document ``Framework``. + languages: Every ``@language:`` keyword, extended with document ``Languages``. + data_types: Every ``@data-type:`` keyword, extended with document ``DataType``. + input_modalities: Every ``@input-modality:`` keyword, extended with the document. + output_modalities: Every ``@output-modality:`` keyword, extended with the document. + model_type: The ``@model-type:`` keyword, or the first document ``ModelTypes`` value. + capabilities: Every ``@capability:`` value plus the document ``Capabilities``, lowercase. + recipes: Every ``@recipe:`` keyword plus the document ``RecipeCollection``. + creation_time: ``CreationTime`` of the version read. + original_creation_time: ``OriginalCreationTime`` from ``ListHubContents``, the date the + model first entered the catalog. ``None`` when read from ``DescribeHubContent``. + search_keywords: The raw ``HubContentSearchKeywords``. + """ + + hub_content_name: str + hub_content_version: Optional[str] = None + display_name: Optional[str] = None + description: Optional[str] = None + tasks: List[str] = field(default_factory=list) + license: Optional[str] = None + provider: Optional[str] = None + model_size: Optional[str] = None + context_window: Optional[str] = None + framework: Optional[str] = None + languages: List[str] = field(default_factory=list) + data_types: List[str] = field(default_factory=list) + input_modalities: List[str] = field(default_factory=list) + output_modalities: List[str] = field(default_factory=list) + model_type: Optional[str] = None + capabilities: List[str] = field(default_factory=list) + recipes: List[HubRecipe] = field(default_factory=list) + creation_time: Optional[datetime] = None + original_creation_time: Optional[datetime] = None + search_keywords: List[str] = field(default_factory=list) + + @property + def customization_supported(self) -> bool: + """``True`` when the hub tags the model with the ``customization`` capability.""" + return CUSTOMIZATION_CAPABILITY in self.capabilities + + @property + def bedrock_supported(self) -> bool: + """``True`` when the hub tags the model with the ``bedrock_console`` capability.""" + return BEDROCK_CAPABILITY in self.capabilities + + def fine_tuning_techniques(self) -> List[str]: + """The techniques of the model's FineTuning recipes (``sft``, ``dpo``, ...), sorted.""" + return sorted({r.technique for r in self.recipes if r.recipe_type == "finetuning"}) + + def supports_technique(self, technique: str) -> bool: + """``True`` when a FineTuning recipe for ``technique`` exists (case-insensitive).""" + return technique.lower() in self.fine_tuning_techniques() + + @classmethod + def from_hub_content(cls, hub_content: Any) -> "HubModelCatalogEntry": + """Build an entry from a ``sagemaker.core.resources.HubContent`` or a boto response. + + Accepts the resource object that ``HubContent.get``, ``HubContent.get_all``, and + ``HubContent.get_all_versions`` return, a ``DescribeHubContent`` response dict, a + ``ListHubContents`` summary dict, or the ``HubContent`` record of a ``Search`` result. + """ + if isinstance(hub_content, Mapping): + return cls._from_boto(hub_content) + return cls._from_fields( + name=_attribute(hub_content, "hub_content_name"), + version=_attribute(hub_content, "hub_content_version"), + display_name=_attribute(hub_content, "hub_content_display_name"), + description=_attribute(hub_content, "hub_content_description"), + keywords=_attribute(hub_content, "hub_content_search_keywords") or [], + document=_parse_document(_attribute(hub_content, "hub_content_document")), + creation_time=_attribute(hub_content, "creation_time"), + original_creation_time=_attribute(hub_content, "original_creation_time"), + ) + + @classmethod + def _from_boto(cls, record: Mapping[str, Any]) -> "HubModelCatalogEntry": + return cls._from_fields( + name=record.get("HubContentName"), + version=record.get("HubContentVersion"), + display_name=record.get("HubContentDisplayName"), + description=record.get("HubContentDescription"), + keywords=record.get("HubContentSearchKeywords") or [], + document=_parse_document(record.get("HubContentDocument")), + creation_time=record.get("CreationTime"), + original_creation_time=record.get("OriginalCreationTime"), + ) + + @classmethod + def _from_fields( + cls, + name: Optional[str], + version: Optional[str], + display_name: Optional[str], + description: Optional[str], + keywords: Iterable[str], + document: Dict[str, Any], + creation_time: Optional[datetime], + original_creation_time: Optional[datetime], + ) -> "HubModelCatalogEntry": + if not name: + raise ValueError("A hub content entry must have a HubContentName.") + keyword_list = [k for k in keywords if isinstance(k, str)] + values = _keyword_values(keyword_list) + entry = cls( + hub_content_name=name, + hub_content_version=version, + display_name=display_name, + description=description, + creation_time=creation_time, + original_creation_time=original_creation_time, + search_keywords=list(keyword_list), + ) + for attribute in KEYWORD_PREFIXES: + keyword_hits = values.get(attribute, []) + document_hits = _document_values(document, _DOCUMENT_FIELDS[attribute]) + if attribute in _LIST_FIELDS: + merged = list(keyword_hits) + merged.extend(v for v in document_hits if v.lower() not in _lower_set(merged)) + setattr(entry, attribute, merged) + else: + chosen = ( + keyword_hits[0] + if keyword_hits + else (document_hits[0] if document_hits else None) + ) + setattr(entry, attribute, chosen) + capabilities = [v.lower() for v in values.get("capabilities", [])] + capabilities.extend( + v.lower() + for v in _document_values(document, "Capabilities") + if v.lower() not in capabilities + ) + entry.capabilities = capabilities + recipes = [r for r in (HubRecipe.from_keyword(k) for k in keyword_list) if r] + for raw in document.get("RecipeCollection") or []: + if isinstance(raw, Mapping): + recipe = HubRecipe.from_document(raw) + if recipe and recipe not in recipes: + recipes.append(recipe) + entry.recipes = recipes + return entry + + +def _attribute(obj: Any, name: str) -> Any: + value = getattr(obj, name, None) + if value is None or type(value).__name__ == "Unassigned": + return None + return value + + +def _parse_document(document: Any) -> Dict[str, Any]: + if isinstance(document, Mapping): + return dict(document) + if isinstance(document, str) and document: + try: + parsed = json.loads(document) + except ValueError: + return {} + return parsed if isinstance(parsed, dict) else {} + return {} + + +def _keyword_values(keywords: Iterable[str]) -> Dict[str, List[str]]: + """Group keyword values by catalog attribute, in hub order.""" + values: Dict[str, List[str]] = {} + for keyword in keywords: + lowered = keyword.lower() + if lowered.startswith(CAPABILITY_PREFIX): + values.setdefault("capabilities", []).append(keyword[len(CAPABILITY_PREFIX) :]) + continue + for attribute, prefix in KEYWORD_PREFIXES.items(): + if lowered.startswith(prefix): + value = keyword[len(prefix) :] + if value: + values.setdefault(attribute, []).append(value) + break + return values + + +def _document_values(document: Mapping[str, Any], field_name: str) -> List[str]: + value = document.get(field_name) + items = value if isinstance(value, list) else [value] + return [item for item in items if isinstance(item, str) and item] + + +def _lower_set(values: Iterable[str]) -> set: + return {v.lower() for v in values} + + +def catalog_entries(hub_contents: Iterable[Any]) -> List[HubModelCatalogEntry]: + """Build a ``HubModelCatalogEntry`` for every item of ``hub_contents``. + + ``hub_contents`` can be a ``HubContent.get_all(...)`` iterator, a list of + ``ListHubContents`` summaries, or a list of ``Search`` ``HubContent`` records. + """ + return [HubModelCatalogEntry.from_hub_content(item) for item in hub_contents] diff --git a/sagemaker-core/tests/unit/jumpstart/hub/test_catalog.py b/sagemaker-core/tests/unit/jumpstart/hub/test_catalog.py new file mode 100644 index 0000000000..2381bfe8e2 --- /dev/null +++ b/sagemaker-core/tests/unit/jumpstart/hub/test_catalog.py @@ -0,0 +1,215 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""``HubModelCatalogEntry`` reads a hub model's catalog metadata from keywords and document. + +Before this class every consumer split ``HubContentSearchKeywords`` by prefix itself: +``jumpstart/search.py``, ``recipe_utils.py``, ``training_session/_utils.py``, and four +scripts in the SageMaker agent skills, each with its own partial field list and none +handling the 50-keyword cap that drops ``@model-size:`` and ``@context-window:`` on models +with many languages. +""" + +from __future__ import absolute_import + +import datetime +import json + +import pytest + +from sagemaker.core.jumpstart.hub.catalog import ( + HubModelCatalogEntry, + HubRecipe, + catalog_entries, +) +from sagemaker.core.resources import HubContent + +KEYWORDS = [ + "@task:text generation", + "@task:summarization", + "@license:apache-2.0", + "@provider:Meta", + "@model-size:1b-10b", + "@context-window:128k", + "@framework:huggingface", + "@language:english", + "@language:german", + "@data-type:text", + "@input-modality:text", + "@output-modality:text", + "@model-type:open_weights", + "@capability:customization", + "@capability:bedrock_console", + "@recipe:finetuning_sft_lora", + "@recipe:finetuning_dpo_lora", + "@recipe:finetuning_cpt", + "@recipe:evaluation_llmasjudge", + "@huggingface-id:meta-llama/Llama-3.2-1B", + "unknown keyword", +] +ORIGINAL = datetime.datetime(2023, 6, 1, tzinfo=datetime.timezone.utc) +CREATED = datetime.datetime(2024, 3, 1, tzinfo=datetime.timezone.utc) + + +def _summary(**overrides): + summary = { + "HubContentName": "meta-textgeneration-llama-3-2-1b", + "HubContentVersion": "2.1.0", + "HubContentDisplayName": "Llama 3.2 1B", + "HubContentDescription": "A small Llama.", + "HubContentSearchKeywords": list(KEYWORDS), + "CreationTime": CREATED, + "OriginalCreationTime": ORIGINAL, + } + summary.update(overrides) + return summary + + +def test_entry_reads_every_keyword_prefix_from_a_list_summary(): + entry = HubModelCatalogEntry.from_hub_content(_summary()) + + assert entry.hub_content_name == "meta-textgeneration-llama-3-2-1b" + assert entry.hub_content_version == "2.1.0" + assert entry.display_name == "Llama 3.2 1B" + assert entry.tasks == ["text generation", "summarization"] + assert entry.license == "apache-2.0" + assert entry.provider == "Meta" + assert entry.model_size == "1b-10b" + assert entry.context_window == "128k" + assert entry.framework == "huggingface" + assert entry.languages == ["english", "german"] + assert entry.data_types == ["text"] + assert entry.input_modalities == ["text"] + assert entry.output_modalities == ["text"] + assert entry.model_type == "open_weights" + assert entry.capabilities == ["customization", "bedrock_console"] + assert entry.customization_supported is True + assert entry.bedrock_supported is True + assert entry.creation_time == CREATED + assert entry.original_creation_time == ORIGINAL + assert entry.search_keywords == KEYWORDS + + +def test_recipes_parse_type_technique_and_optional_strategy(): + entry = HubModelCatalogEntry.from_hub_content(_summary()) + + assert HubRecipe("finetuning", "sft", "lora") in entry.recipes + assert HubRecipe("finetuning", "cpt", None) in entry.recipes + assert HubRecipe("evaluation", "llmasjudge", None) in entry.recipes + assert entry.fine_tuning_techniques() == ["cpt", "dpo", "sft"] + assert entry.supports_technique("SFT") is True + assert entry.supports_technique("rlvr") is False + + +def test_recipe_keyword_parser_ignores_other_keywords_and_malformed_recipes(): + assert HubRecipe.from_keyword("@task:text generation") is None + assert HubRecipe.from_keyword("@recipe:finetuning") is None + assert HubRecipe.from_keyword("@recipe:FineTuning_SFT_LoRA") == HubRecipe( + "finetuning", "sft", "lora" + ) + + +def test_document_fills_fields_the_keyword_cap_dropped(): + document = { + "Task": "Speech Recognition", + "License": "MIT", + "Provider": "OpenAI", + "ModelSize": "1.5B", + "ContextWindow": "30s", + "Languages": ["English", "French"], + "Capabilities": ["CUSTOMIZATION"], + "RecipeCollection": [ + {"Type": "FineTuning", "CustomizationTechnique": "SFT", "Strategy": "LORA"}, + {"Type": "Evaluation", "EvaluationType": "ASR"}, + ], + } + entry = HubModelCatalogEntry.from_hub_content( + _summary( + HubContentSearchKeywords=["@task:speech recognition", "@language:english"], + HubContentDocument=json.dumps(document), + ) + ) + + # A keyword is authoritative; the document only fills what is missing. + assert entry.tasks == ["speech recognition"] + assert entry.languages == ["english", "French"] + assert entry.license == "MIT" + assert entry.provider == "OpenAI" + assert entry.model_size == "1.5B" + assert entry.context_window == "30s" + assert entry.capabilities == ["customization"] + assert entry.recipes == [ + HubRecipe("finetuning", "sft", "lora"), + HubRecipe("evaluation", "asr", None), + ] + + +def test_keyword_wins_over_document_for_scalar_fields(): + entry = HubModelCatalogEntry.from_hub_content( + _summary( + HubContentSearchKeywords=["@license:apache-2.0", "@model-size:7b"], + HubContentDocument=json.dumps({"License": "Apache 2.0", "ModelSize": "8B"}), + ) + ) + + assert entry.license == "apache-2.0" + assert entry.model_size == "7b" + + +def test_malformed_document_and_missing_keywords_give_an_empty_entry(): + entry = HubModelCatalogEntry.from_hub_content( + {"HubContentName": "bare", "HubContentDocument": "not json"} + ) + + assert entry.tasks == [] + assert entry.license is None + assert entry.capabilities == [] + assert entry.recipes == [] + assert entry.customization_supported is False + assert entry.original_creation_time is None + + +def test_entry_requires_a_name(): + with pytest.raises(ValueError, match="HubContentName"): + HubModelCatalogEntry.from_hub_content({"HubContentSearchKeywords": []}) + + +def test_entry_reads_a_hub_content_resource(): + resource = HubContent( + hub_name="SageMakerPublicHub", + hub_content_type="Model", + hub_content_name="meta-textgeneration-llama-3-2-1b", + hub_content_version="2.1.0", + hub_content_search_keywords=["@task:text generation", "@recipe:finetuning_sft_lora"], + hub_content_document=json.dumps({"License": "Llama 3.2 Community"}), + creation_time=CREATED, + ) + + entry = HubModelCatalogEntry.from_hub_content(resource) + + assert entry.hub_content_name == "meta-textgeneration-llama-3-2-1b" + assert entry.tasks == ["text generation"] + assert entry.license == "Llama 3.2 Community" + assert entry.fine_tuning_techniques() == ["sft"] + assert entry.creation_time == CREATED + # Unassigned() on the resource reads as None, not as the sentinel object. + assert entry.display_name is None + assert entry.original_creation_time is None + + +def test_catalog_entries_maps_every_item(): + entries = catalog_entries([_summary(), _summary(HubContentName="other")]) + + assert [entry.hub_content_name for entry in entries] == [ + "meta-textgeneration-llama-3-2-1b", + "other", + ]