refactor(tailor): fix type hint in tailor (#88)

jina-ai · Oct 6, 2021 · 1956a3d · 1956a3d
1 parent 68fc783
commit 1956a3d
Show file tree

Hide file tree

Showing 5 changed files with 16 additions and 17 deletions.
diff --git a/finetuner/__init__.py b/finetuner/__init__.py
@@ -9,8 +9,9 @@
 from typing import Optional, overload, TYPE_CHECKING
 
 
-from .helper import AnyDNN, DocumentArrayLike
-from .tuner.fit import TunerReturnType
+if TYPE_CHECKING:
+    from .helper import AnyDNN, DocumentArrayLike
+    from .tuner.fit import TunerReturnType
 
 # fit interface generated from Labeler + Tuner
 # overload_inject_fit_tailor_tuner_start
@@ -22,8 +23,8 @@
 # overload_inject_fit_labeler_tuner_start
 @overload
 def fit(
-    embed_model: AnyDNN,
-    train_data: DocumentArrayLike,
+    embed_model: 'AnyDNN',
+    train_data: 'DocumentArrayLike',
     clear_labels_on_start: bool = False,
     port_expose: Optional[int] = None,
     runtime_backend: str = 'thread',
@@ -40,20 +41,20 @@ def fit(
 # overload_inject_fit_tuner_start
 @overload
 def fit(
-    embed_model: AnyDNN,
-    train_data: DocumentArrayLike,
-    eval_data: Optional[DocumentArrayLike] = None,
+    embed_model: 'AnyDNN',
+    train_data: 'DocumentArrayLike',
+    eval_data: Optional['DocumentArrayLike'] = None,
     epochs: int = 10,
     batch_size: int = 256,
     head_layer: str = 'CosineLayer',
-) -> TunerReturnType:
+) -> 'TunerReturnType':
     ...
 
 
 # overload_inject_fit_tuner_end
 
 
-def fit(*args, **kwargs) -> Optional[TunerReturnType]:
+def fit(*args, **kwargs) -> Optional['TunerReturnType']:
     if kwargs.get('interactive', False):
         kwargs.pop('interactive')
         from .labeler.fit import fit

diff --git a/finetuner/tailor/base.py b/finetuner/tailor/base.py
@@ -3,8 +3,6 @@
     Optional,
 )
 
-from jina.logging.logger import JinaLogger
-
 from ..helper import AnyDNN, EmbeddingLayerInfo
 
 

diff --git a/finetuner/tailor/keras/__init__.py b/finetuner/tailor/keras/__init__.py
@@ -1,7 +1,7 @@
 from tensorflow.keras import Model
 
 from ..base import BaseTailor
-from ...helper import AnyDNN, EmbeddingLayerInfo
+from ...helper import EmbeddingLayerInfo
 
 
 class KerasTailor(BaseTailor):

diff --git a/finetuner/tailor/paddle/__init__.py b/finetuner/tailor/paddle/__init__.py
@@ -1,6 +1,6 @@
-from typing import Tuple
-from copy import deepcopy
 from collections import OrderedDict
+from copy import deepcopy
+from typing import Tuple
 
 import numpy as np
 import paddle

diff --git a/finetuner/tailor/pytorch/__init__.py b/finetuner/tailor/pytorch/__init__.py
@@ -1,13 +1,13 @@
-from typing import Tuple
-from copy import deepcopy
 from collections import OrderedDict
+from copy import deepcopy
+from typing import Tuple
 
 import numpy as np
 import torch
 from torch import nn
 
 from ..base import BaseTailor
-from ...helper import AnyDNN, is_list_int, EmbeddingLayerInfo
+from ...helper import is_list_int, EmbeddingLayerInfo
 
 
 class PytorchTailor(BaseTailor):