Add support for model statistics

This commit is contained in:
allegroai
2022-07-08 17:39:41 +03:00
parent 0c15169668
commit adc1825843
5 changed files with 109 additions and 32 deletions

View File

@@ -1,5 +1,5 @@
from datetime import datetime
from typing import Callable, Tuple
from typing import Callable, Tuple, Sequence, Dict
from apiserver.apierrors import errors
from apiserver.apimodels.models import ModelTaskPublishResponse
@@ -128,3 +128,33 @@ class ModelBLL:
)
return unarchived
@classmethod
def get_model_stats(
cls, company: str, model_ids: Sequence[str],
) -> Dict[str, dict]:
if not model_ids:
return {}
result = Model.aggregate(
[
{
"$match": {
"company": {"$in": [None, "", company]},
"_id": {"$in": model_ids},
}
},
{
"$addFields": {
"labels_count": {"$size": {"$objectToArray": "$labels"}}
}
},
{
"$project": {"labels_count": 1},
},
]
)
return {
r.pop("_id"): r
for r in result
}