85 lines
2.1 KiB
Python
85 lines
2.1 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from common import iter_model_card_paths, read_yaml_mapping
|
|
|
|
REQUIRED_FIELDS = {
|
|
"id",
|
|
"name",
|
|
"type",
|
|
"task",
|
|
"language",
|
|
"source",
|
|
"license",
|
|
"status",
|
|
"storage_path",
|
|
"deployment",
|
|
}
|
|
|
|
VALID_STATUSES = {"candidate", "draft", "staging", "production", "archived"}
|
|
|
|
|
|
def load_yaml(path: Path) -> dict:
|
|
return read_yaml_mapping(path)
|
|
|
|
|
|
def validate_card(path: Path) -> list[str]:
|
|
errors: list[str] = []
|
|
|
|
try:
|
|
data = load_yaml(path)
|
|
except Exception as exc:
|
|
return [f"{path}: cannot parse YAML: {exc}"]
|
|
|
|
missing = sorted(REQUIRED_FIELDS - set(data))
|
|
for field in missing:
|
|
errors.append(f"{path}: missing required field `{field}`")
|
|
|
|
status = data.get("status")
|
|
if status is not None and status not in VALID_STATUSES:
|
|
errors.append(
|
|
f"{path}: invalid status `{status}`, expected one of {sorted(VALID_STATUSES)}"
|
|
)
|
|
|
|
for field in ("task", "language"):
|
|
value = data.get(field)
|
|
if value is not None and not isinstance(value, list):
|
|
errors.append(f"{path}: `{field}` must be a list")
|
|
|
|
deployment = data.get("deployment")
|
|
if deployment is not None and not isinstance(deployment, dict):
|
|
errors.append(f"{path}: `deployment` must be a mapping")
|
|
|
|
return errors
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Validate model cards.")
|
|
parser.add_argument("--include-examples", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
cards = iter_model_card_paths(include_examples=args.include_examples)
|
|
if not cards:
|
|
print("No model cards found.")
|
|
return 0
|
|
|
|
errors: list[str] = []
|
|
for card in cards:
|
|
errors.extend(validate_card(card))
|
|
|
|
if errors:
|
|
print("Model card validation failed:")
|
|
for error in errors:
|
|
print(f"- {error}")
|
|
return 1
|
|
|
|
print(f"Validated {len(cards)} model card(s).")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|