Spaces:
Paused
Paused
File size: 6,195 Bytes
9792ea7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | # -*- coding: utf-8 -*-
"""The embedding model card class."""
from __future__ import annotations
import copy
from typing import Literal, Self, Type
import yaml
from pydantic import BaseModel, Field
class EmbeddingModelCard(BaseModel):
"""A card describing an embedding model's capabilities.
Mirrors :class:`~agentscope.model.ModelCard` but tailored for
embedding models. Uses ``input_types`` / ``output_types`` to
describe model capabilities, and ``parameter_schema`` (built from
the embedding class's ``Parameters`` + YAML ``parameter_overrides``)
to tell the frontend which knobs the user can adjust.
The output type ``application/x-embedding`` indicates that the
model produces dense vector embeddings.
"""
type: Literal["embedding_model"] = "embedding_model"
"""The card type, always ``"embedding_model"``."""
name: str = Field(description="The model name used in API calls.")
"""The model name (e.g. ``"text-embedding-3-small"``)."""
label: str = Field(description="Human-readable label for the frontend.")
"""Display label (e.g. ``"Text Embedding 3 Small"``)."""
status: Literal["active", "deprecated", "sunset"] = Field(
default="active",
description="The model lifecycle status.",
)
"""The model status."""
input_types: list[str] = Field(
default=["text/plain"],
description="Supported input media types.",
)
"""Supported input types (e.g. ``["text/plain"]``,
``["text/plain", "image/jpeg", "image/png"]``)."""
output_types: list[str] = Field(
default=["application/x-embedding"],
description="Supported output media types.",
)
"""Output types. ``application/x-embedding`` for vector output."""
dimensions: int = Field(
...,
description="Default output vector dimensions for this model.",
gt=0,
)
"""The default output dimensions for this model.
First-class top-level field — kept outside of
:attr:`parameter_schema` so that callers can rely on a strongly
typed ``int`` rather than the soft ``parameter_schema['properties']
['dimensions']['default']`` lookup.
"""
supported_dimensions: list[int] | None = Field(
default=None,
description=(
"If set, the only dimensions this model can produce. "
"``None`` means dimensions are fixed at "
":attr:`dimensions` and cannot be overridden."
),
)
"""Optional set of allowed output dimensions.
Set for Matryoshka-style models (e.g. OpenAI's
``text-embedding-3-*``) that can be truncated to a smaller size.
``None`` indicates a fixed-dimension model.
"""
context_size: int | None = Field(
default=None,
description="Maximum input length (in tokens) per request.",
gt=0,
)
"""Maximum input context size, if known."""
parameter_schema: dict = Field(
default_factory=dict,
description=(
"JSON Schema for user-configurable parameters "
"(built from the Parameters class + YAML overrides)."
),
)
"""The parameter schema sent to the frontend for form rendering.
Empty ``properties`` means nothing to configure (e.g. fixed
dimensions)."""
parameter_overrides: dict[str, dict] = Field(
default_factory=dict,
description="Raw parameter overrides from the YAML file.",
)
"""The raw parameter overrides, preserved for reference."""
@classmethod
def from_yaml(
cls,
yaml_path: str,
parameter_class: Type[BaseModel],
) -> Self:
"""Load an embedding model card from a YAML file.
Merges the base ``parameter_class`` JSON Schema with
``parameter_overrides`` from the YAML — identical to the
approach used by :meth:`~agentscope.model.ModelCard.from_yaml`.
Args:
yaml_path (`str`):
Path to the YAML file.
parameter_class (`Type[BaseModel]`):
The ``Parameters`` class from the embedding model subclass.
Returns:
`EmbeddingModelCard`: The loaded model card.
"""
with open(yaml_path, "r", encoding="utf-8") as f:
config = yaml.safe_load(f)
if "dimensions" not in config:
raise ValueError(
f"Embedding model card {yaml_path!r} is missing the "
f"required top-level 'dimensions' field.",
)
# Build parameter schema from the Parameters class
base_schema = parameter_class.model_json_schema()
properties = copy.deepcopy(base_schema.get("properties", {}))
# Apply parameter_overrides (same logic as ModelCard.from_yaml)
overrides = config.get("parameter_overrides", {})
for param_name, override in overrides.items():
if override is None:
# null means remove
properties.pop(param_name, None)
continue
if isinstance(override, dict):
if override.get("hidden"):
properties.pop(param_name, None)
continue
# Simple dict merge
if param_name in properties:
properties[param_name] = {
**properties[param_name],
**override,
}
final_schema = {
"type": "object",
"properties": properties,
"required": base_schema.get("required", []),
}
return cls(
name=config["name"],
label=config["label"],
status=config.get("status", "active"),
input_types=config.get("input_types", ["text/plain"]),
output_types=config.get(
"output_types",
["application/x-embedding"],
),
dimensions=config["dimensions"],
supported_dimensions=config.get("supported_dimensions"),
context_size=config.get("context_size"),
parameter_schema=final_schema,
parameter_overrides=overrides,
)
|