WIP refactor benchmarking: extract Engine

Author: maximpavliv

dlclive/benchmark.py

@@ -25,7 +25,7 @@
 from dlclive import DLCLive
 from dlclive import VERSION
 from dlclive import __file__ as dlcfile
-from dlclive.factory import Engine
+from dlclive.engine import Engine
 from dlclive.utils import decode_fourcc
 
 

dlclive/engine.py

@@ -0,0 +1,14 @@
+from enum import Enum
+
+class Engine(Enum):
+    TENSORFLOW = "tensorflow"
+    PYTORCH = "pytorch"
+
+    @classmethod
+    def from_model_type(cls, model_type: str) -> "Engine":
+        if model_type.lower() == "pytorch":
+            return cls.PYTORCH
+        elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"):
+            return cls.TENSORFLOW
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")

dlclive/factory.py

@@ -5,6 +5,7 @@
 from typing import Literal
 
 from dlclive.core.runner import BaseRunner
+from dlclive.engine import Engine
 
 
 def build_runner(
@@ -54,19 +55,3 @@ def build_runner(
 def filter_keys(valid: set[str], kwargs: dict) -> dict:
     """Filters the keys in kwargs, only keeping those in valid."""
     return {k: v for k, v in kwargs.items() if k in valid}
-
-
-from enum import Enum
-
-class Engine(Enum):
-    TENSORFLOW = "tensorflow"
-    PYTORCH = "pytorch"
-
-    @classmethod
-    def from_model_type(cls, model_type: str) -> "Engine":
-        if model_type.lower() == "pytorch":
-            return cls.PYTORCH
-        elif model_type.lower() in ("tensorflow", "base", "tensorrt", "lite"):
-            return cls.TENSORFLOW
-        else:
-            raise ValueError(f"Unknown model type: {model_type}")