feat(fit): add tailor to top-level fit function

finetuner/__init__.py

@@ -6,59 +6,91 @@
 __default_tag_key__ = 'finetuner'
 
 # define the high-level API: fit()
-from typing import Optional, overload, TYPE_CHECKING
-
+from typing import Optional, overload, TYPE_CHECKING, Tuple
 
 if TYPE_CHECKING:
  from .helper import AnyDNN, DocumentArrayLike, TunerReturnType
 
-# fit interface generated from Labeler + Tuner
-# overload_inject_fit_tailor_tuner_start
 
-# overload_inject_fit_tailor_tuner_end
+# fit interface generated from Tuner
+@overload
+def fit(
+ model: 'AnyDNN', #: must be an embedding model
+ train_data: 'DocumentArrayLike',
+ eval_data: Optional['DocumentArrayLike'] = None,
+ epochs: int = 10,
+ batch_size: int = 256,
+ head_layer: str = 'CosineLayer',
+) -> 'TunerReturnType':
+ ...
+
+
+# fit interface derived from Tailor + Tuner
+@overload
+def fit(
+ model: 'AnyDNN',
+ train_data: 'DocumentArrayLike',
+ eval_data: Optional['DocumentArrayLike'] = None,
+ epochs: int = 10,
+ batch_size: int = 256,
+ head_layer: str = 'CosineLayer',
+ to_embedding_model: bool = True, #: below are tailor args
+ input_size: Optional[Tuple[int, ...]] = None,
+ input_dtype: str = 'float32',
+ layer_name: Optional[str] = None,
+ output_dim: Optional[int] = None,
+ freeze: bool = False,
+) -> 'TunerReturnType':
+ ...
 
 
-# fit interface generated from Labeler + Tuner
-# overload_inject_fit_labeler_tuner_start
+# fit interface from Labeler + Tuner
 @overload
 def fit(
- embed_model: 'AnyDNN',
+ model: 'AnyDNN', #: must be an embedding model
  train_data: 'DocumentArrayLike',
+ interactive: bool = True, #: below are labeler args
  clear_labels_on_start: bool = False,
  port_expose: Optional[int] = None,
  runtime_backend: str = 'thread',
- interactive: bool = True,
  head_layer: str = 'CosineLayer',
 ) -> None:
  ...
 
 
-# overload_inject_fit_labeler_tuner_end
-
-
-# fit interface generated from Tuner
-# overload_inject_fit_tuner_start
+# fit interface from Labeler + Tailor + Tuner
 @overload
 def fit(
- embed_model: 'AnyDNN',
+ model: 'AnyDNN',
  train_data: 'DocumentArrayLike',
- eval_data: Optional['DocumentArrayLike'] = None,
- epochs: int = 10,
- batch_size: int = 256,
+ interactive: bool = True, #: below are labeler args
+ clear_labels_on_start: bool = False,
+ port_expose: Optional[int] = None,
+ runtime_backend: str = 'thread',
  head_layer: str = 'CosineLayer',
-) -> 'TunerReturnType':
+ to_embedding_model: bool = True, #: below are tailor args
+ input_size: Optional[Tuple[int, ...]] = None,
+ input_dtype: str = 'float32',
+ layer_name: Optional[str] = None,
+ output_dim: Optional[int] = None,
+ freeze: bool = False,
+) -> None:
  ...
 
 
-# overload_inject_fit_tuner_end
+def fit(
+ model: 'AnyDNN', train_data: 'DocumentArrayLike', *args, **kwargs
+) -> Optional['TunerReturnType']:
+ if kwargs.get('to_embedding_model', False):
+ from .tailor import to_embedding_model
 
+ model = to_embedding_model(model, *args, **kwargs)
 
-def fit(*args, **kwargs) -> Optional['TunerReturnType']:
  if kwargs.get('interactive', False):
  from .labeler import fit
 
- return fit(*args, **kwargs)
+ return fit(model, train_data, *args, **kwargs)
  else:
  from .tuner import fit
 
- return fit(*args, **kwargs)
+ return fit(model, train_data, *args, **kwargs)

finetuner/helper.py

@@ -52,6 +52,6 @@ def get_framework(dnn_model: AnyDNN) -> str:
  )
 
 
-def is_list_int(tp) -> bool:
- """Return True if the input is a list of integers."""
+def is_seq_int(tp) -> bool:
+ """Return True if the input is a sequence of integers."""
  return tp and isinstance(tp, Sequence) and all(isinstance(p, int) for p in tp)

finetuner/tailor/__init__.py

@@ -10,6 +10,7 @@ def to_embedding_model(
  layer_name: Optional[str] = None,
  output_dim: Optional[int] = None,
  freeze: bool = False,
+ **kwargs
 ) -> AnyDNN:
  f_type = get_framework(model)
 

finetuner/tailor/paddle/__init__.py

@@ -9,7 +9,7 @@
 from paddle import nn, Tensor
 
 from ..base import BaseTailor
-from ...helper import is_list_int, EmbeddingLayerInfoType, AnyDNN
+from ...helper import is_seq_int, EmbeddingLayerInfoType, AnyDNN
 
 
 class PaddleTailor(BaseTailor):
@@ -105,7 +105,7 @@ def hook(layer, input, output):
  if (
  not output_shape
  or len(output_shape) != 2
- or not is_list_int(output_shape)
+ or not is_seq_int(output_shape)
  or summary[layer]['cls_name'] in self._model.__class__.__name__
  ):
  continue

finetuner/tailor/pytorch/__init__.py

@@ -9,7 +9,7 @@
 from torch import nn
 
 from ..base import BaseTailor
-from ...helper import is_list_int, EmbeddingLayerInfoType, AnyDNN
+from ...helper import is_seq_int, EmbeddingLayerInfoType, AnyDNN
 
 
 class PytorchTailor(BaseTailor):
@@ -92,7 +92,7 @@ def hook(module, input, output):
  if (
  not output_shape
  or len(output_shape) != 2
- or not is_list_int(output_shape)
+ or not is_seq_int(output_shape)
  or summary[layer]['cls_name'] in self._model.__class__.__name__
  ):
  continue