Add typing to the model wrapper

vitalii-bulyzhyn · vitalii-bulyzhyn · commit 01de0ffe61f8 · 2022-10-04T15:53:40.000-03:00
diff --git a/sdk/diffgram.egg-info/SOURCES.txt b/sdk/diffgram.egg-info/SOURCES.txt
@@ -36,6 +36,7 @@ diffgram/label/test_label.py
 diffgram/member/__init__.py
 diffgram/models/__init__.py
 diffgram/models/base_model.py
+diffgram/models/model_interfaces.py
 diffgram/pytorch_diffgram/__init__.py
 diffgram/pytorch_diffgram/diffgram_pytorch_dataset.py
 diffgram/regular/__init__.py
diff --git a/sdk/diffgram/__init__.py b/sdk/diffgram/__init__.py
@@ -5,4 +5,4 @@
 from diffgram.core.core import Project
 from diffgram.file.file import File
 from diffgram.task.task import Task
-from diffgram.models.base_model import DiffgramBaseModel
+from diffgram.models.base_model import DiffgramBaseModel, Instance
diff --git a/sdk/diffgram/models/base_model.py b/sdk/diffgram/models/base_model.py
@@ -1,10 +1,10 @@
-from fastapi import FastAPI
-
+from .model_interfaces import DiffgramFile, Instance
+from typing import List
 class DiffgramBaseModel():
     def init(self):
         pass
 
-    def infere(self):
+    def infere(self, file: DiffgramFile) -> List[Instance]:
         raise NotImplementedError
 
     def get_schema(self):
@@ -14,10 +14,21 @@ def ping(self):
         pass
 
     def serve(self, app):
-        @app.get("/infere")
-        async def predict():
+        @app.post("/infere")
+        async def predict(file: DiffgramFile):
+            predictions = self.infere(file)
+
+            if not isinstance(predictions, List):
+                raise ValueError('infere should return List of type Instance') 
+
+            for prediction in predictions:
+                res = isinstance(prediction, Instance)
+                if not res:
+                    raise ValueError('infere should return List of type Instance') 
+
             return {
-                "message": "Infere route"
+                "file": file,
+                "predictions": predictions
             }
 
         @app.get("/get_schema")
diff --git a/sdk/diffgram/models/model_interfaces.py b/sdk/diffgram/models/model_interfaces.py
@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+from typing import List
+
+class DiffgramFile(BaseModel):
+    id: int
+    type: str
+
+class Instance(BaseModel):
+    id: int
diff --git a/sdk/samples/model_server.py b/sdk/samples/model_server.py
@@ -1,6 +1,14 @@
-from diffgram import DiffgramBaseModel
+from diffgram import DiffgramBaseModel, Instance
 from fastapi import FastAPI
 
-app = FastAPI()
+class MyTestModel(DiffgramBaseModel):
+    def infere(self, file):
+        return [
+            Instance(id=1)
+        ]
+
+    def get_schema(self):
+        return super().get_schema()
 
-DiffgramBaseModel().serve(app)
+app = FastAPI()
+MyTestModel().serve(app)
