Get Model Names

1import json
2import sys
3
4import boto3
5
6if len(sys.argv) < 2:
7    print("Usage: python get_model_names.py <hub-name> [region]")
8    sys.exit(1)
9
10hub_name = sys.argv[1]
11region_name = sys.argv[2] if len(sys.argv) > 2 else None
12
13sm_client = boto3.client("sagemaker", region_name=region_name)
14
15# Retrieve all models with pagination
16all_contents = []
17next_token = None
18
19while True:
20    params = {"HubName": hub_name, "HubContentType": "Model", "MaxResults": 100}
21
22    if next_token:
23        params["NextToken"] = next_token
24
25    response = sm_client.list_hub_contents(**params)
26    all_contents.extend(response.get("HubContentSummaries", []))
27
28    next_token = response.get("NextToken")
29    if not next_token:
30        break
31
32# Filter for customization-capable models
33customization_models = [
34    content
35    for content in all_contents
36    if "@capability:customization" in content.get("HubContentSearchKeywords", [])
37]
38
39model_names = [m.get("HubContentName") for m in customization_models]
40
41print(json.dumps(model_names))