fix(api): levelup save and display to top-level (#176)

* fix(api): levelup save and display to top-level

* fix(api): levelup save and display to top-level

README.md

@@ -82,7 +82,7 @@ Perfect! Now `embed_model` and `labeled_data` are given by you already, simply d
 ```python
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     embed_model,
     train_data=labeled_data
 )
@@ -96,7 +96,7 @@ You can use Finetuner to interactive label data and train `embed_model` as below
 ```python
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     embed_model,
     train_data=unlabeled_data,
     interactive=True
@@ -111,7 +111,7 @@ worries, Finetuner can convert your model into an embedding model and train it v
 ```python
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     general_model,
     train_data=labeled_data,
     to_embedding_model=True,
@@ -127,7 +127,7 @@ worries, Finetuner can help you train an embedding model with interactive labeli
 ```python
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     general_model,
     train_data=unlabeled_data,
     interactive=True,

docs/basics/fit.md

@@ -3,4 +3,85 @@
 ```{include} ../index.md
 :start-after: <!-- start fit-method -->
 :end-before: <!-- end fit-method -->
+```
+
+Note, `finetuner.fit` always returns two objects in a tuple: the tuned model and a summary statistic of the fitting procedure. You can save the model via `finetuner.save(model)`. You can save the summary object by calling `.save()`; or plot it via `.plot()` and see how the loss on training & evaluation data changes over time.
+
+## Save model
+
+```python
+import finetuner
+
+finetuner.save(model, './saved-model')
+```
+
+```{caution}
+Depending on your framework, `save` can behave differently. Tensorflow/Keras saves model architecture with the parameters, whereas PyTorch & Paddle only saves the trained parameters.
+```
+
+## Display model
+
+```python
+import finetuner
+
+finetuner.display(model)
+```
+
+```{caution}
+Depending on your framework, `display` may require different argument for rendering the model correctly. In PyTorch & Paddle, you will also need to give the `input_size` and sometimes `input_dtype` to correctly render the model.
+```
+
+More information can be {ref}`found here<display-method>`.
+
+## Example
+
+```python
+import torch
+from finetuner.toydata import generate_qa_match
+import finetuner
+
+class LastCell(torch.nn.Module):
+    def forward(self, x):
+        out, _ = x
+        return out[:, -1, :]
+
+embed_model = torch.nn.Sequential(
+    torch.nn.Embedding(num_embeddings=5000, embedding_dim=64),
+    torch.nn.LSTM(64, 64, bidirectional=True, batch_first=True),
+    LastCell(),
+    torch.nn.Linear(in_features=2 * 64, out_features=32))
+
+model, summary = finetuner.fit(
+    embed_model,
+    train_data=lambda: generate_qa_match(num_neg=5, max_seq_len=10),
+    eval_data=lambda: generate_qa_match(num_neg=5, max_seq_len=10),
+    epochs=3,
+)
+
+finetuner.display(model, input_size=(100,), input_dtype='long')
+
+finetuner.save(model, './saved-model')
+summary.plot('fit.png')
+```
+
+```console
+⠼  Epoch 1/3 ━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0:00:00 estimating... 
+⠼       DONE ━━━━━━━━━━━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0:00:01 10.2 step/s train: 1.25 | Eval Loss: 1.06
+⠇       DONE ━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0:00:01 100% ETA: 0 seconds train: 1.14 | Eval Loss: 1.02
+⠹       DONE ━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 0:00:01 100% ETA: 0 seconds train: 1.10 | Eval Loss: 0.98
+                                                                    
+  name          output_shape_display         nb_params   trainable  
+ ────────────────────────────────────────────────────────────────── 
+  embedding_1   [100, 64]                    320000      True       
+  lstm_2        [[[2, 2, 64], [2, 2, 64]]]   66560       False      
+  lastcell_3    [128]                        0           False      
+  linear_4      [32]                         4128        True       
+                                                                    
+Green layers can be used as embedding layers, whose name can be used as 
+layer_name in to_embedding_model(...).
+```
+
+```{figure} fit-plot.png
+:align: center
+:width: 80%
 ```

docs/components/tailor.md

@@ -37,9 +37,10 @@ Here, `model` is the general model with loaded weights; `layer_name` is the sele
 
 In general, you do not need to call `to_embedding_model` manually. You can use it directly via `finetuner.fit(..., to_embedding_model=True)`
 
+(display-method)=
 ## `display` method
 
-Tailor also provides a helper function `finetuner.tailor.display()` that gives a table summary of a Keras/PyTorch/Paddle model.
+Tailor also provides a helper function `finetuner.display()` that gives a table summary of a Keras/PyTorch/Paddle model.
 
 Let's see how to use them in action.
 

docs/components/tuner.md

@@ -65,7 +65,7 @@ Although siamese and triplet loss works on pair and triplet inputs respectively,
 
 ## `save` method
 
-After a model is tuned, you can save it by calling `finetuner.tuner.save(model, save_path)`.
+After a model is tuned, you can save it by calling `finetuner.save(model, save_path)`.
 
 
 ## Examples

docs/index.md

@@ -136,7 +136,7 @@ Perfect! Now `embed_model` and `train_data` are already provided by you, simply
 ```python
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     embed_model,
     train_data=train_data
 )
@@ -159,7 +159,7 @@ emphasize-lines: 6
 ---
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     embed_model,
     train_data=unlabeled_data,
     interactive=True
@@ -183,7 +183,7 @@ emphasize-lines: 6, 7
 ---
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     general_model,
     train_data=labeled_data,
     to_embedding_model=True,
@@ -208,7 +208,7 @@ emphasize-lines: 6, 7
 ---
 import finetuner
 
-finetuner.fit(
+tuned_model, _ = finetuner.fit(
     general_model,
     train_data=labeled_data,
     interactive=True,

finetuner/__init__.py

@@ -108,8 +108,13 @@ def fit(
 
         # TODO: atm return will never hit as labeler UI hangs the
         #  flow via `.block()`
-        fit(model, train_data, *args, **kwargs)
+        return model, fit(model, train_data, *args, **kwargs)
     else:
         from .tuner import fit
 
         return model, fit(model, train_data, *args, **kwargs)
+
+
+# level them up to the top-level
+from .tuner import save
+from .tailor import display