52 lines
1.4 KiB
Python
52 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from common import ROOT, iter_model_card_paths, read_yaml_mapping
|
|
|
|
|
|
INDEX_PATH = ROOT / "registry" / "index.json"
|
|
|
|
|
|
def simplify_card(path: Path, data: dict) -> dict:
|
|
deployment = data.get("deployment") or {}
|
|
return {
|
|
"id": data.get("id"),
|
|
"name": data.get("name"),
|
|
"type": data.get("type"),
|
|
"status": data.get("status"),
|
|
"task": data.get("task") or [],
|
|
"language": data.get("language") or [],
|
|
"source": data.get("source"),
|
|
"upstream_id": data.get("upstream_id"),
|
|
"license": data.get("license"),
|
|
"storage_path": data.get("storage_path"),
|
|
"format": data.get("format"),
|
|
"quantization": data.get("quantization"),
|
|
"runtime": deployment.get("runtime"),
|
|
"served_model_name": deployment.get("served_model_name"),
|
|
"card_path": str(path.relative_to(ROOT)).replace("\\", "/"),
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
models = []
|
|
for path in iter_model_card_paths():
|
|
models.append(simplify_card(path, read_yaml_mapping(path)))
|
|
|
|
index = {
|
|
"schema_version": 1,
|
|
"models": models,
|
|
}
|
|
|
|
INDEX_PATH.write_text(
|
|
json.dumps(index, ensure_ascii=False, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(f"Wrote {INDEX_PATH.relative_to(ROOT)} with {len(models)} model(s).")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|