import re
import logging
from typing import List, Iterator, Optional
from sagemaker.core.helper.session_helper import Session
from sagemaker.core.resources import HubContent
logger = logging.getLogger(__name__)
class _ExpressionNode:
"""Base class for expression AST nodes."""
def evaluate(self, keywords: List[str]) -> bool:
"""Evaluate this node against the given keywords."""
raise NotImplementedError
class _AndNode(_ExpressionNode):
"""AND logical operator node."""
def __init__(self, left: _ExpressionNode, right: _ExpressionNode):
self.left = left
self.right = right
def evaluate(self, keywords: List[str]) -> bool:
return self.left.evaluate(keywords) and self.right.evaluate(keywords)
class _OrNode(_ExpressionNode):
"""OR logical operator node."""
def __init__(self, left: _ExpressionNode, right: _ExpressionNode):
self.left = left
self.right = right
def evaluate(self, keywords: List[str]) -> bool:
return self.left.evaluate(keywords) or self.right.evaluate(keywords)
class _NotNode(_ExpressionNode):
"""NOT logical operator node."""
def __init__(self, operand: _ExpressionNode):
self.operand = operand
def evaluate(self, keywords: List[str]) -> bool:
return not self.operand.evaluate(keywords)
class _PatternNode(_ExpressionNode):
"""Pattern matching node for keywords with wildcard support."""
def __init__(self, pattern: str):
self.pattern = pattern.strip('"').strip("'")
def evaluate(self, keywords: List[str]) -> bool:
"""Check if any keyword matches this pattern."""
for keyword in keywords:
if self._matches_pattern(keyword, self.pattern):
return True
return False
def _matches_pattern(self, keyword: str, pattern: str) -> bool:
"""Check if a keyword matches a pattern with wildcard support."""
if pattern.startswith("*") and pattern.endswith("*"):
# Contains pattern: *text*
stripped = pattern.strip("*")
return stripped in keyword
elif pattern.startswith("*"):
# Ends with pattern: *text
stripped = pattern[1:]
return keyword.endswith(stripped)
elif pattern.endswith("*"):
# Starts with pattern: text*
stripped = pattern[:-1]
return keyword.startswith(stripped)
else:
# Exact match
return keyword == pattern
class _Filter:
"""
A filter that evaluates logical expressions against a list of keyword strings.
Supports logical operators (AND, OR, NOT), parentheses for grouping, and wildcard patterns
(e.g., `text-*`, `*ai`, `@task:foo`).
Example:
filt = _Filter("(@framework:huggingface OR text-*) AND NOT deprecated")
filt.match(["@framework:huggingface", "text-generation"]) # Returns True
"""
def __init__(self, expression: str) -> None:
"""
Initialize the filter with a string expression.
Args:
expression (str): A logical expression to evaluate against keywords.
Supports AND, OR, NOT, parentheses, and wildcard patterns (*).
"""
self.expression: str = expression
self._ast: Optional[_ExpressionNode] = None
def match(self, keywords: List[str]) -> bool:
"""
Evaluate the filter expression against a list of keywords.
Args:
keywords (List[str]): A list of keyword strings to test.
Returns:
bool: True if the expression evaluates to True for the given keywords, else False.
"""
try:
if self._ast is None:
self._ast = self._parse_expression(self.expression)
return self._ast.evaluate(keywords)
except Exception:
return False
def _parse_expression(self, expr: str) -> _ExpressionNode:
"""
Parse the logical filter expression into an AST.
Args:
expr (str): The raw expression to parse.
Returns:
_ExpressionNode: Root node of the parsed expression AST.
"""
tokens = self._tokenize(expr)
result, _ = self._parse_or_expression(tokens, 0)
return result
def _tokenize(self, expr: str) -> List[str]:
"""Tokenize the expression into logical operators, keywords, and parentheses."""
return re.findall(r"\bAND\b|\bOR\b|\bNOT\b|[^\s()]+|\(|\)", expr, flags=re.IGNORECASE)
def _parse_or_expression(self, tokens: List[str], pos: int) -> tuple[_ExpressionNode, int]:
"""Parse OR expression (lowest precedence)."""
left, pos = self._parse_and_expression(tokens, pos)
while pos < len(tokens) and tokens[pos].upper() == "OR":
pos += 1 # Skip OR token
right, pos = self._parse_and_expression(tokens, pos)
left = _OrNode(left, right)
return left, pos
def _parse_and_expression(self, tokens: List[str], pos: int) -> tuple[_ExpressionNode, int]:
"""Parse AND expression (medium precedence)."""
left, pos = self._parse_not_expression(tokens, pos)
while pos < len(tokens) and tokens[pos].upper() == "AND":
pos += 1 # Skip AND token
right, pos = self._parse_not_expression(tokens, pos)
left = _AndNode(left, right)
return left, pos
def _parse_not_expression(self, tokens: List[str], pos: int) -> tuple[_ExpressionNode, int]:
"""Parse NOT expression (highest precedence)."""
if pos < len(tokens) and tokens[pos].upper() == "NOT":
pos += 1 # Skip NOT token
operand, pos = self._parse_primary_expression(tokens, pos)
return _NotNode(operand), pos
else:
return self._parse_primary_expression(tokens, pos)
def _parse_primary_expression(self, tokens: List[str], pos: int) -> tuple[_ExpressionNode, int]:
"""Parse primary expression (parentheses or pattern)."""
if pos >= len(tokens):
raise ValueError("Unexpected end of expression")
token = tokens[pos]
if token == "(":
pos += 1 # Skip opening parenthesis
expr, pos = self._parse_or_expression(tokens, pos)
if pos >= len(tokens) or tokens[pos] != ")":
raise ValueError("Missing closing parenthesis")
pos += 1 # Skip closing parenthesis
return expr, pos
elif token == ")":
raise ValueError("Unexpected closing parenthesis")
else:
# Pattern token
return _PatternNode(token), pos + 1
def _list_all_hub_models(hub_name: str, sm_client: Session) -> Iterator[HubContent]:
"""
Retrieve all model entries from the specified hub and yield them one by one.
This function paginates through the SageMaker Hub API to retrieve all published models of type "Model"
and yields them as `HubContent` objects.
Args:
hub_name (str): The name of the hub to query.
sm_client (Session): The SageMaker session.
Yields:
HubContent: A `HubContent` object representing a single model entry from the hub.
"""
next_token = None
while True:
# Prepare the request parameters
params = {"HubName": hub_name, "HubContentType": "Model", "MaxResults": 100}
# Add NextToken if it exists
if next_token:
params["NextToken"] = next_token
# Make the API call
response = sm_client.list_hub_contents(**params)
# Yield each content summary
for content in response["HubContentSummaries"]:
yield HubContent(
hub_name=hub_name,
hub_content_arn=content["HubContentArn"],
hub_content_type="Model",
hub_content_name=content["HubContentName"],
hub_content_version=content["HubContentVersion"],
hub_content_description=content.get("HubContentDescription", ""),
hub_content_search_keywords=content.get("HubContentSearchKeywords", []),
)
# Check if there are more results
next_token = response.get("NextToken", None)
if not next_token or len(response["HubContentSummaries"]) == 0:
break # Exit the loop if there are no more pages
[docs]
def search_public_hub_models(
query: str,
hub_name: Optional[str] = "SageMakerPublicHub",
sagemaker_session: Optional[Session] = None,
) -> List[HubContent]:
"""
Search and filter models from hub using a keyword expression.
Args:
query (str): A logical expression used to filter models by keywords.
Example: "@task:text-generation AND NOT @framework:legacy"
hub_name (Optional[str]): The name of the hub to query. Defaults to "SageMakerPublicHub".
sagemaker_session (Optional[Session]): An optional SageMaker `Session` object. If not provided,
a default session will be created and a warning will be logged.
Returns:
List[HubContent]: A list of filtered `HubContent` model objects that match the query.
"""
if sagemaker_session is None:
sagemaker_session = Session()
logger.warning("SageMaker session not provided. Using default Session.")
sm_client = sagemaker_session.sagemaker_client
models = _list_all_hub_models(hub_name, sm_client)
filt = _Filter(query)
results: List[HubContent] = []
for model in models:
keywords = model.hub_content_search_keywords
normalized_keywords = [kw.replace(" ", "-") for kw in keywords]
if filt.match(normalized_keywords):
results.append(model)
return results