feat: expand Model API Pydantic types

Signed-off-by: Miguel Brandão <[email protected]> Signed-off-by: Panos Vagenas <[email protected]> Co-authored-by: Miguel Brandão <[email protected]>
DS4SD · Jul 21, 2023 · ead7246 · ead7246
1 parent b3e7bda
commit ead7246
Show file tree

Hide file tree

Showing 12 changed files with 213 additions and 114 deletions.
diff --git a/deepsearch/model/base/controller.py b/deepsearch/model/base/controller.py
@@ -2,32 +2,42 @@
 from typing import Optional
 
 from deepsearch.model.base.model import BaseDSModel
-from deepsearch.model.base.types import BaseModelConfig
-from deepsearch.model.server.inference_types import ControllerInput, ControllerOutput
+from deepsearch.model.base.types import BaseModelConfig, BaseModelMetadata
+from deepsearch.model.server.inference_types import (
+    AppModelInfoOutput,
+    CtrlPredInput,
+    CtrlPredOutput,
+)
 
 
 class BaseController(ABC):
     _config: Optional[BaseModelConfig] = None
 
-    def get_info(self) -> dict:
-        model = self._get_model()
+    @abstractmethod
+    def get_info(self) -> AppModelInfoOutput:
+        raise NotImplementedError()
+
+    def _get_api_version(self) -> str:
+        return "v1"
+
+    def _get_metadata(self) -> BaseModelMetadata:
         cfg = self._get_config()
-        result = {  # TODO refactor with pydantic
-            "definitions": {
-                "apiVersion": "v1",
-                "kind": cfg.kind,
-                "spec": model.get_definition_spec(),
-            }
-        }
-        return result
+        return BaseModelMetadata(
+            name=cfg.name,
+            version=cfg.version,
+            url=cfg.url,
+            author=cfg.author,
+            description=cfg.description,
+            expected_compute_time=cfg.expected_compute_time,
+        )
 
     def _get_config(self):
         if self._config is None:
             self._config = self._get_model().get_config()
         return self._config
 
     @abstractmethod
-    def dispatch_predict(self, spec: ControllerInput) -> ControllerOutput:
+    def dispatch_predict(self, spec: CtrlPredInput) -> CtrlPredOutput:
         raise NotImplementedError()
 
     @abstractmethod

diff --git a/deepsearch/model/base/model.py b/deepsearch/model/base/model.py
@@ -7,17 +7,3 @@ class BaseDSModel(ABC):
     @abstractmethod
     def get_config(self) -> BaseModelConfig:
         raise NotImplementedError()
-
-    def get_definition_spec(self) -> dict:
-        cfg = self.get_config()
-        spec = {  # TODO refactor with pydantic
-            "metadata": {
-                "name": cfg.name,
-                "version": cfg.version,
-                "url": cfg.url,
-                "author": cfg.author,
-                "description": cfg.description,
-                "expected_compute_time": cfg.expected_compute_time,
-            },
-        }
-        return spec
diff --git a/deepsearch/model/base/types.py b/deepsearch/model/base/types.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 from enum import Enum
-from typing import Any, Optional
+from typing import Any, Dict, Optional
 
 from pydantic import BaseModel, Extra, Field, PositiveFloat
 
@@ -34,18 +34,36 @@ class Metadata(StrictModel):
     annotations: Annotations
 
 
-class BaseInfReq(StrictModel):
+class BaseAppPredInput(StrictModel):
     apiVersion: str
     kind: Kind
     metadata: Metadata
     spec: Any
 
 
-class BaseModelConfig(StrictModel):
-    kind: Kind
+class BaseModelMetadata(StrictModel):
     name: str
     version: str
     url: Optional[str] = None
     author: Optional[str] = None
     description: Optional[str] = None
     expected_compute_time: Optional[PositiveFloat] = None
+
+
+class BaseModelConfig(BaseModelMetadata):
+    kind: Kind
+
+
+class ModelInfoOutputDefsSpec(BaseModel):
+    definition: Dict
+    metadata: BaseModelMetadata
+
+
+class CtrlInfoOutputDefs(BaseModel):
+    apiVersion: str
+    kind: Kind
+    spec: ModelInfoOutputDefsSpec
+
+
+class CtrlInfoOutput(BaseModel):
+    definitions: CtrlInfoOutputDefs
diff --git a/deepsearch/model/examples/dummy_nlp_annotator/model.py b/deepsearch/model/examples/dummy_nlp_annotator/model.py
@@ -5,6 +5,7 @@
 from deepsearch.model.base.types import Kind
 from deepsearch.model.kinds.nlp.model import BaseNLPModel
 from deepsearch.model.kinds.nlp.types import (
+    AnnotateEntitiesEntry,
     AnnotateEntitiesOutput,
     AnnotatePropertiesOutput,
     AnnotateRelationshipsOutput,
@@ -41,18 +42,18 @@ def annotate_batched_entities(
             results.append(
                 {
                     k: [
-                        {
-                            "type": k,
-                            "match": f"a '{k}' match in '{item}'",
-                            "original": f"a '{k}' original in '{item}'",
-                            "range": [1, 5],
-                        },
-                        {
-                            "type": k,
-                            "match": f"another '{k}' match in '{item}'",
-                            "original": f"another '{k}' original in '{item}'",
-                            "range": [12, 42],
-                        },
+                        AnnotateEntitiesEntry(
+                            type=k,
+                            match=f"a '{k}' match in '{item}'",
+                            original=f"a '{k}' original in '{item}'",
+                            range=[1, 5],
+                        ),
+                        AnnotateEntitiesEntry(
+                            type=k,
+                            match=f"another '{k}' match in '{item}'",
+                            original=f"another '{k}' original in '{item}'",
+                            range=[12, 42],
+                        ),
                     ]
                     for k in _entity_names
                 }

diff --git a/deepsearch/model/examples/simple_geo_nlp_annotator/model.py b/deepsearch/model/examples/simple_geo_nlp_annotator/model.py
@@ -157,7 +157,7 @@ def annotate_batched_entities(
 
     def _annotate_entities_in_item(
         self, object_type: str, item: str, entity_names: Optional[List[str]]
-    ) -> List[dict]:
+    ) -> List:
         # In this case entity_names is never None, however since BaseAnnotator defines the signature of this method as
         # Optionally having entity names we must ensure that they are defined.
         if entity_names is None:
@@ -200,7 +200,7 @@ def annotate_batched_relationships(
                 if relation in self.relationship_names:
                     result[relation] = self._rel_annots[
                         relation
-                    ].annotate_relationships_text(text, entity_map)
+                    ].annotate_relationships_text(text, entity_map, relation)
 
             if result:
                 results.append(result)

diff --git a/deepsearch/model/kinds/nlp/controller.py b/deepsearch/model/kinds/nlp/controller.py
@@ -8,27 +8,48 @@
     FindEntitiesText,
     FindPropertiesText,
     FindRelationshipsText,
-    NLPEntitiesControllerOutput,
     NLPEntitiesReqSpec,
-    NLPPropertiesControllerOutput,
+    NLPEntsCtrlPredOuput,
+    NLPInfoOutput,
+    NLPInfoOutputDefinitions,
+    NLPInfoOutputDefinitionsSpec,
+    NLPModelMetadata,
     NLPPropertiesReqSpec,
-    NLPRelationshipsControllerOutput,
+    NLPPropsCtrlPredOutput,
     NLPRelationshipsReqSpec,
+    NLPRelsCtrlPredOutput,
 )
-from deepsearch.model.server.inference_types import ControllerInput, ControllerOutput
+from deepsearch.model.server.inference_types import CtrlPredInput, CtrlPredOutput
 
 
 class NLPController(BaseController):
     def __init__(self, model: BaseNLPModel):
         self._model = model
 
+    def get_info(self) -> NLPInfoOutput:
+        cfg = self._model.get_nlp_config()
+        metadata = NLPModelMetadata(
+            supported_object_types=cfg.supported_types,
+            **self._get_metadata().dict(),  # passing parent metadata dict as kwargs
+        )
+        spec = NLPInfoOutputDefinitionsSpec(
+            definition=cfg.labels,
+            metadata=metadata,
+        )
+        definitions = NLPInfoOutputDefinitions(
+            apiVersion=self._get_api_version(),
+            kind=self.get_kind(),
+            spec=spec,
+        )
+        return NLPInfoOutput(definitions=definitions)
+
     def get_kind(self) -> str:
         return Kind.NLPModel
 
     def _get_model(self) -> BaseDSModel:
         return self._model
 
-    def dispatch_predict(self, spec: ControllerInput) -> ControllerOutput:
+    def dispatch_predict(self, spec: CtrlPredInput) -> CtrlPredOutput:
         cfg = self._model.get_nlp_config()
         type_ok = True
 
@@ -43,7 +64,7 @@ def dispatch_predict(self, spec: ControllerInput) -> ControllerOutput:
                 items=spec.findEntities.texts,
                 entity_names=spec.findEntities.entityNames,
             )
-            return NLPEntitiesControllerOutput(entities=entities)
+            return NLPEntsCtrlPredOuput(entities=entities)
         elif (
             isinstance(spec, NLPRelationshipsReqSpec)
             and isinstance(spec.findRelationships, FindRelationshipsText)
@@ -55,7 +76,7 @@ def dispatch_predict(self, spec: ControllerInput) -> ControllerOutput:
                 entities=spec.findRelationships.entities,
                 relationship_names=spec.findRelationships.relationshipNames,
             )
-            return NLPRelationshipsControllerOutput(relationships=relationships)
+            return NLPRelsCtrlPredOutput(relationships=relationships)
         elif (
             isinstance(spec, NLPPropertiesReqSpec)
             and isinstance(spec.findProperties, FindPropertiesText)
@@ -71,7 +92,7 @@ def dispatch_predict(self, spec: ControllerInput) -> ControllerOutput:
                 entities=entities,
                 property_names=spec.findProperties.propertyNames,
             )
-            return NLPPropertiesControllerOutput(properties=properties)
+            return NLPPropsCtrlPredOutput(properties=properties)
         elif not type_ok:
             raise HTTPException(
                 status_code=status.HTTP_400_BAD_REQUEST,

diff --git a/deepsearch/model/kinds/nlp/model.py b/deepsearch/model/kinds/nlp/model.py
@@ -1,5 +1,4 @@
 from abc import abstractmethod
-from copy import deepcopy
 from typing import List, Optional
 
 from deepsearch.model.base.model import BaseDSModel
@@ -13,8 +12,6 @@
 
 
 class BaseNLPModel(BaseDSModel):
-    _cached_def_spec: dict = {}
-
     @abstractmethod
     def annotate_batched_entities(
         self,
@@ -44,16 +41,6 @@ def annotate_batched_properties(
     ) -> AnnotatePropertiesOutput:
         raise NotImplementedError()
 
-    def get_definition_spec(self) -> dict:
-        cfg = self.get_nlp_config()
-        if not self._cached_def_spec:
-            self._cached_def_spec = deepcopy(super().get_definition_spec())
-            self._cached_def_spec["definition"] = cfg.labels
-            self._cached_def_spec["metadata"][
-                "supported_object_types"
-            ] = cfg.supported_types
-        return self._cached_def_spec
-
     @abstractmethod
     def get_nlp_config(self) -> NLPConfig:
         raise NotImplementedError()

diff --git a/deepsearch/model/kinds/nlp/types.py b/deepsearch/model/kinds/nlp/types.py
@@ -1,7 +1,16 @@
 from enum import Enum
-from typing import List, Literal, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
-from deepsearch.model.base.types import BaseInfReq, BaseModelConfig, Kind, StrictModel
+from deepsearch.model.base.types import (
+    BaseAppPredInput,
+    BaseModelConfig,
+    BaseModelMetadata,
+    CtrlInfoOutput,
+    CtrlInfoOutputDefs,
+    Kind,
+    ModelInfoOutputDefsSpec,
+    StrictModel,
+)
 
 
 class NLPType(str, Enum):
@@ -48,7 +57,7 @@ class NLPRelationshipsReqSpec(StrictModel):
 ]
 
 
-class NLPRequest(BaseInfReq):
+class NLPAppPredInput(BaseAppPredInput):
     kind: Literal[Kind.NLPModel]
     spec: NLPReqSpec
 
@@ -80,33 +89,60 @@ class AnnotationLabels(StrictModel):
     properties: List[PropertyLabel]
 
 
-# TODO Annotate*Input pydantic models needed?
+class AnnotateEntitiesEntry(StrictModel):
+    type: str
+    match: str
+    original: str
+    range: List[int]
 
-AnnotateEntitiesOutput = List[dict]  # TODO provide real implementation
-AnnotateRelationshipsOutput = List[dict]  # TODO provide real implementation
-AnnotatePropertiesOutput = List[dict]  # TODO provide real implementation
 
+class AnnotateRelationshipsEntry(StrictModel):
+    header: list
+    data: list
 
-class NLPEntitiesControllerOutput(StrictModel):
+
+AnnotateEntitiesOutput = List[Dict[str, List[AnnotateEntitiesEntry]]]
+AnnotateRelationshipsOutput = List[Dict[str, AnnotateRelationshipsEntry]]
+AnnotatePropertiesOutput = List[Dict]  # TODO specify
+
+
+class NLPEntsCtrlPredOuput(StrictModel):
     entities: AnnotateEntitiesOutput
 
 
-class NLPRelationshipsControllerOutput(StrictModel):
+class NLPRelsCtrlPredOutput(StrictModel):
     relationships: AnnotateRelationshipsOutput
 
 
-class NLPPropertiesControllerOutput(StrictModel):
+class NLPPropsCtrlPredOutput(StrictModel):
     properties: AnnotatePropertiesOutput
 
 
-NLPControllerOutput = Union[
-    NLPEntitiesControllerOutput,
-    NLPRelationshipsControllerOutput,
-    NLPPropertiesControllerOutput,
+NLPCtrlPredOutput = Union[
+    NLPEntsCtrlPredOuput,
+    NLPRelsCtrlPredOutput,
+    NLPPropsCtrlPredOutput,
 ]
 
 
 class NLPConfig(BaseModelConfig):
     kind: Literal[Kind.NLPModel]
     supported_types: List[NLPType]
     labels: AnnotationLabels
+
+
+class NLPModelMetadata(BaseModelMetadata):
+    supported_object_types: List[Literal["text", "table", "image"]]
+
+
+class NLPInfoOutputDefinitionsSpec(ModelInfoOutputDefsSpec):
+    metadata: NLPModelMetadata
+
+
+class NLPInfoOutputDefinitions(CtrlInfoOutputDefs):
+    kind: Literal[Kind.NLPModel]
+    spec: NLPInfoOutputDefinitionsSpec
+
+
+class NLPInfoOutput(CtrlInfoOutput):
+    definitions: NLPInfoOutputDefinitions