Repository: voidful/TFkit (language: python)

Issue:
Refactor
refactor dataset to allow large data loading


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 7f4b33c..3cf0f2d 100644
--- a/README.md
+++ b/README.md
@@ -29,13 +29,15 @@
 </p>
 
 ## What is it
-TFKit is a tool kit mainly for language generation.  
-It leverages the use of transformers on many tasks with different models in this all-in-one framework.   
-All you need is a little change of config.  
+
+TFKit is a tool kit mainly for language generation.
+It leverages the use of transformers on many tasks with different models in this all-in-one framework.
+All you need is a little change of config.
 
 ## Task Supported
-With transformer models - BERT/ALBERT/T5/BART......  
-|  |  |
+
+With transformer models - BERT/ALBERT/T5/BART......
+| | |
 |-|-|
 | Text Generation | :memo: seq2seq language model |
 | Text Generation | :pen: causal language model |
@@ -43,23 +45,29 @@ With transformer models - BERT/ALBERT/T5/BART......
 | Text Generation | :pencil: onebyone generation model |
 
 # Getting Started
-Learn more from the [document](https://voidful.github.io/TFkit/).  
+
+Learn more from the [document](https://voidful.github.io/TFkit/).
 
 ## How To Use
 
 ### Step 0: Install
+
 Simple installation from PyPI
+
 ```bash
 pip install git+https://github.com/voidful/TFkit.git@refactor-dataset
 ```
 
 ### Step 1: Prepare dataset in csv format
+
 [Task format](https://voidful.tech/TFkit/tasks/)
-``` 
+
+```
 input, target
 ```
 
 ### Step 2: Train model
+
 ```bash
 tfkit-train \
 --task clas \
@@ -73,6 +81,7 @@ tfkit-train \
 ```
 
 ### Step 3: Evaluate
+
 ```bash
 tfkit-eval \
 --task roberta_sentiment_classificer/1.pt \
@@ -81,45 +90,50 @@ tfkit-eval \
 ```
 
 ## Advanced features
+
 <details>
   <summary>Multi-task training </summary>
 
-  ```bash
-  tfkit-train \
-    --task clas clas \
-    --config xlm-roberta-base \
-    --train training_data_taskA.csv training_data_taskB.csv \
-    --test testing_data_taskA.csv testing_data_taskB.csv \
-    --lr 4e-5 \
-    --maxlen 384 \
-    --epoch 10 \
-    --savedir roberta_sentiment_classificer_multi_task
-  ```
+```bash
+tfkit-train \
+  --task clas clas \
+  --config xlm-roberta-base \
+  --train training_data_taskA.csv training_data_taskB.csv \
+  --test testing_data_taskA.csv testing_data_taskB.csv \
+  --lr 4e-5 \
+  --maxlen 384 \
+  --epoch 10 \
+  --savedir roberta_sentiment_classificer_multi_task
+```
+
 </details>
 
 ## Not maintained task
+
 Due to time constraints, the following tasks are temporarily not supported
-|  |  |
+| | |
 |-|-|
-| Classification  | :label: multi-class and multi-label classification |
-| Question Answering  | :page_with_curl: extractive qa |
-| Question Answering  | :radio_button: multiple-choice qa |
-| Tagging  | :eye_speech_bubble: sequence level tagging / sequence level with crf  |
+| Classification | :label: multi-class and multi-label classification |
+| Question Answering | :page_with_curl: extractive qa |
+| Question Answering | :radio_button: multiple-choice qa |
+| Tagging | :eye_speech_bubble: sequence level tagging / sequence level with crf |
 | Self-supervise Learning | :diving_mask: mask language model |
 
 ## Supplement
-- [transformers models list](https://huggingface.co/models): you can find any pretrained models here   
-- [nlprep](https://github.com/voidful/NLPrep): download and preprocessing data in one line     
-- [nlp2go](https://github.com/voidful/nlp2go): create demo api as quickly as possible.
 
+- [transformers models list](https://huggingface.co/models): you can find any pretrained models here
+- [nlprep](https://github.com/voidful/NLPrep): download and preprocessing data in one line
+- [nlp2go](https://github.com/voidful/nlp2go): create demo api as quickly as possible.
 
 ## Contributing
+
 Thanks for your interest.There are many ways to contribute to this project. Get started [here](https://github.com/voidful/tfkit/blob/master/CONTRIBUTING.md).
 
 ## License ![PyPI - License](https://img.shields.io/github/license/voidful/tfkit)
 
-* [License](https://github.com/voidful/tfkit/blob/master/LICENSE)
+- [License](https://github.com/voidful/tfkit/blob/master/LICENSE)
 
 ## Icons reference
-Icons modify from <a href="http://www.freepik.com/" title="Freepik">Freepik</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>      
-Icons modify from <a href="https://www.flaticon.com/authors/nikita-golubev" title="Nikita Golubev">Nikita Golubev</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>      
+
+Icons modify from <a href="http://www.freepik.com/" title="Freepik">Freepik</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>
+Icons modify from <a href="https://www.flaticon.com/authors/nikita-golubev" title="Nikita Golubev">Nikita Golubev</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>
diff --git a/docs/index.md b/docs/index.md
index 85d4234..d9a3d1d 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -22,37 +22,40 @@
 ## Getting started
 
 ### Installing via pip
+
 ```bash
 pip install tfkit
 ```
 
-* You can use tfkit for model training and evaluation with `tfkit-train` and `tfkit-eval`.
+- You can use tfkit for model training and evaluation with `tfkit-train` and `tfkit-eval`.
 
 ### Running TFKit on the task you wanted
 
 ### First step - prepare your dataset
+
 The key to combine different task together is to make different task with same data format.
 
-**notice**  
+**notice**
 
-* All data will be in csv format - tfkit will use **csv** for all task, normally it will have two columns, first columns is the input of models, the second column is the output of models.
-* Plane text with no tokenization - there is no need to tokenize text before training, or do re-calculating for tokenization, tfkit will handle it for you.
-* No header is needed.
+- All data will be in csv format - tfkit will use **csv** for all task, normally it will have two columns, first columns is the input of models, the second column is the output of models.
+- Plane text with no tokenization - there is no need to tokenize text before training, or do re-calculating for tokenization, tfkit will handle it for you.
+- No header is needed.
 
 For example, a sentiment classification dataset will be like:
+
 ```csv
 how dare you,negative
 ```
 
-!!! hint 
-    For the detail and example format on different, you can check [here](tasks/) 
+!!! hint
+For the detail and example format on different, you can check [here](tasks/)
 
-!!! hint 
-    nlprep is a tool for data split/preprocessing/argumentation, it can help you to create ready to train data for tfkit, check [here](https://github.com/voidful/NLPrep)
+!!! hint
+nlprep is a tool for data split/preprocessing/argumentation, it can help you to create ready to train data for tfkit, check [here](https://github.com/voidful/NLPrep)
 
 ### Second step - model training
 
-Using `tfkit-train` for model training, you can use 
+Using `tfkit-train` for model training, you can use
 
 Before training a model, there is something you need to clarify:
 
@@ -60,10 +63,11 @@ Before training a model, there is something you need to clarify:
 - `--config` what pretrained model you want to use？ you can go [https://huggingface.co/models](https://huggingface.co/models) to search for available pretrained models.
 - `--train` and `--test` training and testing dataset path, which is in csv format.
 - `--savedir` model saving directory, default will be in '/checkpoints' folder
-  
+
 you can leave the rest to the default config, or use `tfkit-train -h` to more configuration.
 
 An example about training a sentiment classifier:
+
 ```bash
 tfkit-train \
 --task clas \
@@ -78,55 +82,64 @@ tfkit-train \
 
 #### Third step - model eval
 
-Using `tfkit-eval` for model evaluation.   
-- `--model` saved model's path.  
-- `--metric` the evaluation metric eg: emf1, nlg(BLEU/ROUGE), clas(confusion matrix).  
-- `--valid` validation data, also in csv format.  
-- `--panel` a input panel for model specific parameter.  
+Using `tfkit-eval` for model evaluation.
+
+- `--model` saved model's path.
+- `--metric` the evaluation metric eg: emf1, nlg(BLEU/ROUGE), clas(confusion matrix).
+- `--valid` validation data, also in csv format.
+- `--panel` a input panel for model specific parameter.
 
 for more configuration detail, you may use `tfkit-eval -h`.
 
-After evaluate, It will print evaluate result in your console, and also generate three report for debugging.  
-- `*_score.csv` overall score, it is the copy of the console result.  
-- `*each_data_score.csv` score on each data, 3 column `predicted,targets,score`, ranked from the lowest to the highest.  
-- `*predicted.csv` csv file include 3 column `input,predicted,targets`.  
+After evaluate, It will print evaluate result in your console, and also generate three report for debugging.
+
+- `*_score.csv` overall score, it is the copy of the console result.
+- `*each_data_score.csv` score on each data, 3 column `predicted,targets,score`, ranked from the lowest to the highest.
+- `*predicted.csv` csv file include 3 column `input,predicted,targets`.
 
-!!! hint 
-    nlp2go is a tool for demonstration, with CLI and Restful interface. check [here](https://github.com/voidful/nlp2go) 
+!!! hint
+nlp2go is a tool for demonstration, with CLI and Restful interface. check [here](https://github.com/voidful/nlp2go)
 
 ### Example
+
 #### Use distilbert to train NER Model
+
 ```bash
 nlprep --dataset tag_clner  --outdir ./clner_row --util s2t
-tfkit-train --batch 10 --epoch 3 --lr 5e-6 --train ./clner_row/train --test ./clner_row/test --maxlen 512 --task tag --config distilbert-base-multilingual-cased 
-nlp2go --task ./checkpoints/3.pt  --cli     
+tfkit-train --batch 10 --epoch 3 --lr 5e-6 --train ./clner_row/train --test ./clner_row/test --maxlen 512 --task tag --config distilbert-base-multilingual-cased
+nlp2go --task ./checkpoints/3.pt  --cli
 ```
 
 #### Use Albert to train DRCD Model Model
+
 ```bash
-nlprep --dataset qa_zh --outdir ./zhqa/   
+nlprep --dataset qa_zh --outdir ./zhqa/
 tfkit-train --maxlen 512 --savedir ./drcd_qa_model/ --train ./zhqa/drcd-train --test ./zhqa/drcd-test --task qa --config voidful/albert_chinese_small  --cache
-nlp2go --task ./drcd_qa_model/3.pt --cli 
+nlp2go --task ./drcd_qa_model/3.pt --cli
 ```
 
 #### Use Albert to train both DRCD Model and NER Model
+
 ```bash
 nlprep --dataset tag_clner  --outdir ./clner_row --util s2t
-nlprep --dataset qa_zh --outdir ./zhqa/ 
+nlprep --dataset qa_zh --outdir ./zhqa/
 tfkit-train --maxlen 300 --savedir ./mt-qaner --train ./clner_row/train ./zhqa/drcd-train --test ./clner_row/test ./zhqa/drcd-test --task tag qa --config voidful/albert_chinese_small
-nlp2go --task ./mt-qaner/3.pt --cli 
+nlp2go --task ./mt-qaner/3.pt --cli
 ```
 
 **You can also try tfkit in Google Colab: [![Google Colab](https://colab.research.google.com/assets/colab-badge.svg "tfkit")](https://colab.research.google.com/drive/1hqaTKxd3VtX2XkvjiO0FMtY-rTZX30MJ?usp=sharing)**
 
 ## Contributing
+
 Thanks for your interest.There are many ways to contribute to this project. Get started [here](https://github.com/voidful/tfkit/blob/master/CONTRIBUTING.md).
 
-## License 
+## License
+
 ![PyPI - License](https://img.shields.io/github/license/voidful/tfkit)
 
-* [License](https://github.com/voidful/tfkit/blob/master/LICENSE)
+- [License](https://github.com/voidful/tfkit/blob/master/LICENSE)
 
 ## Icons reference
-Icons modify from <a href="http://www.freepik.com/" title="Freepik">Freepik</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>      
-Icons modify from <a href="https://www.flaticon.com/authors/nikita-golubev" title="Nikita Golubev">Nikita Golubev</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>      
+
+Icons modify from <a href="http://www.freepik.com/" title="Freepik">Freepik</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>
+Icons modify from <a href="https://www.flaticon.com/authors/nikita-golubev" title="Nikita Golubev">Nikita Golubev</a> from <a href="https://www.flaticon.com/" title="Flaticon">www.flaticon.com</a>
diff --git a/requirements.txt b/requirements.txt
index e461b61..6b2aa15 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -13,4 +13,4 @@ pandas
 accelerate>=0.5.1
 joblib
 pandas>=1.4.1
-scikit-learn
\ No newline at end of file
+scikit-learn
diff --git a/tfkit/__init__.py b/tfkit/__init__.py
index ebfc013..cc3e04f 100644
--- a/tfkit/__init__.py
+++ b/tfkit/__init__.py
@@ -1,5 +1,5 @@
-import tfkit.utility
 import tfkit.dump
-import tfkit.train
 import tfkit.eval
-from tfkit.task import *
\ No newline at end of file
+import tfkit.train
+import tfkit.utility
+from tfkit.task import *
diff --git a/tfkit/dump.py b/tfkit/dump.py
index 9cf99f4..b4727e2 100644
--- a/tfkit/dump.py
+++ b/tfkit/dump.py
@@ -1,9 +1,12 @@
 import argparse
 import sys
 
-from transformers import AutoModelWithLMHead, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
+from transformers import AutoModelForSeq2SeqLM
+from transformers import AutoModelForSequenceClassification
+from transformers import AutoModelWithLMHead
 
-from tfkit.utility.model import load_trained_model, add_tokens_to_pretrain
+from tfkit.utility.model import add_tokens_to_pretrain
+from tfkit.utility.model import load_trained_model
 
 
 def parse_dump_args(args):
@@ -14,36 +17,43 @@ def parse_dump_args(args):
 
 
 def main(arg=None):
-    arg = parse_dump_args(sys.argv[1:]) if arg is None else parse_dump_args(arg)
-    model, model_type, model_class, model_info, model_preprocessor = load_trained_model(arg.get('model'))
+    arg = parse_dump_args(
+        sys.argv[1:]) if arg is None else parse_dump_args(arg)
+    model, model_type, model_class, model_info, model_preprocessor = load_trained_model(
+        arg.get("model"))
     tokenizer = model.tokenizer
     pretrained_config = model_info.get("model_config")
-    if model_type == 'clm' and "gpt" in pretrained_config:
-        hf_model = AutoModelWithLMHead.from_pretrained(model_info.get("model_config"))
+    if model_type == "clm" and "gpt" in pretrained_config:
+        hf_model = AutoModelWithLMHead.from_pretrained(
+            model_info.get("model_config"))
         hf_model.eval()
         hf_model.transformer = model.pretrained
         hf_model.lm_head.weight = model.model.weight
         hf_model.config.tie_word_embeddings = False
-        hf_model, tokenizer = add_tokens_to_pretrain(hf_model, tokenizer, model_info.get('add_tokens', []))
-        hf_model.save_pretrained(arg.get('dumpdir'))
-    elif model_type == 'seq2seq' and "bart" in pretrained_config:
-        hf_model = AutoModelForSeq2SeqLM.from_pretrained(model_info.get("model_config"))
+        hf_model, tokenizer = add_tokens_to_pretrain(
+            hf_model, tokenizer, model_info.get("add_tokens", []))
+        hf_model.save_pretrained(arg.get("dumpdir"))
+    elif model_type == "seq2seq" and "bart" in pretrained_config:
+        hf_model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_info.get("model_config"))
         hf_model.eval()
         hf_model.lm_head = model.model
         hf_model.model = model.pretrained
         hf_model.config.tie_word_embeddings = False
         hf_model.config.tie_encoder_decoder = False
-        hf_model, tokenizer = add_tokens_to_pretrain(hf_model, tokenizer, model_info.get('add_tokens', []))
-        hf_model.save_pretrained(arg.get('dumpdir'))
-    elif model_type == 'clas':
-        hf_model = AutoModelForSequenceClassification.from_pretrained(model_info.get("model_config"))
+        hf_model, tokenizer = add_tokens_to_pretrain(
+            hf_model, tokenizer, model_info.get("add_tokens", []))
+        hf_model.save_pretrained(arg.get("dumpdir"))
+    elif model_type == "clas":
+        hf_model = AutoModelForSequenceClassification.from_pretrained(
+            model_info.get("model_config"))
         hf_model.classifier.weight = model.classifier_list[0].weight
-        hf_model.save_pretrained(arg.get('dumpdir'))
+        hf_model.save_pretrained(arg.get("dumpdir"))
     else:
-        model.pretrained.save_pretrained(arg.get('dumpdir'))
+        model.pretrained.save_pretrained(arg.get("dumpdir"))
 
-    tokenizer.save_pretrained(arg.get('dumpdir'))
-    print('==================')
+    tokenizer.save_pretrained(arg.get("dumpdir"))
+    print("==================")
     print("Finish model dump.")
 
 
diff --git a/tfkit/eval.py b/tfkit/eval.py
index b741cb5..01967d7 100644
--- a/tfkit/eval.py
+++ b/tfkit/eval.py
@@ -10,49 +10,83 @@
 from tqdm.auto import tqdm
 
 from tfkit.utility.eval_metric import EvalMetric
-from tfkit.utility.model import load_trained_model, load_predict_parameter
+from tfkit.utility.model import load_predict_parameter
+from tfkit.utility.model import load_trained_model
 
-transformers_logger = logging.getLogger('transformers')
+transformers_logger = logging.getLogger("transformers")
 transformers_logger.setLevel(logging.CRITICAL)
 
 
 def parse_eval_args(args):
     parser = argparse.ArgumentParser()
     group = parser.add_mutually_exclusive_group(required=True)
-    group.add_argument("--model", nargs='+', type=str, help="evaluation model")
-    parser.add_argument("--config", type=str, help='pre-trained task path after add token')
-    parser.add_argument("--metric", required=True, type=str, choices=['emf1', 'nlg', 'clas', 'er'],
-                        help="evaluate metric")
-    parser.add_argument("--valid", required=True, type=str, nargs='+', help="evaluate data path")
-    parser.add_argument("--tag", type=str, help="evaluate task tag for select multi-task task")
-    parser.add_argument("--print", action='store_true', help="print each pair of evaluate data")
-    parser.add_argument("--panel", action='store_true', help="enable panel to input argument")
+    group.add_argument("--model", nargs="+", type=str, help="evaluation model")
+    parser.add_argument("--config",
+                        type=str,
+                        help="pre-trained task path after add token")
+    parser.add_argument(
+        "--metric",
+        required=True,
+        type=str,
+        choices=["emf1", "nlg", "clas", "er"],
+        help="evaluate metric",
+    )
+    parser.add_argument("--valid",
+                        required=True,
+                        type=str,
+                        nargs="+",
+                        help="evaluate data path")
+    parser.add_argument("--tag",
+                        type=str,
+                        help="evaluate task tag for select multi-task task")
+    parser.add_argument("--print",
+                        action="store_true",
+                        help="print each pair of evaluate data")
+    parser.add_argument("--panel",
+                        action="store_true",
+                        help="enable panel to input argument")
 
     input_arg, model_arg = parser.parse_known_args(args)
     input_arg = {k: v for k, v in vars(input_arg).items() if v is not None}
-    model_arg = {k.replace("--", ""): v for k, v in zip(model_arg[:-1:2], model_arg[1::2])}
+    model_arg = {
+        k.replace("--", ""): v
+        for k, v in zip(model_arg[:-1:2], model_arg[1::2])
+    }
     return input_arg, model_arg
 
 
 def main(arg=None):
     with torch.no_grad():
-        eval_arg, model_arg = parse_eval_args(sys.argv[1:]) if arg is None else parse_eval_args(arg)
-        models = eval_arg.get('model', [])
+        eval_arg, model_arg = (parse_eval_args(sys.argv[1:])
+                               if arg is None else parse_eval_args(arg))
+        models = eval_arg.get("model", [])
         for m in models:
             for f in nlp2.get_files_from_dir(m):
-                if '.pt' == f[-3:]:
+                if ".pt" == f[-3:]:
                     models.append(f)
         for model_path in models:
             start_time = time.time()
-            valid = eval_arg.get('valid')[0]
-            model, model_type, model_class, model_info, preprocessor = load_trained_model(model_path,
-                                                                                          pretrained_config=eval_arg.get(
-                                                                                              'config'),
-                                                                                          tag=eval_arg.get('tag'))
-            predict_parameter = load_predict_parameter(model, model_arg, eval_arg.get('panel'))
-
-            if 'decodenum' in predict_parameter and int(predict_parameter['decodenum']) > 1:
-                eval_metrics = [EvalMetric(model.tokenizer) for _ in range(int(predict_parameter['decodenum']))]
+            valid = eval_arg.get("valid")[0]
+            (
+                model,
+                model_type,
+                model_class,
+                model_info,
+                preprocessor,
+            ) = load_trained_model(
+                model_path,
+                pretrained_config=eval_arg.get("config"),
+                tag=eval_arg.get("tag"),
+            )
+            predict_parameter = load_predict_parameter(model, model_arg,
+                                                       eval_arg.get("panel"))
+
+            if ("decodenum" in predict_parameter
+                    and int(predict_parameter["decodenum"]) > 1):
+                eval_metrics = [
+                    EvalMetric(model.tokenizer)
+                    for _ in range(int(predict_parameter["decodenum"]))
+                ]
             else:
                 eval_metrics = [EvalMetric(model.tokenizer)]
 
@@ -64,80 +98,109 @@ def main(arg=None):
             get_data_item = preprocessor.read_file_to_data(valid)
             for chunk in tqdm(get_data_item):
                 for i in chunk:
-                    input = i['input']
-                    target = i['target']
-                    predict_parameter.update({'input': input})
+                    input = i["input"]
+                    target = i["target"]
+                    predict_parameter.update({"input": input})
                     result, result_dict = model.predict(**predict_parameter)
                     for eval_pos, eval_metric in enumerate(eval_metrics):
                         # predicted can be list of string or string
                         # target should be list of string
                         predicted = result
                         processed_target = target
-                        if 'qa' in model_type:
-                            processed_target = " ".join(input.split(" ")[int(target[0]): int(target[1])])
+                        if "qa" in model_type:
+                            processed_target = " ".join(
+                                input.split(" ")
+                                [int(target[0]):int(target[1])])
                             if len(result) > 0:
-                                predicted = result[0][0] if isinstance(result[0], list) else result[0]
+                                predicted = (result[0][0] if isinstance(
+                                    result[0], list) else result[0])
                             else:
-                                predicted = ''
-                        elif 'onebyone' in model_type or 'seq2seq' in model_type or 'clm' in model_type:
+                                predicted = ""
+                        elif ("onebyone" in model_type
+                              or "seq2seq" in model_type
+                              or "clm" in model_type):
                             processed_target = target
                             if len(result) < eval_pos:
-                                print("Decode size smaller than decode num:", result_dict['label_map'])
+                                print(
+                                    "Decode size smaller than decode num:",
+                                    result_dict["label_map"],
+                                )
                             predicted = result[eval_pos]
-                        elif 'once' in model_type:
+                        elif "once" in model_type:
                             processed_target = target
                             predicted = result[eval_pos]
-                        elif 'mask' in model_type:
+                        elif "mask" in model_type:
                             processed_target = target.split(" ")
                             predicted = result
-                        elif 'tag' in model_type:
-                            predicted = " ".join([list(d.values())[0] for d in result_dict[0]['label_map']])
+                        elif "tag" in model_type:
+                            predicted = " ".join([
+                                list(d.values())[0]
+                                for d in result_dict[0]["label_map"]
+                            ])
                             processed_target = target[0].split(" ")
                             predicted = predicted.split(" ")
-                        if eval_arg.get('print'):
-                            print('===eval===')
+                        if eval_arg.get("print"):
+                            print("===eval===")
                             print("input: ", input)
                             print("target: ", processed_target)
                             print("predicted: ", predicted)
-                            print('==========')
+                            print("==========")
 
-                        eval_metric.add_record(input, predicted, processed_target, eval_arg.get('metric'))
+                        eval_metric.add_record(input, predicted,
+                                               processed_target,
+                                               eval_arg.get("metric"))
 
             for eval_pos, eval_metric in enumerate(eval_metrics):
                 argtype = "_dataset" + valid.replace("/", "_").replace(".", "")
-                if 'decodenum' in predict_parameter and int(predict_parameter['decodenum']) > 1:
+                if ("decodenum" in predict_parameter
+                        and int(predict_parameter["decodenum"]) > 1):
                     argtype += "_num_" + str(eval_pos)
-                if 'mode' in predict_parameter:
-                    para_mode = predict_parameter['mode'][0] if isinstance(predict_parameter['mode'], list) else \
-                        predict_parameter['mode'].lower()
+                if "mode" in predict_parameter:
+                    para_mode = (predict_parameter["mode"][0] if isinstance(
+                        predict_parameter["mode"], list) else
+                                 predict_parameter["mode"].lower())
                     argtype += "_mode_" + str(para_mode)
-                if 'filtersim' in predict_parameter:
-                    argtype += "_filtersim_" + str(predict_parameter['filtersim'])
+                if "filtersim" in predict_parameter:
+                    argtype += "_filtersim_" + str(
+                        predict_parameter["filtersim"])
                 outfile_name = model_path + argtype
 
-                with open(outfile_name + "_predicted.csv", "w", encoding='utf8') as f:
+                with open(outfile_name + "_predicted.csv",
+                          "w",
+                          encoding="utf8") as f:
                     writer = csv.writer(f)
-                    records = eval_metric.get_record(eval_arg.get('metric'))
-                    writer.writerow(['input', 'predicted', 'targets'])
-                    for i, p, t in zip(records['ori_input'], records['ori_predicted'], records['ori_target']):
+                    records = eval_metric.get_record(eval_arg.get("metric"))
+                    writer.writerow(["input", "predicted", "targets"])
+                    for i, p, t in zip(
+                            records["ori_input"],
+                            records["ori_predicted"],
+                            records["ori_target"],
+                    ):
                         writer.writerow([i, p, t])
                 print("write result at:", outfile_name)
 
-                with open(outfile_name + "_each_data_score.csv", "w", encoding='utf8') as edsf:
+                with open(outfile_name + "_each_data_score.csv",
+                          "w",
+                          encoding="utf8") as edsf:
                     eds = csv.writer(edsf)
-                    with open(outfile_name + "_score.csv", "w", encoding='utf8') as f:
-                        for i in eval_metric.cal_score(eval_arg.get('metric')):
-                            f.write("TASK: " + str(i[0]) + " , " + str(eval_pos) + '\n')
-                            f.write(str(i[1]) + '\n')
+                    with open(outfile_name + "_score.csv",
+                              "w",
+                              encoding="utf8") as f:
+                        for i in eval_metric.cal_score(eval_arg.get("metric")):
+                            f.write("TASK: " + str(i[0]) + " , " +
+                                    str(eval_pos) + "\n")
+                            f.write(str(i[1]) + "\n")
                             eds.writerows(i[2])
 
                 print("write score at:", outfile_name)
 
-                for i in eval_metric.cal_score(eval_arg.get('metric')):
+                for i in eval_metric.cal_score(eval_arg.get("metric")):
                     print("TASK: ", i[0], eval_pos)
                     print(i[1])
 
-            print(f"=== Execution time: {timedelta(seconds=(time.time() - start_time))} ===")
+            print(
+                f"=== Execution time: {timedelta(seconds=(time.time() - start_time))} ==="
+            )
 
 
 if __name__ == "__main__":
diff --git a/tfkit/task/__init__.py b/tfkit/task/__init__.py
index 0f423aa..a0c4b8c 100644
--- a/tfkit/task/__init__.py
+++ b/tfkit/task/__init__.py
@@ -1,3 +1,6 @@
-import os, pkgutil
+import os
+import pkgutil
 
-__all__ = list(module for _, module, _ in pkgutil.iter_modules([os.path.dirname(__file__)]))
+__all__ = list(
+    module
+    for _, module, _ in pkgutil.iter_modules([os.path.dirname(__file__)]))
diff --git a/tfkit/task/clas/__init__.py b/tfkit/task/clas/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/clas/__init__.py
+++ b/tfkit/task/clas/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/clas/model.py b/tfkit/task/clas/model.py
index 11ee655..d532a56 100644
--- a/tfkit/task/clas/model.py
+++ b/tfkit/task/clas/model.py
@@ -2,20 +2,28 @@
 import sys
 
 import torch
-from tfkit.utility.predictor import ClassificationPredictor
 from torch import nn
+from torch import sigmoid
+from torch import softmax
+
+from tfkit.task.clas import Preprocessor
+from tfkit.utility.loss import BCEFocalLoss
+from tfkit.utility.loss import FocalLoss
+from tfkit.utility.predictor import ClassificationPredictor
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-from torch import softmax, sigmoid
-from tfkit.task.clas import Preprocessor
-from tfkit.utility.loss import FocalLoss, BCEFocalLoss
-
 
 class Model(nn.Module):
 
-    def __init__(self, tokenizer, pretrained, tasks_detail, maxlen=512, dropout=0.1, **kwargs):
+    def __init__(self,
+                 tokenizer,
+                 pretrained,
+                 tasks_detail,
+                 maxlen=512,
+                 dropout=0.1,
+                 **kwargs):
         super().__init__()
         self.tokenizer = tokenizer
         self.pretrained = pretrained
@@ -28,7 +36,8 @@ def __init__(self, tokenizer, pretrained, tasks_detail, maxlen=512, dropout=0.1,
         self.tasks_detail = tasks_detail
         self.classifier_list = nn.ModuleList()
         for task, labels in tasks_detail.items():
-            self.classifier_list.append(nn.Linear(self.pretrained.config.hidden_size, len(labels)))
+            self.classifier_list.append(
+                nn.Linear(self.pretrained.config.hidden_size, len(labels)))
             self.tasks[task] = len(self.classifier_list) - 1
         self.maxlen = maxlen
 
@@ -57,7 +66,8 @@ def mean_pooling(self, model_output, attention_mask):
         :param attention_mask:
         :return:
         """
-        input_mask_expanded = attention_mask.unsqueeze(-1).expand(model_output.size()).float()
+        input_mask_expanded = (attention_mask.unsqueeze(-1).expand(
+            model_output.size()).float())
         input_mask_expanded[input_mask_expanded < 0] = 0
         sum_embeddings = torch.sum(model_output * input_mask_expanded, 1)
         sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
@@ -65,17 +75,15 @@ def mean_pooling(self, model_output, attention_mask):
 
     def forward(self, batch_data, eval=False, **kwargs):
         # covert input to correct data type
-        tasks = batch_data['task']
-        tasks = [bytes(t).decode(encoding="utf-8", errors="ignore") for t in tasks]
-        inputs = torch.as_tensor(batch_data['input'])
-        targets = torch.as_tensor(batch_data['target'])
-        masks = torch.as_tensor(batch_data['mask'])
+        tasks = batch_data["task"]
+        tasks = [
+            bytes(t).decode(encoding="utf-8", errors="ignore") for t in tasks
+        ]
+        inputs = torch.as_tensor(batch_data["input"])
+        targets = torch.as_tensor(batch_data["target"])
+        masks = torch.as_tensor(batch_data["mask"])
         # define model output
-        result_dict = {
-            'max_item': [],
-            'prob_list': [],
-            'label_prob': []
-        }
+        result_dict = {"max_item": [], "prob_list": [], "label_prob": []}
 
         result_logits = []
         result_labels = []
@@ -85,32 +93,38 @@ def forward(self, batch_data, eval=False, **kwargs):
             task_labels = self.tasks_detail[task]
 
             output = self.pretrained(input.unsqueeze(0), mask.unsqueeze(0))[0]
-            pooled_output = self.dropout(self.mean_pooling(output, mask.unsqueeze(0)))
+            pooled_output = self.dropout(
+                self.mean_pooling(output, mask.unsqueeze(0)))
             classifier_output = self.classifier_list[task_id](pooled_output)
-            reshaped_logit = classifier_output.view(-1, len(task_labels))  # 0 for cls position
+            # 0 for cls position
+            reshaped_logit = classifier_output.view(-1, len(task_labels))
             result_logits.append(reshaped_logit)
             if not eval:
                 target = targets[p]
                 result_labels.append(target)
             else:
-                if 'multi_label' in task:
+                if "multi_label" in task:
                     reshaped_logit = sigmoid(reshaped_logit)
                 else:
                     reshaped_logit = softmax(reshaped_logit, dim=1)
                 logit_prob = reshaped_logit[0].data.tolist()
                 logit_label = dict(zip(task_labels, logit_prob))
-                result_dict['label_prob'].append({task: logit_label})
-                if 'multi_label' in task:
-                    result_dict['max_item'].append({task: [k for k, v in logit_label.items() if v > 0.5]})
+                result_dict["label_prob"].append({task: logit_label})
+                if "multi_label" in task:
+                    result_dict["max_item"].append(
+                        {task: [k for k, v in logit_label.items() if v > 0.5]})
                 else:
-                    result_dict['max_item'].append({task: [task_labels[logit_prob.index(max(logit_prob))]]})
+                    result_dict["max_item"].append({
+                        task: [task_labels[logit_prob.index(max(logit_prob))]]
+                    })
 
         if eval:
             outputs = result_dict
         else:
             loss = 0
-            for logit, labels, task in zip(result_logits, result_labels, tasks):
-                if 'multi_label' in task:
+            for logit, labels, task in zip(result_logits, result_labels,
+                                           tasks):
+                if "multi_label" in task:
                     loss += self.loss_fct_mt(logit, labels.type_as(logit))
                 else:
                     loss += self.loss_fct(logit, labels)
diff --git a/tfkit/task/clas/preprocessor.py b/tfkit/task/clas/preprocessor.py
index 1e9d7ae..1b06be4 100644
--- a/tfkit/task/clas/preprocessor.py
+++ b/tfkit/task/clas/preprocessor.py
@@ -1,5 +1,6 @@
 import torch
 from sklearn.preprocessing import MultiLabelBinarizer
+
 from tfkit.utility import tok
 from tfkit.utility.datafile import get_multiclas_data_from_file
 from tfkit.utility.preprocess import GeneralNLPPreprocessor
@@ -11,26 +12,28 @@ def read_file_to_data(self, path):
         return get_multiclas_data_from_file(path)
 
     def preprocess_component_convert_to_id(self, item, **param_dict):
-        item['input'] = self.tokenizer.convert_tokens_to_ids(item['input'])
+        item["input"] = self.tokenizer.convert_tokens_to_ids(item["input"])
         yield item
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        tinput, task = item['input'], item['task']
-        row_dict = {'task': list(task.encode("utf-8"))}
-        tokenized_input_id = [tok.tok_begin_id(tokenizer)] + tinput + [tok.tok_sep_id(tokenizer)]
+        tinput, task = item["input"], item["task"]
+        row_dict = {"task": list(task.encode("utf-8"))}
+        tokenized_input_id = ([tok.tok_begin_id(tokenizer)] + tinput +
+                              [tok.tok_sep_id(tokenizer)])
         mask_id = [1] * len(tokenized_input_id)
-        tokenized_input_id.extend([tokenizer.pad_token_id] * (maxlen - len(tokenized_input_id)))
+        tokenized_input_id.extend([tokenizer.pad_token_id] *
+                                  (maxlen - len(tokenized_input_id)))
         mask_id.extend([-1] * (maxlen - len(mask_id)))
-        row_dict['input'] = tokenized_input_id
-        row_dict['mask'] = mask_id
-        row_dict['target'] = [-1]
-        if 'target' in item:
-            target = item['target']
-            if 'multi_label' in task:
-                mlb = MultiLabelBinarizer(classes=item['task_dict'][task])
+        row_dict["input"] = tokenized_input_id
+        row_dict["mask"] = mask_id
+        row_dict["target"] = [-1]
+        if "target" in item:
+            target = item["target"]
+            if "multi_label" in task:
+                mlb = MultiLabelBinarizer(classes=item["task_dict"][task])
                 tar = mlb.fit_transform(target)
                 tokenize_label = tar
             else:
-                tokenize_label = [item['task_dict'][task].index(target[0])]
-            row_dict['target'] = tokenize_label
+                tokenize_label = [item["task_dict"][task].index(target[0])]
+            row_dict["target"] = tokenize_label
         return {key: torch.tensor(value) for key, value in row_dict.items()}
diff --git a/tfkit/task/clm/__init__.py b/tfkit/task/clm/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/clm/__init__.py
+++ b/tfkit/task/clm/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/clm/model.py b/tfkit/task/clm/model.py
index 8053923..2e4d9ec 100644
--- a/tfkit/task/clm/model.py
+++ b/tfkit/task/clm/model.py
@@ -1,24 +1,27 @@
 import os
 import sys
 
+import torch
+from torch import nn
+from torch.nn.functional import softmax
+
 from tfkit.task.clm import Preprocessor
 from tfkit.utility.predictor import AutoRegressivePredictor
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import torch
-from torch import nn
-from torch.nn.functional import softmax
-
 
 class Model(nn.Module):
+
     def __init__(self, tokenizer, pretrained, maxlen=512, **kwargs):
         super().__init__()
         self.tokenizer = tokenizer
         self.pretrained = pretrained
-        self.vocab_size = max(self.pretrained.config.vocab_size, self.tokenizer.__len__())
-        self.model = nn.Linear(self.pretrained.config.hidden_size, self.vocab_size)
+        self.vocab_size = max(self.pretrained.config.vocab_size,
+                              self.tokenizer.__len__())
+        self.model = nn.Linear(self.pretrained.config.hidden_size,
+                               self.vocab_size)
         self.maxlen = maxlen
         predictor = AutoRegressivePredictor(self, Preprocessor)
         self.predictor = predictor
@@ -28,9 +31,14 @@ def clean_cache(self):
         self.encoder_outputs = None
         self.past_key_values = None
 
-    def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwargs):
-        inputs = batch_data['input']
-        masks = batch_data['mask']
+    def forward(self,
+                batch_data,
+                eval=False,
+                beamsearch=False,
+                max_return=1,
+                **kwargs):
+        inputs = batch_data["input"]
+        masks = batch_data["mask"]
         tokens_tensor = torch.as_tensor(inputs)
         mask_tensors = torch.as_tensor(masks)
 
@@ -39,24 +47,34 @@ def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwar
 
         if eval:
             result_dict = {}
-            start = batch_data['start'][0]
-            softmax_score = softmax(prediction_scores[0][start], dim=-1).flatten()
+            start = batch_data["start"][0]
+            softmax_score = softmax(prediction_scores[0][start],
+                                    dim=-1).flatten()
             max_item_id = torch.argmax(softmax_score, -1).item()
             max_item_prob = softmax_score[max_item_id].item()
-            result_dict['max_item'] = (self.tokenizer.convert_ids_to_tokens(max_item_id), max_item_prob)
+            result_dict["max_item"] = (
+                self.tokenizer.convert_ids_to_tokens(max_item_id),
+                max_item_prob,
+            )
             if max_return > 1:
                 topK = torch.topk(softmax_score, max_return)
-                prob_result = [(self.tokenizer.convert_ids_to_tokens(tid), prob) for prob, tid in
-                               zip(topK.values.data.tolist(), topK.indices.data.tolist())]
-                result_dict['prob_list'] = softmax_score.data.tolist()[:max_return]
-                result_dict['label_prob'] = prob_result
+                prob_result = [
+                    (self.tokenizer.convert_ids_to_tokens(tid), prob)
+                    for prob, tid in zip(topK.values.data.tolist(),
+                                         topK.indices.data.tolist())
+                ]
+                result_dict["prob_list"] = softmax_score.data.tolist(
+                )[:max_return]
+                result_dict["label_prob"] = prob_result
             outputs = result_dict
         else:
-            targets = batch_data['target']
+            targets = batch_data["target"]
             loss_tensors = torch.as_tensor(targets)
-            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                      loss_tensors.view(-1))
+            loss_fct = nn.CrossEntropyLoss(
+                ignore_index=-1)  # -1 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.vocab_size),
+                loss_tensors.view(-1))
 
             outputs = masked_lm_loss
         return outputs
diff --git a/tfkit/task/clm/preprocessor.py b/tfkit/task/clm/preprocessor.py
index dba495f..dee1163 100644
--- a/tfkit/task/clm/preprocessor.py
+++ b/tfkit/task/clm/preprocessor.py
@@ -5,41 +5,47 @@
 
 
 class Preprocessor(GeneralNLPPreprocessor):
+
     def read_file_to_data(self, path):
         return get_gen_data_from_file(path)
 
     def preprocess_component_convert_to_id(self, item, **param_dict):
-        tokenized_input, target = item['input'], item.get('target', None)
+        tokenized_input, target = item["input"], item.get("target", None)
         tokenized_target = self.tokenizer.tokenize(target) if target else None
         previous = item.get("previous", [])
         if tokenized_target is None:
 
-            yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                   'previous': self.tokenizer.convert_tokens_to_ids(previous)}
+            yield {
+                "input": self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                "previous": self.tokenizer.convert_tokens_to_ids(previous),
+            }
         else:
-            yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                   'previous': self.tokenizer.convert_tokens_to_ids(previous),
-                   'target': self.tokenizer.convert_tokens_to_ids(tokenized_target)}
+            yield {
+                "input": self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                "previous": self.tokenizer.convert_tokens_to_ids(previous),
+                "target":
+                self.tokenizer.convert_tokens_to_ids(tokenized_target),
+            }
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        t_input_id, previous = item['input'], item['previous']
+        t_input_id, previous = item["input"], item["previous"]
         row_dict = {}
-        if 'target' in item:
-            target = item['target']
+        if "target" in item:
+            target = item["target"]
             t_target_id = [self.tok_pad_id] * len(t_input_id)
             mask_id = [0] * (len(t_target_id))
             t_target_id += target + [self.tok_sep_id]
             mask_id += [1] * (len(target + [self.tok_sep_id]))
 
-            row_dict['start'] = [len(t_input_id)]
+            row_dict["start"] = [len(t_input_id)]
             t_input_id += [self.tok_bos_id] + target
             mask_id = [1] * (len(t_input_id))
-            row_dict['target'] = t_target_id
+            row_dict["target"] = t_target_id
         else:
             t_prev_id = [self.tok_sep_id] + previous
             t_input_id.extend(t_prev_id)
             mask_id = [1] * (len(t_input_id))
-            row_dict['start'] = [len(t_input_id) - 1]
-        row_dict['input'] = t_input_id
-        row_dict['mask'] = mask_id
+            row_dict["start"] = [len(t_input_id) - 1]
+        row_dict["input"] = t_input_id
+        row_dict["mask"] = mask_id
         return {key: torch.tensor(value) for key, value in row_dict.items()}
diff --git a/tfkit/task/once/__init__.py b/tfkit/task/once/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/once/__init__.py
+++ b/tfkit/task/once/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/once/model.py b/tfkit/task/once/model.py
index 31e49af..c9bc405 100644
--- a/tfkit/task/once/model.py
+++ b/tfkit/task/once/model.py
@@ -2,24 +2,27 @@
 import sys
 from collections import defaultdict
 
+from torch.nn.functional import softmax
+
 from tfkit.task.once import Preprocessor
+from tfkit.utility.loss import *
 from tfkit.utility.predictor import NonAutoRegressivePredictor
+from tfkit.utility.tok import *
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-from torch.nn.functional import softmax
-from tfkit.utility.loss import *
-from tfkit.utility.tok import *
-
 
 class Model(nn.Module):
+
     def __init__(self, tokenizer, pretrained, maxlen=512, tasks_detail=None):
         super().__init__()
         self.tokenizer = tokenizer
         self.pretrained = pretrained
-        self.vocab_size = max(self.pretrained.config.vocab_size, self.tokenizer.__len__())
-        self.model = nn.Linear(self.pretrained.config.hidden_size, self.vocab_size)
+        self.vocab_size = max(self.pretrained.config.vocab_size,
+                              self.tokenizer.__len__())
+        self.model = nn.Linear(self.pretrained.config.hidden_size,
+                               self.vocab_size)
         self.maxlen = maxlen
 
         predictor = NonAutoRegressivePredictor(self, Preprocessor)
@@ -31,10 +34,10 @@ def clean_cache(self):
         self.past_key_values = None
 
     def forward(self, batch_data, eval=False, max_return=1, **kwargs):
-        inputs = batch_data['input']
-        masks = batch_data['mask']
-        starts = batch_data['start']
-        ends = batch_data['end']
+        inputs = batch_data["input"]
+        masks = batch_data["mask"]
+        starts = batch_data["start"]
+        ends = batch_data["end"]
         tokens_tensor = torch.as_tensor(inputs)
         mask_tensors = torch.as_tensor(masks)
 
@@ -44,11 +47,11 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
 
         if eval:
             result_dict = {
-                'max_item': [],
-                'label_prob': defaultdict(list),
-                'prob_list': []
+                "max_item": [],
+                "label_prob": defaultdict(list),
+                "prob_list": [],
             }
-            start = batch_data['start'][0]
+            start = batch_data["start"][0]
             stop = False
             topK_ids = [[]] * max_return
             topK_probs = [1] * max_return
@@ -58,7 +61,9 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
                 max_item_prob = softmax_score[max_item_id].item()
                 if max_return > 1:
                     topK = torch.topk(softmax_score, max_return)
-                    for k, (prob, tid) in enumerate(zip(topK.values.data.tolist(), topK.indices.data.tolist())):
+                    for k, (prob, tid) in enumerate(
+                            zip(topK.values.data.tolist(),
+                                topK.indices.data.tolist())):
                         topK_ids[k].append(tid)
                         topK_probs[k] *= prob
                 else:
@@ -68,23 +73,28 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
                 if tok_sep_id(self.tokenizer) == max_item_id:
                     stop = True
                 start += 1
-            result_dict['prob_list'] = topK_probs
-            result_dict['label_prob'] = [[self.tokenizer.decode(ids), prob] for ids, prob in
-                                         zip(topK_ids, topK_probs)]
-            result_dict['max_item'] = [i[0] for i in result_dict['label_prob']]
+            result_dict["prob_list"] = topK_probs
+            result_dict["label_prob"] = [[
+                self.tokenizer.decode(ids), prob
+            ] for ids, prob in zip(topK_ids, topK_probs)]
+            result_dict["max_item"] = [i[0] for i in result_dict["label_prob"]]
             outputs = result_dict
         else:
-            targets = batch_data['target']
-            negative_targets = batch_data['ntarget']
+            targets = batch_data["target"]
+            negative_targets = batch_data["ntarget"]
             loss_tensors = torch.as_tensor(targets)
             negativeloss_tensors = torch.as_tensor(negative_targets)
-            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                      loss_tensors.view(-1))
+            loss_fct = nn.CrossEntropyLoss(
+                ignore_index=-1)  # -1 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.vocab_size),
+                loss_tensors.view(-1))
             if not torch.all(negativeloss_tensors.eq(-1)).item():
                 negative_loss_fct = NegativeCElLoss()
-                negative_loss = negative_loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                                  negativeloss_tensors.view(-1))
+                negative_loss = negative_loss_fct(
+                    prediction_scores.view(-1, self.vocab_size),
+                    negativeloss_tensors.view(-1),
+                )
                 masked_lm_loss += negative_loss
             outputs = masked_lm_loss
 
diff --git a/tfkit/task/once/preprocessor.py b/tfkit/task/once/preprocessor.py
index b44f325..39723ea 100644
--- a/tfkit/task/once/preprocessor.py
+++ b/tfkit/task/once/preprocessor.py
@@ -4,30 +4,53 @@
 
 
 class Preprocessor(GeneralNLPPreprocessor):
+
     def read_file_to_data(self, path):
         return get_gen_data_from_file(path)
 
     def set_global_parameters(self):
         self.tokenize_target = True
 
-    def preprocess_component_convert_to_id(self, item, likelihood=['none', 'pos', 'neg', 'both'], **param_dict):
-        likelihood = likelihood[0] if isinstance(likelihood, list) else likelihood
-        tokenized_input, tokenized_target, n_target = item['input'], item.get('target', None), item.get('ntarget', None)
-        yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-               'target': self.tokenizer.convert_tokens_to_ids(tokenized_target)}
+    def preprocess_component_convert_to_id(self,
+                                           item,
+                                           likelihood=[
+                                               "none", "pos", "neg", "both"
+                                           ],
+                                           **param_dict):
+        likelihood = likelihood[0] if isinstance(likelihood,
+                                                 list) else likelihood
+        tokenized_input, tokenized_target, n_target = (
+            item["input"],
+            item.get("target", None),
+            item.get("ntarget", None),
+        )
+        yield {
+            "input": self.tokenizer.convert_tokens_to_ids(tokenized_input),
+            "target": self.tokenizer.convert_tokens_to_ids(tokenized_target),
+        }
         if "neg" in likelihood:
             # formatting neg data in csv
             if n_target is None:
                 ntext_arr = [
-                    tok.tok_sep(self.tokenizer) + self.tokenizer.convert_tokens_to_string(tokenized_target)]
+                    tok.tok_sep(self.tokenizer) +
+                    self.tokenizer.convert_tokens_to_string(tokenized_target)
+                ]
             elif tok.tok_sep(self.tokenizer) in n_target:
-                ntext_arr = [ntext.strip() for ntext in n_target.split(tok.tok_sep(self.tokenizer))]
+                ntext_arr = [
+                    ntext.strip()
+                    for ntext in n_target.split(tok.tok_sep(self.tokenizer))
+                ]
             else:
                 ntext_arr = [n_target.strip()]
             for neg_text in ntext_arr:
-                yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                       'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
-                       'ntarget': self.tokenizer.convert_tokens_to_ids(neg_text)}
+                yield {
+                    "input":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                    "target":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                    "ntarget":
+                    self.tokenizer.convert_tokens_to_ids(neg_text),
+                }
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
         tok_pad = tok.tok_pad_id(tokenizer)
@@ -36,29 +59,30 @@ def postprocess(self, item, tokenizer, maxlen, **kwargs):
         tok_mask = tok.tok_mask_id(tokenizer)
 
         row_dict = {}
-        t_input_id = item['input']
+        t_input_id = item["input"]
         encoder_mask_id = [1] * (len(t_input_id))
         encoder_mask_id.extend([0] * (maxlen - len(encoder_mask_id)))
         target_start = len(t_input_id)
         target_end = maxlen
         target_length = target_end - target_start
         t_input_id.extend([tok_pad] * (maxlen - len(t_input_id)))
-        if 'target' in item and item['target'] is not None:
-            target = item['target'] + [tok_sep]
+        if "target" in item and item["target"] is not None:
+            target = item["target"] + [tok_sep]
             target.extend([-1] * (maxlen - len(target)))
-            row_dict['target'] = target
-            row_dict['ntarget'] = [-1] * maxlen
-            if 'ntarget' in item and len(item['ntarget'].strip()) > 0:
-                tokenized_ntarget_id = item['ntarget']
-                tokenized_ntarget_id.extend([-1] * (maxlen - len(tokenized_ntarget_id)))
+            row_dict["target"] = target
+            row_dict["ntarget"] = [-1] * maxlen
+            if "ntarget" in item and len(item["ntarget"].strip()) > 0:
+                tokenized_ntarget_id = item["ntarget"]
+                tokenized_ntarget_id.extend(
+                    [-1] * (maxlen - len(tokenized_ntarget_id)))
                 if len(tokenized_ntarget_id) <= maxlen:
-                    row_dict['ntarget'] = tokenized_ntarget_id
+                    row_dict["ntarget"] = tokenized_ntarget_id
 
         input_length = min(maxlen, target_start * 3)
-        row_dict['input'] = t_input_id
-        row_dict['mask'] = encoder_mask_id
-        row_dict['start'] = target_start
-        row_dict['end'] = maxlen
-        row_dict['input_length'] = input_length
-        row_dict['target_length'] = target_length
+        row_dict["input"] = t_input_id
+        row_dict["mask"] = encoder_mask_id
+        row_dict["start"] = target_start
+        row_dict["end"] = maxlen
+        row_dict["input_length"] = input_length
+        row_dict["target_length"] = target_length
         return row_dict
diff --git a/tfkit/task/oncectc/model.py b/tfkit/task/oncectc/model.py
index fcded05..125f1bc 100644
--- a/tfkit/task/oncectc/model.py
+++ b/tfkit/task/oncectc/model.py
@@ -1,20 +1,21 @@
-import sys
 import os
+import sys
 from collections import defaultdict
 
+from torch.nn.functional import softmax
+
 from tfkit.task.once import Preprocessor
+from tfkit.utility.loss import *
+from tfkit.utility.loss import SeqCTCLoss
 from tfkit.utility.predictor import NonAutoRegressivePredictor
+from tfkit.utility.tok import *
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-from torch.nn.functional import softmax
-from tfkit.utility.loss import *
-from tfkit.utility.tok import *
-from tfkit.utility.loss import SeqCTCLoss
-
 
 class Model(nn.Module):
+
     def __init__(self, tokenizer, pretrained, maxlen=512, tasks_detail=None):
         super().__init__()
         self.tokenizer = tokenizer
@@ -23,10 +24,13 @@ def __init__(self, tokenizer, pretrained, maxlen=512, tasks_detail=None):
         self.blank_token = "<BLANK>"
         self.tokenizer.add_tokens(self.blank_token)
         self.pretrained.resize_token_embeddings(len(tokenizer))
-        self.blank_index = self.tokenizer.convert_tokens_to_ids([self.blank_token])[0]
+        self.blank_index = self.tokenizer.convert_tokens_to_ids(
+            [self.blank_token])[0]
         self.loss = SeqCTCLoss(blank_index=self.blank_index)
-        self.vocab_size = max(self.pretrained.config.vocab_size, self.tokenizer.__len__())
-        self.model = nn.Linear(self.pretrained.config.hidden_size, self.vocab_size)
+        self.vocab_size = max(self.pretrained.config.vocab_size,
+                              self.tokenizer.__len__())
+        self.model = nn.Linear(self.pretrained.config.hidden_size,
+                               self.vocab_size)
         predictor = NonAutoRegressivePredictor(self, Preprocessor)
         self.predictor = predictor
         self.predict = predictor.predict
@@ -36,10 +40,10 @@ def clean_cache(self):
         self.past_key_values = None
 
     def forward(self, batch_data, eval=False, max_return=1, **kwargs):
-        inputs = batch_data['input']
-        masks = batch_data['mask']
-        starts = batch_data['start']
-        ends = batch_data['end']
+        inputs = batch_data["input"]
+        masks = batch_data["mask"]
+        starts = batch_data["start"]
+        ends = batch_data["end"]
         tokens_tensor = torch.as_tensor(inputs)
         mask_tensors = torch.as_tensor(masks)
 
@@ -47,23 +51,26 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
         sequence_output = output[0]
         prediction_scores = self.model(sequence_output)
         batch_size = list(tokens_tensor.shape)[0]
-        prediction_scores = prediction_scores.view(batch_size, -1, self.vocab_size)
+        prediction_scores = prediction_scores.view(batch_size, -1,
+                                                   self.vocab_size)
 
         if eval:
             result_dict = {
-                'max_item': [],
-                'label_prob': defaultdict(list),
-                'prob_list': []
+                "max_item": [],
+                "label_prob": defaultdict(list),
+                "prob_list": [],
             }
-            start = batch_data['start'][0]
+            start = batch_data["start"][0]
             topK_ids = [[]] * max_return
             topK_probs = [1] * max_return
 
             pscore = prediction_scores.detach().cpu()
             predicted_indexs = pscore.argmax(2).tolist()[0]
-            predicted_tokens = self.tokenizer.convert_ids_to_tokens(predicted_indexs)
+            predicted_tokens = self.tokenizer.convert_ids_to_tokens(
+                predicted_indexs)
             output = []
-            for pos, (predicted_index, predicted_token) in enumerate(zip(predicted_indexs, predicted_tokens)):
+            for pos, (predicted_index, predicted_token) in enumerate(
+                    zip(predicted_indexs, predicted_tokens)):
                 if len(output) > 0 and predicted_index == output[-1]:
                     continue
                 if predicted_token == self.blank_token:
@@ -78,7 +85,9 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
                 max_item_prob = softmax_score[max_item_id].item()
                 if max_return > 1:
                     topK = torch.topk(softmax_score, max_return)
-                    for k, (prob, tid) in enumerate(zip(topK.values.data.tolist(), topK.indices.data.tolist())):
+                    for k, (prob, tid) in enumerate(
+                            zip(topK.values.data.tolist(),
+                                topK.indices.data.tolist())):
                         topK_ids[k].append(tid)
                         topK_probs[k] *= prob
                 else:
@@ -86,16 +95,17 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
                     topK_probs[0] *= max_item_prob
                 start += 1
 
-            result_dict['prob_list'] = topK_probs
-            result_dict['label_prob'] = [[self.tokenizer.decode(ids), prob] for ids, prob in
-                                         zip(topK_ids, topK_probs)]
-            result_dict['max_item'] = [i[0] for i in result_dict['label_prob']]
+            result_dict["prob_list"] = topK_probs
+            result_dict["label_prob"] = [[
+                self.tokenizer.decode(ids), prob
+            ] for ids, prob in zip(topK_ids, topK_probs)]
+            result_dict["max_item"] = [i[0] for i in result_dict["label_prob"]]
             outputs = result_dict
         else:
-            targets = batch_data['target']
-            negative_targets = batch_data['ntarget']
-            input_lengths = batch_data['input_length']
-            target_lengths = batch_data['target_length']
+            targets = batch_data["target"]
+            negative_targets = batch_data["ntarget"]
+            input_lengths = batch_data["input_length"]
+            target_lengths = batch_data["target_length"]
 
             target_tensors = torch.as_tensor(targets)
             input_length_tensors = torch.as_tensor(input_lengths)
@@ -103,18 +113,24 @@ def forward(self, batch_data, eval=False, max_return=1, **kwargs):
 
             loss_tensors = torch.as_tensor(targets)
             negativeloss_tensors = torch.as_tensor(negative_targets)
-            ctc_lm_loss = self.loss(prediction_scores,
-                                    input_length_tensors,
-                                    target_tensors.view(batch_size, -1),
-                                    target_length_tensors)
-
-            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                      loss_tensors.view(-1))
+            ctc_lm_loss = self.loss(
+                prediction_scores,
+                input_length_tensors,
+                target_tensors.view(batch_size, -1),
+                target_length_tensors,
+            )
+
+            loss_fct = nn.CrossEntropyLoss(
+                ignore_index=-1)  # -1 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.vocab_size),
+                loss_tensors.view(-1))
             if not torch.all(negativeloss_tensors.eq(-1)).item():
                 negative_loss_fct = NegativeCElLoss()
-                negative_loss = negative_loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                                  negativeloss_tensors.view(-1))
+                negative_loss = negative_loss_fct(
+                    prediction_scores.view(-1, self.vocab_size),
+                    negativeloss_tensors.view(-1),
+                )
                 masked_lm_loss += negative_loss
             outputs = ctc_lm_loss + masked_lm_loss
 
diff --git a/tfkit/task/qa/__init__.py b/tfkit/task/qa/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/qa/__init__.py
+++ b/tfkit/task/qa/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/qa/model.py b/tfkit/task/qa/model.py
index 7218078..ec0bd06 100644
--- a/tfkit/task/qa/model.py
+++ b/tfkit/task/qa/model.py
@@ -1,20 +1,25 @@
 import os
 import sys
 
-from tfkit.utility.predictor import QuestionAnsweringPredictor
-
-dir_path = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
-
 import torch
 import torch.nn as nn
 from torch.nn.functional import softmax
+
 from tfkit.task.qa.preprocessor import Preprocessor
+from tfkit.utility.predictor import QuestionAnsweringPredictor
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
 
 class Model(nn.Module):
 
-    def __init__(self, tokenizer, pretrained, maxlen=128, dropout=0.1, **kwargs):
+    def __init__(self,
+                 tokenizer,
+                 pretrained,
+                 maxlen=128,
+                 dropout=0.1,
+                 **kwargs):
         super().__init__()
         self.tokenizer = tokenizer
         self.pretrained = pretrained
@@ -34,10 +39,10 @@ def __init__(self, tokenizer, pretrained, maxlen=128, dropout=0.1, **kwargs):
         self.predict = predictor.predict
 
     def forward(self, batch_data, eval=False, **kwargs):
-        print("batch_data",batch_data)
-        inputs = torch.as_tensor(batch_data['input'])
-        masks = torch.as_tensor(batch_data['mask'])
-        targets = torch.as_tensor(batch_data['target'])
+        print("batch_data", batch_data)
+        inputs = torch.as_tensor(batch_data["input"])
+        masks = torch.as_tensor(batch_data["mask"])
+        targets = torch.as_tensor(batch_data["target"])
         start_positions, end_positions = targets.split(1, dim=1)
         start_positions = start_positions.squeeze(1)
         end_positions = end_positions.squeeze(1)
@@ -49,18 +54,23 @@ def forward(self, batch_data, eval=False, **kwargs):
         end_logits = end_logits.squeeze(-1)
 
         if eval:
-            result_dict = {
-                'label_prob_all': [],
-                'label_map': []
-            }
+            result_dict = {"label_prob_all": [], "label_map": []}
             reshaped_start_logits = softmax(start_logits, dim=1)
             reshaped_end_logits = softmax(end_logits, dim=1)
             start_prob = reshaped_start_logits.data.tolist()[0]
             end_prob = reshaped_end_logits.data.tolist()[0]
-            result_dict['label_prob_all'].append({'start': dict(zip(range(len(start_prob)), start_prob)),
-                                                  'end': dict(zip(range(len(end_prob)), end_prob))})
-            result_dict['label_map'].append({'start': start_prob.index(max(start_prob)),
-                                             'end': end_prob.index(max(end_prob))})
+            result_dict["label_prob_all"].append({
+                "start":
+                dict(zip(range(len(start_prob)), start_prob)),
+                "end":
+                dict(zip(range(len(end_prob)), end_prob)),
+            })
+            result_dict["label_map"].append({
+                "start":
+                start_prob.index(max(start_prob)),
+                "end":
+                end_prob.index(max(end_prob)),
+            })
             outputs = result_dict
         else:
             start_loss = self.loss_fct(start_logits, start_positions)
diff --git a/tfkit/task/qa/preprocessor.py b/tfkit/task/qa/preprocessor.py
index 3ed4d5c..1763049 100644
--- a/tfkit/task/qa/preprocessor.py
+++ b/tfkit/task/qa/preprocessor.py
@@ -1,34 +1,37 @@
 import nlp2
-import tfkit.utility.tok as tok
 import torch
+
+import tfkit.utility.tok as tok
 from tfkit.utility.datafile import get_qa_data_from_file
 from tfkit.utility.preprocess import GeneralNLPPreprocessor
 
 
 class Preprocessor(GeneralNLPPreprocessor):
+
     def read_file_to_data(self, path):
         return get_qa_data_from_file(path)
 
     def preprocess_component_prepare_input(self, item):
         mapping_index = []
         pos = 1  # cls as start 0
-        input_text_list = nlp2.split_sentence_to_array(item['input'])
+        input_text_list = nlp2.split_sentence_to_array(item["input"])
         for i in input_text_list:
             for _ in range(len(self.tokenizer.tokenize(i))):
                 if _ < 1:
-                    mapping_index.append({'char': i, 'pos': pos})
+                    mapping_index.append({"char": i, "pos": pos})
                 pos += 1
-        item['mapping_index'] = mapping_index
+        item["mapping_index"] = mapping_index
         return item
 
     def preprocess_component_convert_to_id(self, item, **param_dict):
-        input_text, target = item['input'], item.get('target', None)
-        tokenized_input = [tok.tok_begin(self.tokenizer)] + input_text + [tok.tok_sep(self.tokenizer)]
+        input_text, target = item["input"], item.get("target", None)
+        tokenized_input = ([tok.tok_begin(self.tokenizer)] + input_text +
+                           [tok.tok_sep(self.tokenizer)])
         input_id = self.tokenizer.convert_tokens_to_ids(tokenized_input)
-        start_index = item['input_index'][0]
-        end_index = item['input_index'][1]
+        start_index = item["input_index"][0]
+        end_index = item["input_index"][1]
         if target:
-            item['target'] = [0, 0]
+            item["target"] = [0, 0]
             target_start, target_end = target
             ori_start = target_start = int(target_start)
             ori_end = target_end = int(target_end)
@@ -42,25 +45,23 @@ def preprocess_component_convert_to_id(self, item, **param_dict):
             #     target_start = 0
             #     target_end = 0
             # else:
-            for map_pos, map_tok in enumerate(item['mapping_index'][start_index:]):
-                if start_index < map_tok['pos'] <= end_index:
-                    length = len(self.tokenizer.tokenize(map_tok['char']))
+            for map_pos, map_tok in enumerate(
+                    item["mapping_index"][start_index:]):
+                if start_index < map_tok["pos"] <= end_index:
+                    length = len(self.tokenizer.tokenize(map_tok["char"]))
                     if map_pos < ori_start:
                         target_start += length - 1
                     if map_pos < ori_end:
                         target_end += length - 1
-            item['target'] = [target_start + 1, target_end + 1]  # cls +1
+            item["target"] = [target_start + 1, target_end + 1]  # cls +1
 
-        item['input'] = input_id
-        item['mask'] = [1] * len(input_id)
-        item['raw_input'] = tokenized_input
+        item["input"] = input_id
+        item["mask"] = [1] * len(input_id)
+        item["raw_input"] = tokenized_input
         yield item
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        row_dict = {
-            'input': item['input'],
-            'mask': item['mask']
-        }
-        if 'target' in item:
-            row_dict['target'] = item['target']
+        row_dict = {"input": item["input"], "mask": item["mask"]}
+        if "target" in item:
+            row_dict["target"] = item["target"]
         return {key: torch.tensor(value) for key, value in row_dict.items()}
diff --git a/tfkit/task/seq2seq/__init__.py b/tfkit/task/seq2seq/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/seq2seq/__init__.py
+++ b/tfkit/task/seq2seq/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/seq2seq/model.py b/tfkit/task/seq2seq/model.py
index 1f84ac2..3020477 100644
--- a/tfkit/task/seq2seq/model.py
+++ b/tfkit/task/seq2seq/model.py
@@ -1,35 +1,41 @@
+import copy
 import os
 import sys
 
+import torch
+from torch import nn
+from torch.autograd.grad_mode import F
+from torch.nn.functional import softmax
 from transformers import AutoModel
 
 from tfkit.task.seq2seq import Preprocessor
+from tfkit.utility.loss import NegativeCElLoss
+from tfkit.utility.loss import SelfKDLoss
 from tfkit.utility.model import tie_encoder_decoder_weights
 from tfkit.utility.predictor import AutoRegressivePredictor
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-from torch.autograd.grad_mode import F
-import torch
-from torch import nn
-from torch.nn.functional import softmax
-from tfkit.utility.loss import NegativeCElLoss, SelfKDLoss
-import copy
-
 
 class Model(nn.Module):
-    def __init__(self, tokenizer, pretrained, maxlen=512, selfkd=False, **kwargs):
+
+    def __init__(self,
+                 tokenizer,
+                 pretrained,
+                 maxlen=512,
+                 selfkd=False,
+                 **kwargs):
         super().__init__()
         self.maxlen = maxlen
         self.tokenizer = tokenizer
         self.pretrained = pretrained
         self.selfkd = selfkd
         init_weight = None
-        if hasattr(pretrained, 'decoder'):
+        if hasattr(pretrained, "decoder"):
             self.decoder_model = None
             decoder_hidden_size = pretrained.config.hidden_size
-            if hasattr(pretrained, 'shared'):
+            if hasattr(pretrained, "shared"):
                 init_weight = copy.deepcopy(pretrained.shared.weight)
         else:
             decoder_config = copy.deepcopy(pretrained.config)
@@ -37,13 +43,17 @@ def __init__(self, tokenizer, pretrained, maxlen=512, selfkd=False, **kwargs):
             decoder_config.add_cross_attention = True
             self.decoder_model = AutoModel.from_config(decoder_config)
             tie_encoder_decoder_weights(
-                self.pretrained, self.decoder_model,
-                self.decoder_model.base_model_prefix
+                self.pretrained,
+                self.decoder_model,
+                self.decoder_model.base_model_prefix,
             )
             decoder_hidden_size = decoder_config.hidden_size
 
-        self.vocab_size = max(self.pretrained.config.vocab_size, self.tokenizer.__len__())
-        self.model = nn.Linear(decoder_hidden_size, self.vocab_size, bias=False)
+        self.vocab_size = max(self.pretrained.config.vocab_size,
+                              self.tokenizer.__len__())
+        self.model = nn.Linear(decoder_hidden_size,
+                               self.vocab_size,
+                               bias=False)
         if init_weight is not None:
             self.model.weight = init_weight
         self.encoder_outputs = None
@@ -56,14 +66,20 @@ def clean_cache(self):
         self.encoder_outputs = None
         self.past_key_values = None
 
-    def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwargs):
-        input_tensors = torch.as_tensor(batch_data['input'])
-        prev_tensors = torch.as_tensor(batch_data['prev'])
-        encoder_mask_tensors = torch.as_tensor(batch_data['encoder_mask'])
-        decoder_mask_tensors = torch.as_tensor(batch_data['decoder_mask'])
+    def forward(self,
+                batch_data,
+                eval=False,
+                beamsearch=False,
+                max_return=1,
+                **kwargs):
+        input_tensors = torch.as_tensor(batch_data["input"])
+        prev_tensors = torch.as_tensor(batch_data["prev"])
+        encoder_mask_tensors = torch.as_tensor(batch_data["encoder_mask"])
+        decoder_mask_tensors = torch.as_tensor(batch_data["decoder_mask"])
         if self.decoder_model:
             if not eval or self.encoder_outputs is None:
-                outputs = self.pretrained(input_tensors, attention_mask=encoder_mask_tensors)
+                outputs = self.pretrained(input_tensors,
+                                          attention_mask=encoder_mask_tensors)
                 self.encoder_outputs = outputs[0]
             # Decoder
             prediction = self.decoder_model(
@@ -74,8 +90,8 @@ def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwar
                 use_cache=True,
                 return_dict=True,
             )
-            prediction_output = prediction['last_hidden_state']
-            prediction_all_hidden = prediction.get('hidden_states')
+            prediction_output = prediction["last_hidden_state"]
+            prediction_all_hidden = prediction.get("hidden_states")
         else:
             if eval and self.encoder_outputs is not None and not beamsearch:
                 prev_tensors = prev_tensors[..., -1:]
@@ -85,7 +101,7 @@ def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwar
                     past_key_values=self.past_key_values,
                     output_hidden_states=self.selfkd,
                     use_cache=True,
-                    return_dict=True
+                    return_dict=True,
                 )
             elif eval and self.encoder_outputs is not None:
                 prev_tensors = prev_tensors[..., -1:]
@@ -95,7 +111,7 @@ def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwar
                     past_key_values=self.past_key_values,
                     output_hidden_states=self.selfkd,
                     use_cache=True,
-                    return_dict=True
+                    return_dict=True,
                 )
             else:
                 prediction = self.pretrained(
@@ -106,60 +122,74 @@ def forward(self, batch_data, eval=False, beamsearch=False, max_return=1, **kwar
                     past_key_values=self.past_key_values,
                     output_hidden_states=self.selfkd,
                     use_cache=True,
-                    return_dict=True
+                    return_dict=True,
                 )
-            prediction_output = prediction['last_hidden_state']
-            prediction_all_hidden = prediction.get('decoder_hidden_states')
+            prediction_output = prediction["last_hidden_state"]
+            prediction_all_hidden = prediction.get("decoder_hidden_states")
             if eval and self.encoder_outputs is None:
-                self.encoder_outputs = (prediction['encoder_last_hidden_state'])
+                self.encoder_outputs = prediction["encoder_last_hidden_state"]
             if eval and not beamsearch:
-                self.past_key_values = prediction.get('past_key_values', None)
+                self.past_key_values = prediction.get("past_key_values", None)
         prediction_scores = self.model(prediction_output)
         if eval:
             result_dict = {}
             softmax_score = softmax(prediction_scores[0][0], dim=0)
             max_item_id = torch.argmax(softmax_score, -1).item()
             max_item_prob = softmax_score[max_item_id].item()
-            result_dict['max_item'] = (self.tokenizer.convert_ids_to_tokens(max_item_id), max_item_prob)
+            result_dict["max_item"] = (
+                self.tokenizer.convert_ids_to_tokens(max_item_id),
+                max_item_prob,
+            )
             if max_return > 1:
                 topK = torch.topk(softmax_score, max_return)
-                prob_result = [(self.tokenizer.convert_ids_to_tokens(tid), prob) for prob, tid in
-                               zip(topK.values.data.tolist(), topK.indices.data.tolist())]
-                result_dict['prob_list'] = softmax_score.data.tolist()[:max_return]
-                result_dict['label_prob'] = prob_result
+                prob_result = [
+                    (self.tokenizer.convert_ids_to_tokens(tid), prob)
+                    for prob, tid in zip(topK.values.data.tolist(),
+                                         topK.indices.data.tolist())
+                ]
+                result_dict["prob_list"] = softmax_score.data.tolist(
+                )[:max_return]
+                result_dict["label_prob"] = prob_result
             outputs = result_dict
         else:
-            targets = batch_data['target']
-            negative_targets = batch_data['ntarget']
+            targets = batch_data["target"]
+            negative_targets = batch_data["ntarget"]
             loss_tensors = torch.as_tensor(targets)
-            loss_fct = nn.CrossEntropyLoss(ignore_index=-1)  # -1 index = padding token
+            loss_fct = nn.CrossEntropyLoss(
+                ignore_index=-1)  # -1 index = padding token
             lm_loss = loss_fct(prediction_scores.view(-1, self.vocab_size),
                                loss_tensors.view(-1))
             if self.selfkd:
                 selfkdloss_fct = SelfKDLoss(ignore_index=-1)
                 for decoder_hidden in prediction_all_hidden[:-1]:
                     student = self.model(decoder_hidden)
-                    lm_loss += selfkdloss_fct(student.view(-1, self.vocab_size),
-                                              prediction_scores.view(-1, self.vocab_size), loss_tensors.view(-1))
-            if 'btarget' in batch_data:
-                backtran_tensors = torch.as_tensor(batch_data['btarget'])
+                    lm_loss += selfkdloss_fct(
+                        student.view(-1, self.vocab_size),
+                        prediction_scores.view(-1, self.vocab_size),
+                        loss_tensors.view(-1),
+                    )
+            if "btarget" in batch_data:
+                backtran_tensors = torch.as_tensor(batch_data["btarget"])
                 if not torch.all(backtran_tensors.eq(-1)).item():
                     backtran_predation = self.pretrained(
                         input_ids=backtran_tensors,
                         output_hidden_states=True,
-                        return_dict=True
+                        return_dict=True,
                     )
-                    backtran_hidden = backtran_predation[
-                        'encoder_last_hidden_state'] if 'encoder_last_hidden_state' in backtran_predation else \
-                        prediction[
-                            'last_hidden_state']
-                    backtran_loss = F.cosine_similarity(self.encoder_hidden, backtran_hidden).mean()
+                    backtran_hidden = (
+                        backtran_predation["encoder_last_hidden_state"]
+                        if "encoder_last_hidden_state" in backtran_predation
+                        else prediction["last_hidden_state"])
+                    backtran_loss = F.cosine_similarity(
+                        self.encoder_hidden, backtran_hidden).mean()
                     lm_loss += backtran_loss
             negativeloss_tensors = torch.as_tensor(negative_targets)
             if not torch.all(negativeloss_tensors.eq(-1)).item():
                 negative_loss_fct = NegativeCElLoss(ignore_index=-1)
-                negative_loss = negative_loss_fct(prediction_scores.view(-1, self.vocab_size),
-                                                  negativeloss_tensors.view(-1))
+                negative_loss = negative_loss_fct(
+                    prediction_scores.view(-1, self.vocab_size),
+                    negativeloss_tensors.view(-1),
+                )
                 lm_loss += negative_loss
             outputs = lm_loss
         return outputs
diff --git a/tfkit/task/seq2seq/preprocessor.py b/tfkit/task/seq2seq/preprocessor.py
index cfee630..c8701a7 100644
--- a/tfkit/task/seq2seq/preprocessor.py
+++ b/tfkit/task/seq2seq/preprocessor.py
@@ -6,88 +6,132 @@
 
 
 class Preprocessor(GeneralNLPPreprocessor):
+
     def read_file_to_data(self, path):
         return get_gen_data_from_file(path)
 
     def set_global_parameters(self):
         self.tokenize_target = True
 
-    def preprocess_component_convert_to_id(self, item, likelihood=['none', 'pos', 'neg', 'both'], **param_dict):
-        likelihood = likelihood[0] if isinstance(likelihood, list) else likelihood
-        tokenized_input, tokenized_target, n_target, b_target = item['input'], \
-                                                      item.get('target', None), \
-                                                      item.get('ntarget', None), \
-                                                      item.get('btarget', None)
+    def preprocess_component_convert_to_id(self,
+                                           item,
+                                           likelihood=[
+                                               "none", "pos", "neg", "both"
+                                           ],
+                                           **param_dict):
+        likelihood = likelihood[0] if isinstance(likelihood,
+                                                 list) else likelihood
+        tokenized_input, tokenized_target, n_target, b_target = (
+            item["input"],
+            item.get("target", None),
+            item.get("ntarget", None),
+            item.get("btarget", None),
+        )
         previous = item.get("previous", [])
         if tokenized_target is None:
-            yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                   'previous': self.tokenizer.convert_tokens_to_ids(previous)}
+            yield {
+                "input": self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                "previous": self.tokenizer.convert_tokens_to_ids(previous),
+            }
         elif b_target and len(b_target) > 0:
-            yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                   'previous': self.tokenizer.convert_tokens_to_ids(previous),
-                   'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
-                   'btarget': self.tokenizer.encode(b_target)}
+            yield {
+                "input": self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                "previous": self.tokenizer.convert_tokens_to_ids(previous),
+                "target":
+                self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                "btarget": self.tokenizer.encode(b_target),
+            }
         else:
-            if "neg" in likelihood or 'both' in likelihood:
+            if "neg" in likelihood or "both" in likelihood:
                 # formatting neg data in csv
                 if n_target is None:
                     ntext_arr = [
-                        tok.tok_sep(self.tokenizer) + self.tokenizer.convert_tokens_to_string(tokenized_target)]
+                        tok.tok_sep(self.tokenizer) +
+                        self.tokenizer.convert_tokens_to_string(
+                            tokenized_target)
+                    ]
                 elif tok.tok_sep(self.tokenizer) in n_target:
-                    ntext_arr = [ntext.strip() for ntext in n_target.split(tok.tok_sep(self.tokenizer))]
+                    ntext_arr = [
+                        ntext.strip() for ntext in n_target.split(
+                            tok.tok_sep(self.tokenizer))
+                    ]
                 else:
                     ntext_arr = [n_target.strip()]
                 for neg_text in ntext_arr:
-                    yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                           'previous': self.tokenizer.convert_tokens_to_ids(previous),
-                           'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
-                           'ntarget': self.tokenizer.encode(neg_text)}
+                    yield {
+                        "input":
+                        self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                        "previous":
+                        self.tokenizer.convert_tokens_to_ids(previous),
+                        "target":
+                        self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                        "ntarget":
+                        self.tokenizer.encode(neg_text),
+                    }
             else:
-                yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                       'previous': self.tokenizer.convert_tokens_to_ids(previous),
-                       'target': self.tokenizer.convert_tokens_to_ids(tokenized_target)}
+                yield {
+                    "input":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                    "previous":
+                    self.tokenizer.convert_tokens_to_ids(previous),
+                    "target":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                }
 
             # whole sentence masking
-            if 'pos' in likelihood:
-                yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                       'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
-                       'previous': self.tokenizer.convert_tokens_to_ids(
-                           [tok.tok_mask(self.tokenizer)] * len(tokenized_target))}
-            elif 'both' in likelihood:
+            if "pos" in likelihood:
+                yield {
+                    "input":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                    "target":
+                    self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                    "previous":
+                    self.tokenizer.convert_tokens_to_ids(
+                        [tok.tok_mask(self.tokenizer)] *
+                        len(tokenized_target)),
+                }
+            elif "both" in likelihood:
                 for neg_text in ntext_arr:
-                    yield {'input': self.tokenizer.convert_tokens_to_ids(tokenized_input),
-                           'target': self.tokenizer.convert_tokens_to_ids(tokenized_target),
-                           'previous': self.tokenizer.convert_tokens_to_ids(
-                               [tok.tok_mask(self.tokenizer)] * len(tokenized_target)),
-                           'ntarget': self.tokenizer.encode(neg_text)}
+                    yield {
+                        "input":
+                        self.tokenizer.convert_tokens_to_ids(tokenized_input),
+                        "target":
+                        self.tokenizer.convert_tokens_to_ids(tokenized_target),
+                        "previous":
+                        self.tokenizer.convert_tokens_to_ids(
+                            [tok.tok_mask(self.tokenizer)] *
+                            len(tokenized_target)),
+                        "ntarget":
+                        self.tokenizer.encode(neg_text),
+                    }
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        t_input_id, previous = item['input'], item['previous']
+        t_input_id, previous = item["input"], item["previous"]
         row_dict = {}
-        if 'target' in item:
-            target = item['target']
+        if "target" in item:
+            target = item["target"]
             tokenized_target_id = []
             if len(previous) == len(target):
                 tokenized_prev_id = [self.tok_mask_id] * maxlen
             else:
                 tokenized_prev_id = [self.tok_sep_id] + target
             tokenized_target_id.extend(target + [self.tok_sep_id])
-            row_dict['target'] = tokenized_target_id
-            row_dict['target_pad'] = [-1]
-            row_dict['prev'] = tokenized_prev_id
-            row_dict['ntarget'] = [-1] * maxlen
-            if 'ntarget' in item and len(item['ntarget']) > 0:
-                tokenized_ntarget_id = item['ntarget']
+            row_dict["target"] = tokenized_target_id
+            row_dict["target_pad"] = [-1]
+            row_dict["prev"] = tokenized_prev_id
+            row_dict["ntarget"] = [-1] * maxlen
+            if "ntarget" in item and len(item["ntarget"]) > 0:
+                tokenized_ntarget_id = item["ntarget"]
                 if len(tokenized_ntarget_id) <= maxlen:
-                    row_dict['ntarget'] = tokenized_ntarget_id
-            if 'btarget' in item and len(item['btarget']) > 0:
-                row_dict['btarget'] = tokenizer.encode(item['btarget'])
+                    row_dict["ntarget"] = tokenized_ntarget_id
+            if "btarget" in item and len(item["btarget"]) > 0:
+                row_dict["btarget"] = tokenizer.encode(item["btarget"])
         else:
             tokenized_prev_id = [self.tok_sep_id]
             tokenized_prev_id.extend(previous)
-            row_dict['prev'] = tokenized_prev_id
+            row_dict["prev"] = tokenized_prev_id
 
-        row_dict['input'] = t_input_id
-        row_dict['encoder_mask'] = [1] * len(t_input_id)
-        row_dict['decoder_mask'] = [1] * len(tokenized_prev_id)
+        row_dict["input"] = t_input_id
+        row_dict["encoder_mask"] = [1] * len(t_input_id)
+        row_dict["decoder_mask"] = [1] * len(tokenized_prev_id)
         return {key: torch.tensor(value) for key, value in row_dict.items()}
diff --git a/tfkit/task/tag/__init__.py b/tfkit/task/tag/__init__.py
index e436466..0d70e48 100644
--- a/tfkit/task/tag/__init__.py
+++ b/tfkit/task/tag/__init__.py
@@ -1,2 +1,2 @@
-from .preprocessor import Preprocessor
 from .model import Model
+from .preprocessor import Preprocessor
diff --git a/tfkit/task/tag/model.py b/tfkit/task/tag/model.py
index 42c0fb2..9a914f6 100644
--- a/tfkit/task/tag/model.py
+++ b/tfkit/task/tag/model.py
@@ -1,27 +1,35 @@
 import os
 import sys
-from collections import defaultdict, Counter
+from collections import Counter
+from collections import defaultdict
+
+from torch.distributions import Categorical
+from torch.nn.functional import softmax
 
 from tfkit.task.tag import Preprocessor
+from tfkit.utility.loss import *
 from tfkit.utility.predictor import TaggingPredictor
-from torch.distributions import Categorical
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-from torch.nn.functional import softmax
-from tfkit.utility.loss import *
-
 
 class Model(nn.Module):
 
-    def __init__(self, tokenizer, pretrained, tasks_detail, maxlen=512, dropout=0.2, **kwargs):
+    def __init__(self,
+                 tokenizer,
+                 pretrained,
+                 tasks_detail,
+                 maxlen=512,
+                 dropout=0.2,
+                 **kwargs):
         super().__init__()
         labels = list(tasks_detail.values())[0]
         self.tokenizer = tokenizer
         self.pretrained = pretrained
         self.dropout = nn.Dropout(dropout)
-        self.tagger = nn.Linear(self.pretrained.config.hidden_size, len(labels))
+        self.tagger = nn.Linear(self.pretrained.config.hidden_size,
+                                len(labels))
         self.labels = labels
         self.maxlen = maxlen
         self.loss_fct = FocalLoss()
@@ -40,45 +48,47 @@ def forward(self, batch_data, eval=False, separator=" ", **kwargs):
         # bert embedding
         token_tensor = torch.as_tensor(inputs, dtype=torch.long)
         mask_tensors = torch.as_tensor(masks)
-        bert_output = self.pretrained(token_tensor, attention_mask=mask_tensors)
+        bert_output = self.pretrained(token_tensor,
+                                      attention_mask=mask_tensors)
         res = bert_output[0]
         pooled_output = self.dropout(res)
         reshaped_logits = self.tagger(pooled_output)
 
         if eval:
-            result_dict = {
-                'label_prob_all': [],
-                'label_map': []
-            }
+            result_dict = {"label_prob_all": [], "label_map": []}
 
             ilogit = softmax(reshaped_logits[0], dim=1)
             result_labels = ilogit.data.tolist()
-            start, end = batch_data['pos'][0]
-            token_word_mapping = batch_data['token_word_mapping']
-            for pos, logit_prob in enumerate(result_labels[1:]):  # skip cls and sep
-                if start+pos >= len(token_word_mapping):
+            start, end = batch_data["pos"][0]
+            token_word_mapping = batch_data["token_word_mapping"]
+            # skip cls and sep
+            for pos, logit_prob in enumerate(result_labels[1:]):
+                if start + pos >= len(token_word_mapping):
                     break
-                word = token_word_mapping[start + pos]['word']
-                pos = token_word_mapping[start + pos]['pos']
-                if len(result_dict['label_map']) > pos:
-                    O = Counter(result_dict['label_prob_all'][-1][word])
+                word = token_word_mapping[start + pos]["word"]
+                pos = token_word_mapping[start + pos]["pos"]
+                if len(result_dict["label_map"]) > pos:
+                    O = Counter(result_dict["label_prob_all"][-1][word])
                     N = Counter(dict(zip(self.labels, logit_prob)))
                     mean_prob = {k: v / 2 for k, v in (O + N).items()}
-                    result_dict['label_prob_all'][-1] = {word: mean_prob}
-                    result_dict['label_map'][-1] = {
-                        word: max(mean_prob, key=mean_prob.get)}
+                    result_dict["label_prob_all"][-1] = {word: mean_prob}
+                    result_dict["label_map"][-1] = {
+                        word: max(mean_prob, key=mean_prob.get)
+                    }
                 else:
                     max_index = logit_prob.index(max(logit_prob))
-                    result_dict['label_map'].append({word: self.labels[max_index]})
-                    result_dict['label_prob_all'].append({word: dict(zip(self.labels, logit_prob))})
+                    result_dict["label_map"].append(
+                        {word: self.labels[max_index]})
+                    result_dict["label_prob_all"].append(
+                        {word: dict(zip(self.labels, logit_prob))})
 
-            result_dict['token_word_mapping'] = token_word_mapping[start:end]
+            result_dict["token_word_mapping"] = token_word_mapping[start:end]
             outputs = result_dict
         else:
             targets = batch_data["target"]
             target_tensor = torch.as_tensor(targets, dtype=torch.long)
-            loss = self.loss_fct(reshaped_logits.view(-1, len(self.labels)), target_tensor.view(-1))
+            loss = self.loss_fct(reshaped_logits.view(-1, len(self.labels)),
+                                 target_tensor.view(-1))
             outputs = loss
 
         return outputs
-
diff --git a/tfkit/task/tag/preprocessor.py b/tfkit/task/tag/preprocessor.py
index 4a2e3ff..95673bd 100644
--- a/tfkit/task/tag/preprocessor.py
+++ b/tfkit/task/tag/preprocessor.py
@@ -11,8 +11,8 @@ def read_file_to_data(self, path):
         return get_tag_data_from_file(path)
 
     def preprocess(self, item, **param_dict):
-        input_text, target = item['input'], item.get('target', None)
-        separator = param_dict.get('separator', ' ')
+        input_text, target = item["input"], item.get("target", None)
+        separator = param_dict.get("separator", " ")
         word_token_mapping = []
         token_word_mapping = []
         pos = 0
@@ -21,13 +21,25 @@ def preprocess(self, item, **param_dict):
             tokenize_word = self.tokenizer.tokenize(word)
             for _ in range(len(tokenize_word)):
                 if _ < 1:  # only record first token (one word one record)
-                    word_token_mapping.append({'char': word, 'pos': pos, 'len': len(tokenize_word)})
-                token_word_mapping.append({'tok': tokenize_word[_], 'word': word, 'pos': len(word_token_mapping) - 1})
+                    word_token_mapping.append({
+                        "char": word,
+                        "pos": pos,
+                        "len": len(tokenize_word)
+                    })
+                token_word_mapping.append({
+                    "tok": tokenize_word[_],
+                    "word": word,
+                    "pos": len(word_token_mapping) - 1,
+                })
                 pos += 1
 
-        t_input_list, t_pos_list = tok.handle_exceed(self.tokenizer, input_text, self.parameters['maxlen'] - 2,
-                                                     mode=self.parameters.get('handle_exceed'),
-                                                     keep_after_sep=False)
+        t_input_list, t_pos_list = tok.handle_exceed(
+            self.tokenizer,
+            input_text,
+            self.parameters["maxlen"] - 2,
+            mode=self.parameters.get("handle_exceed"),
+            keep_after_sep=False,
+        )
         preprocessed_data = []
         for t_input, t_pos in zip(t_input_list, t_pos_list):  # -1 for cls
             # ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
@@ -37,50 +49,57 @@ def preprocess(self, item, **param_dict):
 
             if target is not None:
                 target_token = []
-                for input_word, target_label in zip(word_token_mapping, target.split(separator)):
-                    if t_pos[0] <= input_word['pos'] < t_pos[1]:
-                        for _ in range(input_word['len']):
+                for input_word, target_label in zip(word_token_mapping,
+                                                    target.split(separator)):
+                    if t_pos[0] <= input_word["pos"] < t_pos[1]:
+                        for _ in range(input_word["len"]):
                             target_token += [target_label]
 
                 target_id = [target_token[0]] + target_token
 
                 if len(input_id) != len(target_id):
-                    print(list(zip(input.split(separator), target.split(separator))))
+                    print(
+                        list(
+                            zip(input.split(separator),
+                                target.split(separator))))
                     print(self.tokenizer.decode(input_id))
                     print(input_id)
                     print(target_id)
-                    print("input target len not equal ", len(input_id), len(target_id))
+                    print("input target len not equal ", len(input_id),
+                          len(target_id))
                     continue
-                row_dict['target'] = target_id
+                row_dict["target"] = target_id
 
-            row_dict['input'] = input_id
-            row_dict['word_token_mapping'] = word_token_mapping
-            row_dict['token_word_mapping'] = token_word_mapping
-            row_dict['end'] = len(input_id)
-            row_dict['pos'] = t_pos
+            row_dict["input"] = input_id
+            row_dict["word_token_mapping"] = word_token_mapping
+            row_dict["token_word_mapping"] = token_word_mapping
+            row_dict["end"] = len(input_id)
+            row_dict["pos"] = t_pos
             preprocessed_data.append(row_dict)
         return preprocessed_data
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        labels = item['task_dict']
-        print("item['input']",len(item['input']))
-        mask_id = [1] * len(item['input'])
+        labels = item["task_dict"]
+        print("item['input']", len(item["input"]))
+        mask_id = [1] * len(item["input"])
         mask_id.extend([0] * (maxlen - len(mask_id)))
-        item['input'].extend([0] * (self.parameters['maxlen'] - len(item['input'])))
+        item["input"].extend([0] *
+                             (self.parameters["maxlen"] - len(item["input"])))
         row_dict = {
-            'input': item['input'],
-            'mask': mask_id,
-            'pos': item['pos'],
+            "input": item["input"],
+            "mask": mask_id,
+            "pos": item["pos"],
         }
         # 'token_word_mapping': item['token_word_mapping']
-        if 'target' in item:
-            print(labels['tag'])
-            target_id = [labels['tag'].index(i) for i in item['target']]
-            if "O" in labels['tag']:
-                target_id = [labels['tag'].index("O")] + target_id
+        if "target" in item:
+            print(labels["tag"])
+            target_id = [labels["tag"].index(i) for i in item["target"]]
+            if "O" in labels["tag"]:
+                target_id = [labels["tag"].index("O")] + target_id
             else:
                 target_id = [target_id[0]] + target_id
-            target_id.extend([0] * (self.parameters['maxlen'] - len(target_id)))
-            row_dict['target'] = target_id
+            target_id.extend([0] *
+                             (self.parameters["maxlen"] - len(target_id)))
+            row_dict["target"] = target_id
 
         return row_dict
diff --git a/tfkit/train.py b/tfkit/train.py
index a07fbcb..b758469 100644
--- a/tfkit/train.py
+++ b/tfkit/train.py
@@ -1,27 +1,30 @@
 import argparse
-import sys
+import logging
 import os
-from datetime import timedelta
+import sys
 import time
+from datetime import timedelta
 from itertools import zip_longest
 
+import nlp2
 import torch
+from accelerate import Accelerator
 from torch.utils import data
 from tqdm.auto import tqdm
 from transformers import get_linear_schedule_with_warmup
-import nlp2
+
 import tfkit
 import tfkit.utility.tok as tok
 from tfkit.utility.dataloader import dataloader_collate
 from tfkit.utility.dataset import get_dataset
-
 from tfkit.utility.logger import Logger
-from tfkit.utility.model import load_model_class, save_model, load_pretrained_tokenizer, load_pretrained_model, \
-    resize_pretrain_tok
-import logging
-from accelerate import Accelerator
+from tfkit.utility.model import load_model_class
+from tfkit.utility.model import load_pretrained_model
+from tfkit.utility.model import load_pretrained_tokenizer
+from tfkit.utility.model import resize_pretrain_tok
+from tfkit.utility.model import save_model
 
-transformers_logger = logging.getLogger('transformers')
+transformers_logger = logging.getLogger("transformers")
 transformers_logger.setLevel(logging.CRITICAL)
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -30,54 +33,143 @@
 
 def parse_train_args(args):
     parser = argparse.ArgumentParser()
-    exceed_mode = nlp2.function_get_all_arg_with_value(tok.handle_exceed)['mode']
-    parser.add_argument("--batch", type=int, default=20, help="batch size, default 20")
-    parser.add_argument("--lr", type=float, nargs='+', default=[5e-5], help="learning rate, default 5e-5")
-    parser.add_argument("--epoch", type=int, default=10, help="epoch, default 10")
-    parser.add_argument("--maxlen", type=int, default=0, help="max tokenized sequence length, default task max len")
-    parser.add_argument("--handle_exceed", choices=exceed_mode,
-                        help='select ways to handle input exceed max length')
-    parser.add_argument("--savedir", type=str, default="checkpoints/", help="task saving dir, default /checkpoints")
-    parser.add_argument("--add_tokens_freq", type=int, default=0,
-                        help="auto add freq >= x UNK token to word table")
-    parser.add_argument("--add_tokens_file", type=str,
+    exceed_mode = nlp2.function_get_all_arg_with_value(
+        tok.handle_exceed)["mode"]
+    parser.add_argument("--batch",
+                        type=int,
+                        default=20,
+                        help="batch size, default 20")
+    parser.add_argument(
+        "--lr",
+        type=float,
+        nargs="+",
+        default=[5e-5],
+        help="learning rate, default 5e-5",
+    )
+    parser.add_argument("--epoch",
+                        type=int,
+                        default=10,
+                        help="epoch, default 10")
+    parser.add_argument(
+        "--maxlen",
+        type=int,
+        default=0,
+        help="max tokenized sequence length, default task max len",
+    )
+    parser.add_argument(
+        "--handle_exceed",
+        choices=exceed_mode,
+        help="select ways to handle input exceed max length",
+    )
+    parser.add_argument(
+        "--savedir",
+        type=str,
+        default="checkpoints/",
+        help="task saving dir, default /checkpoints",
+    )
+    parser.add_argument(
+        "--add_tokens_freq",
+        type=int,
+        default=0,
+        help="auto add freq >= x UNK token to word table",
+    )
+    parser.add_argument("--add_tokens_file",
+                        type=str,
                         help="add token from a list file")
-    parser.add_argument("--add_tokens_config", type=str,
+    parser.add_argument("--add_tokens_config",
+                        type=str,
                         help="add token from tokenizer config")
-    parser.add_argument("--train", type=str, nargs='+', required=True, help="train dataset path")
-    parser.add_argument("--test", type=str, nargs='+', required=True, help="test dataset path")
-    parser.add_argument("--no_eval", action='store_true', help="not running evaluation")
-    parser.add_argument("--task", type=str, required=True, nargs='+',
-                        choices=tfkit.utility.model.list_all_model(), help="task task")
-    parser.add_argument("--tag", type=str, nargs='+', help="tag to identity task in multi-task")
-    parser.add_argument("--config", type=str, default='bert-base-multilingual-cased', required=True,
-                        help='distilbert-base-multilingual-cased|voidful/albert_chinese_small')
-    parser.add_argument("--tok_config", type=str,
-                        help='tokenizer config')
-    parser.add_argument("--seed", type=int, default=609, help="random seed, default 609")
-    parser.add_argument("--worker", type=int, default=8, help="number of worker on pre-processing, default 8")
-    parser.add_argument("--grad_accum", type=int, default=1, help="gradient accumulation, default 1")
-    parser.add_argument('--tensorboard', action='store_true', help='Turn on tensorboard graphing')
-    parser.add_argument('--wandb', action='store_true', help='Turn on wandb with project name')
-    parser.add_argument("--resume", help='resume training')
-    parser.add_argument("--cache", action='store_true', help='cache training data')
-    parser.add_argument("--panel", action='store_true', help="enable panel to input argument")
+    parser.add_argument("--train",
+                        type=str,
+                        nargs="+",
+                        required=True,
+                        help="train dataset path")
+    parser.add_argument("--test",
+                        type=str,
+                        nargs="+",
+                        required=True,
+                        help="test dataset path")
+    parser.add_argument("--no_eval",
+                        action="store_true",
+                        help="not running evaluation")
+    parser.add_argument(
+        "--task",
+        type=str,
+        required=True,
+        nargs="+",
+        choices=tfkit.utility.model.list_all_model(),
+        help="task task",
+    )
+    parser.add_argument("--tag",
+                        type=str,
+                        nargs="+",
+                        help="tag to identity task in multi-task")
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="bert-base-multilingual-cased",
+        required=True,
+        help="distilbert-base-multilingual-cased|voidful/albert_chinese_small",
+    )
+    parser.add_argument("--tok_config", type=str, help="tokenizer config")
+    parser.add_argument("--seed",
+                        type=int,
+                        default=609,
+                        help="random seed, default 609")
+    parser.add_argument(
+        "--worker",
+        type=int,
+        default=8,
+        help="number of worker on pre-processing, default 8",
+    )
+    parser.add_argument("--grad_accum",
+                        type=int,
+                        default=1,
+                        help="gradient accumulation, default 1")
+    parser.add_argument("--tensorboard",
+                        action="store_true",
+                        help="Turn on tensorboard graphing")
+    parser.add_argument("--wandb",
+                        action="store_true",
+                        help="Turn on wandb with project name")
+    parser.add_argument("--resume", help="resume training")
+    parser.add_argument("--cache",
+                        action="store_true",
+                        help="cache training data")
+    parser.add_argument("--panel",
+                        action="store_true",
+                        help="enable panel to input argument")
 
     input_arg, model_arg = parser.parse_known_args(args)
     input_arg = {k: v for k, v in vars(input_arg).items() if v is not None}
-    model_arg = {k.replace("--", ""): v for k, v in zip(model_arg[:-1:2], model_arg[1::2])}
+    model_arg = {
+        k.replace("--", ""): v
+        for k, v in zip(model_arg[:-1:2], model_arg[1::2])
+    }
 
     return input_arg, model_arg
 
 
 def optimizer(model, lr, total_step):
     optim = torch.optim.AdamW(model.parameters(), lr=lr)
-    scheduler = get_linear_schedule_with_warmup(optim, num_warmup_steps=int(total_step * 0.05),
+    scheduler = get_linear_schedule_with_warmup(optim,
+                                                num_warmup_steps=int(
+                                                    total_step * 0.05),
                                                 num_training_steps=total_step)
     return [optim, scheduler]
 
 
-def model_train(models_list, dataloaders, models_tag, input_arg, epoch, logger, accelerator, fname, add_tokens):
+def model_train(
+    models_list,
+    dataloaders,
+    models_tag,
+    input_arg,
+    epoch,
+    logger,
+    accelerator,
+    fname,
+    add_tokens,
+):
     optims_schs = []
     models = []
     total_iter = 0
@@ -91,23 +183,29 @@ def model_train(models_list, dataloaders, models_tag, input_arg, epoch, logger,
         if not accelerator.state.backend:
             model = torch.nn.DataParallel(model)
         model.train()
-        optim = optimizer(model, input_arg.get('lr')[i] if i < len(input_arg.get('lr')) else input_arg.get('lr')[0],
-                          total_iter_length)
-        model, optim, dataloader = accelerator.prepare(model, optim, dataloader)
+        optim = optimizer(
+            model,
+            input_arg.get("lr")[i]
+            if i < len(input_arg.get("lr")) else input_arg.get("lr")[0],
+            total_iter_length,
+        )
+        model, optim, dataloader = accelerator.prepare(model, optim,
+                                                       dataloader)
         optims_schs.append(optim)
         models.append(model)
         data_iters.append(iter(dataloader))
 
     while not end:
-        for (model, optim_sch, mtag, train_batch) in zip(models, optims_schs, models_tag, data_iters):
+        for (model, optim_sch, mtag,
+             train_batch) in zip(models, optims_schs, models_tag, data_iters):
             optim = optim_sch[0]
             scheduler = optim_sch[1]
             train_batch = next(train_batch, None)
             if train_batch is not None:
                 loss = model(train_batch)
-                loss = loss / input_arg.get('grad_accum')
+                loss = loss / input_arg.get("grad_accum")
                 accelerator.backward(loss.mean())
-                if (total_iter + 1) % input_arg.get('grad_accum') == 0:
+                if (total_iter + 1) % input_arg.get("grad_accum") == 0:
                     optim.step()
                     model.zero_grad()
                     scheduler.step()
@@ -115,23 +213,32 @@ def model_train(models_list, dataloaders, models_tag, input_arg, epoch, logger,
                 logger.write_metric("loss/step", loss.mean().detach(), epoch)
                 if total_iter % 100 == 0 and total_iter != 0:  # monitoring
                     logger.write_log(
-                        f"epoch: {epoch}, tag: {mtag}, task: {model.__class__.__name__}, step: {total_iter}, loss: {t_loss / total_iter if total_iter > 0 else 0}, total:{total_iter_length}")
+                        f"epoch: {epoch}, tag: {mtag}, task: {model.__class__.__name__}, step: {total_iter}, loss: {t_loss / total_iter if total_iter > 0 else 0}, total:{total_iter_length}"
+                    )
                 if total_iter % 50000 == 0 and total_iter != 0:  # cache
-                    save_model(models, input_arg, models_tag, epoch,
-                               f"{fname}_epoch_{epoch}_iter_{total_iter}", logger,
-                               add_tokens=add_tokens,
-                               accelerator=accelerator)
+                    save_model(
+                        models,
+                        input_arg,
+                        models_tag,
+                        epoch,
+                        f"{fname}_epoch_{epoch}_iter_{total_iter}",
+                        logger,
+                        add_tokens=add_tokens,
+                        accelerator=accelerator,
+                    )
             else:
                 end = True
         pbar.update(1)
         total_iter += 1
     pbar.close()
     logger.write_log(
-        f"epoch: {epoch}, step: {total_iter}, loss: {t_loss / total_iter if total_iter > 0 else 0}, total: {total_iter}")
+        f"epoch: {epoch}, step: {total_iter}, loss: {t_loss / total_iter if total_iter > 0 else 0}, total: {total_iter}"
+    )
     return t_loss / total_iter
 
 
-def model_eval(models, dataloaders, fname, input_arg, epoch, logger, accelerator):
+def model_eval(models, dataloaders, fname, input_arg, epoch, logger,
+               accelerator):
     t_loss = 0
     t_length = 0
     for m in models:
@@ -146,7 +253,7 @@ def model_eval(models, dataloaders, fname, input_arg, epoch, logger, accelerator
                 test_batch = next(batch, None)
                 if test_batch is not None:
                     loss = model(test_batch)
-                    loss = loss / input_arg.get('grad_accum')
+                    loss = loss / input_arg.get("grad_accum")
                     t_loss += loss.mean().detach()
                     t_length += 1
                     pbar.update(1)
@@ -160,30 +267,44 @@ def model_eval(models, dataloaders, fname, input_arg, epoch, logger, accelerator
     return avg_t_loss
 
 
-def load_model_and_datas(tokenizer, pretrained, accelerator, model_arg, input_arg):
+def load_model_and_datas(tokenizer, pretrained, accelerator, model_arg,
+                         input_arg):
     models = []
     train_dataset = []
     test_dataset = []
     train_ds_maxlen = 0
     test_ds_maxlen = 0
-    for model_class_name, train_file, test_file in zip_longest(input_arg.get('task'), input_arg.get('train'),
-                                                               input_arg.get('test'),
-                                                               fillvalue=""):
+    for model_class_name, train_file, test_file in zip_longest(
+            input_arg.get("task"),
+            input_arg.get("train"),
+            input_arg.get("test"),
+            fillvalue="",
+    ):
         # get task class
         model_class = load_model_class(model_class_name)
 
         # load dataset
         ds_parameter = {**model_arg, **input_arg}
-        train_ds = get_dataset(train_file, model_class, tokenizer, ds_parameter)
+        train_ds = get_dataset(train_file, model_class, tokenizer,
+                               ds_parameter)
         test_ds = get_dataset(test_file, model_class, tokenizer, ds_parameter)
 
         # load task
-        model = model_class.Model(tokenizer=tokenizer, pretrained=pretrained, tasks_detail=train_ds.task,
-                                  maxlen=input_arg.get('maxlen'), **model_arg)
+        model = model_class.Model(
+            tokenizer=tokenizer,
+            pretrained=pretrained,
+            tasks_detail=train_ds.task,
+            maxlen=input_arg.get("maxlen"),
+            **model_arg,
+        )
 
         # append to max len
-        train_ds_maxlen = train_ds.__len__() if train_ds.__len__() > train_ds_maxlen else train_ds_maxlen
-        test_ds_maxlen = test_ds.__len__() if test_ds.__len__() > test_ds_maxlen else test_ds_maxlen
+        train_ds_maxlen = (train_ds.__len__()
+                           if train_ds.__len__() > train_ds_maxlen else
+                           train_ds_maxlen)
+        test_ds_maxlen = (test_ds.__len__()
+                          if test_ds.__len__() > test_ds_maxlen else
+                          test_ds_maxlen)
 
         train_dataset.append(train_ds)
         test_dataset.append(test_ds)
@@ -193,52 +314,72 @@ def load_model_and_datas(tokenizer, pretrained, accelerator, model_arg, input_ar
 
 
 def main(arg=None):
-    input_arg, model_arg = parse_train_args(sys.argv[1:]) if arg is None else parse_train_args(arg)
+    input_arg, model_arg = (parse_train_args(sys.argv[1:])
+                            if arg is None else parse_train_args(arg))
     accelerator = Accelerator()
-    nlp2.get_dir_with_notexist_create(input_arg.get('savedir'))
-    logger = Logger(savedir=input_arg.get('savedir'), tensorboard=input_arg.get('tensorboard', False),
-                    wandb=input_arg.get('wandb', False), print_fn=accelerator.print)
+    nlp2.get_dir_with_notexist_create(input_arg.get("savedir"))
+    logger = Logger(
+        savedir=input_arg.get("savedir"),
+        tensorboard=input_arg.get("tensorboard", False),
+        wandb=input_arg.get("wandb", False),
+        print_fn=accelerator.print,
+    )
     logger.write_log("Accelerator")
     logger.write_log("=======================")
     logger.write_log(accelerator.state)
     logger.write_log("=======================")
-    nlp2.set_seed(input_arg.get('seed'))
+    nlp2.set_seed(input_arg.get("seed"))
 
-    tokenizer = load_pretrained_tokenizer(input_arg.get('tok_config', input_arg['config']))
-    pretrained = load_pretrained_model(input_arg.get('config'), input_arg.get('task'))
+    tokenizer = load_pretrained_tokenizer(
+        input_arg.get("tok_config", input_arg["config"]))
+    pretrained = load_pretrained_model(input_arg.get("config"),
+                                       input_arg.get("task"))
     pretrained, tokenizer = resize_pretrain_tok(pretrained, tokenizer)
-    if input_arg.get('maxlen') == 0:
-        if hasattr(pretrained.config, 'max_position_embeddings'):
-            input_arg.update({'maxlen': pretrained.config.max_position_embeddings})
+    if input_arg.get("maxlen") == 0:
+        if hasattr(pretrained.config, "max_position_embeddings"):
+            input_arg.update(
+                {"maxlen": pretrained.config.max_position_embeddings})
         else:
-            input_arg.update({'maxlen': 1024})
+            input_arg.update({"maxlen": 1024})
 
     # handling add tokens
     add_tokens = None
-    if input_arg.get('add_tokens_freq', None):
+    if input_arg.get("add_tokens_freq", None):
         logger.write_log("Calculating Unknown Token")
-        add_tokens = tok.get_freqK_unk_token(tokenizer, input_arg.get('train') + input_arg.get('test'),
-                                             input_arg.get('add_tokens_freq'))
-    if input_arg.get('add_tokens_file', None):
+        add_tokens = tok.get_freqK_unk_token(
+            tokenizer,
+            input_arg.get("train") + input_arg.get("test"),
+            input_arg.get("add_tokens_freq"),
+        )
+    if input_arg.get("add_tokens_file", None):
         logger.write_log("Add token from file")
-        add_tokens = nlp2.read_files_into_lines(input_arg.get('add_tokens_file'))
+        add_tokens = nlp2.read_files_into_lines(
+            input_arg.get("add_tokens_file"))
 
-    if input_arg.get('add_tokens_config', None):
+    if input_arg.get("add_tokens_config", None):
         logger.write_log("Add token from config")
-        add_tokens = tok.get_all_tok_from_config(input_arg.get('add_tokens_config'))
+        add_tokens = tok.get_all_tok_from_config(
+            input_arg.get("add_tokens_config"))
 
     if add_tokens:
-        pretrained, tokenizer = tfkit.utility.model.add_tokens_to_pretrain(pretrained, tokenizer, add_tokens,
-                                                                           sample_init=True)
+        pretrained, tokenizer = tfkit.utility.model.add_tokens_to_pretrain(
+            pretrained, tokenizer, add_tokens, sample_init=True)
 
     # load task and data
-    models_tag = input_arg.get('tag') if input_arg.get('tag', None) is not None else [m.lower() + "_" + str(ind) for
-                                                                                      ind, m in
-                                                                                      enumerate(input_arg.get('task'))]
-
-    models, train_dataset, test_dataset, train_ds_maxlen, test_ds_maxlen = load_model_and_datas(tokenizer, pretrained,
-                                                                                                accelerator, model_arg,
-                                                                                                input_arg)
+    models_tag = (input_arg.get("tag")
+                  if input_arg.get("tag", None) is not None else [
+                      m.lower() + "_" + str(ind)
+                      for ind, m in enumerate(input_arg.get("task"))
+                  ])
+
+    (
+        models,
+        train_dataset,
+        test_dataset,
+        train_ds_maxlen,
+        test_ds_maxlen,
+    ) = load_model_and_datas(tokenizer, pretrained, accelerator, model_arg,
+                             input_arg)
     # balance sample for multi-task
     for ds in train_dataset:
         ds.increase_with_sampling(train_ds_maxlen)
@@ -247,68 +388,109 @@ def main(arg=None):
     logger.write_config(input_arg)
     logger.write_log("TRAIN PARAMETER")
     logger.write_log("=======================")
-    [logger.write_log(str(key) + " : " + str(value)) for key, value in input_arg.items()]
+    [
+        logger.write_log(str(key) + " : " + str(value))
+        for key, value in input_arg.items()
+    ]
     logger.write_log("=======================")
 
-    train_dataloaders = [data.DataLoader(dataset=ds,
-                                         batch_size=input_arg.get('batch'),
-                                         shuffle=True,
-                                         pin_memory=False,
-                                         collate_fn=dataloader_collate,
-                                         num_workers=input_arg.get('worker')) for ds in
-                         train_dataset]
-    test_dataloaders = [data.DataLoader(dataset=ds,
-                                        batch_size=input_arg.get('batch'),
-                                        shuffle=False,
-                                        pin_memory=False,
-                                        collate_fn=dataloader_collate,
-                                        num_workers=input_arg.get('worker')) for ds in
-                        test_dataset]
+    train_dataloaders = [
+        data.DataLoader(
+            dataset=ds,
+            batch_size=input_arg.get("batch"),
+            shuffle=True,
+            pin_memory=False,
+            collate_fn=dataloader_collate,
+            num_workers=input_arg.get("worker"),
+        ) for ds in train_dataset
+    ]
+    test_dataloaders = [
+        data.DataLoader(
+            dataset=ds,
+            batch_size=input_arg.get("batch"),
+            shuffle=False,
+            pin_memory=False,
+            collate_fn=dataloader_collate,
+            num_workers=input_arg.get("worker"),
+        ) for ds in test_dataset
+    ]
 
     # loading task back
     start_epoch = 1
-    if input_arg.get('resume'):
-        logger.write_log("Loading back:", input_arg.get('resume'))
-        package = torch.load(input_arg.get('resume'))
-        if 'model_state_dict' in package:
-            models[0].load_state_dict(package['model_state_dict'])
+    if input_arg.get("resume"):
+        logger.write_log("Loading back:", input_arg.get("resume"))
+        package = torch.load(input_arg.get("resume"))
+        if "model_state_dict" in package:
+            models[0].load_state_dict(package["model_state_dict"])
         else:
-            if len(models) != len(package['models']) and not input_arg.get('tag'):
+            if len(models) != len(
+                    package["models"]) and not input_arg.get("tag"):
                 raise Exception(
-                    f"resuming from multi-task task, you should specific which task to use with --tag, from {package['tags']}")
-            elif len(models) != len(package['models']):
-                tags = input_arg.get('tag')
+                    f"resuming from multi-task task, you should specific which task to use with --tag, from {package['tags']}"
+                )
+            elif len(models) != len(package["models"]):
+                tags = input_arg.get("tag")
             else:
-                tags = package['tags']
+                tags = package["tags"]
             for ind, model_tag in enumerate(tags):
-                tag_ind = package['tags'].index(model_tag)
-                models[ind].load_state_dict(package['models'][tag_ind])
-        start_epoch = int(package.get('epoch', 1)) + 1
+                tag_ind = package["tags"].index(model_tag)
+                models[ind].load_state_dict(package["models"][tag_ind])
+        start_epoch = int(package.get("epoch", 1)) + 1
 
     # train/eval loop
-    logger.write_log("training batch : " + str(input_arg.get('batch') * input_arg.get('grad_accum')))
-    for epoch in range(start_epoch, start_epoch + input_arg.get('epoch')):
+    logger.write_log("training batch : " +
+                     str(input_arg.get("batch") * input_arg.get("grad_accum")))
+    for epoch in range(start_epoch, start_epoch + input_arg.get("epoch")):
         start_time = time.time()
-        fname = os.path.join(input_arg.get('savedir'), str(epoch))
+        fname = os.path.join(input_arg.get("savedir"), str(epoch))
 
         logger.write_log(f"=========train at epoch={epoch}=========")
         try:
-            train_avg_loss = model_train(models, train_dataloaders, models_tag, input_arg, epoch, logger, accelerator,
-                                         fname, add_tokens)
+            train_avg_loss = model_train(
+                models,
+                train_dataloaders,
+                models_tag,
+                input_arg,
+                epoch,
+                logger,
+                accelerator,
+                fname,
+                add_tokens,
+            )
             logger.write_metric("train_loss/epoch", train_avg_loss, epoch)
         except KeyboardInterrupt:
-            save_model(models, input_arg, models_tag, epoch, fname + "_interrupt", logger, add_tokens=add_tokens,
-                       accelerator=accelerator)
+            save_model(
+                models,
+                input_arg,
+                models_tag,
+                epoch,
+                fname + "_interrupt",
+                logger,
+                add_tokens=add_tokens,
+                accelerator=accelerator,
+            )
             pass
 
         logger.write_log(f"=========save at epoch={epoch}=========")
-        save_model(models, input_arg, models_tag, epoch, fname, logger, add_tokens=add_tokens, accelerator=accelerator)
-
-        if input_arg.get('no_eval') is False:
+        save_model(
+            models,
+            input_arg,
+            models_tag,
+            epoch,
+            fname,
+            logger,
+            add_tokens=add_tokens,
+            accelerator=accelerator,
+        )
+
+        if input_arg.get("no_eval") is False:
             logger.write_log(f"=========eval at epoch={epoch}=========")
-            eval_avg_loss = model_eval(models, test_dataloaders, fname, input_arg, epoch, logger, accelerator)
+            eval_avg_loss = model_eval(models, test_dataloaders, fname,
+                                       input_arg, epoch, logger, accelerator)
             logger.write_metric("eval_loss/epoch", eval_avg_loss, epoch)
-        logger.write_log(f"=== Epoch execution time: {timedelta(seconds=(time.time() - start_time))} ===")
+        logger.write_log(
+            f"=== Epoch execution time: {timedelta(seconds=(time.time() - start_time))} ==="
+        )
 
 
 if __name__ == "__main__":
diff --git a/tfkit/utility/datafile.py b/tfkit/utility/datafile.py
index 9f68821..f42efd5 100644
--- a/tfkit/utility/datafile.py
+++ b/tfkit/utility/datafile.py
@@ -1,36 +1,37 @@
 import csv
+import warnings
 from collections import defaultdict
 
 import nlp2
 
+from tfkit.utility import tok
+
 
 # ignore sklearn warning
 def warn(*args, **kwargs):
     pass
 
 
-import warnings
-
 warnings.warn = warn
 
-from tfkit.utility import tok
-
 
 def get_multiclas_data_from_file(fpath):
     task_label_dict = defaultdict(list)
-    with open(fpath, 'r') as infile:
+    with open(fpath, "r") as infile:
         reader = csv.DictReader(infile)
         fieldnames = reader.fieldnames
-        headers = ['input'] + ['target_' + str(i) for i in range(len(fieldnames) - 1)]
+        headers = ["input"] + [
+            "target_" + str(i) for i in range(len(fieldnames) - 1)
+        ]
 
         is_multi_label = ""
-        for rows in nlp2.read_csv_chunk(fpath, ','):
+        for rows in nlp2.read_csv_chunk(fpath, ","):
             for row in rows:
                 if tok.UNIVERSAL_SEP in row[1]:
                     is_multi_label = "_multi_label"
                     break
 
-        for rows in nlp2.read_csv_chunk(fpath, ','):
+        for rows in nlp2.read_csv_chunk(fpath, ","):
             for row in rows:
                 start_pos = 1
                 for pos, item in enumerate(row[start_pos:]):
@@ -39,13 +40,16 @@ def get_multiclas_data_from_file(fpath):
                     item = item.strip()
                     if tok.UNIVERSAL_SEP in item:
                         for i in item.split(tok.UNIVERSAL_SEP):
-                            task_label_dict[task].append(i) if i not in task_label_dict[task] else task_label_dict[task]
+                            task_label_dict[task].append(
+                                i) if i not in task_label_dict[
+                                    task] else task_label_dict[task]
                     else:
-                        task_label_dict[task].append(item) if item not in task_label_dict[task] else task_label_dict[
-                            task]
+                        task_label_dict[task].append(
+                            item) if item not in task_label_dict[
+                                task] else task_label_dict[task]
                     task_label_dict[task].sort()
 
-        for rows in nlp2.read_csv_chunk(fpath, ','):
+        for rows in nlp2.read_csv_chunk(fpath, ","):
             chunk = []
             for row in rows:
                 start_pos = 1
@@ -53,64 +57,88 @@ def get_multiclas_data_from_file(fpath):
                     pos += start_pos
                     task = headers[0] + "_" + headers[pos] + is_multi_label
                     item = item.strip()
-                    targets = item.split(tok.UNIVERSAL_SEP) if tok.UNIVERSAL_SEP in item else [item]
-                    targets = [task_label_dict[task][task_label_dict[task].index(target)] for target in targets]
+                    targets = (item.split(tok.UNIVERSAL_SEP)
+                               if tok.UNIVERSAL_SEP in item else [item])
+                    targets = [
+                        task_label_dict[task][task_label_dict[task].index(
+                            target)] for target in targets
+                    ]
                     input = row[0]
-                    chunk.append({"task": task, "input": input, "target": targets})
+                    chunk.append({
+                        "task": task,
+                        "input": input,
+                        "target": targets
+                    })
             yield chunk
         return task_label_dict
 
 
 def get_clas_data_from_file(fpath):
     task_label_dict = defaultdict(list)
-    task = 'clas'
+    task = "clas"
     task_label_dict[task] = []
-    for rows in nlp2.read_csv_chunk(fpath, ','):
+    for rows in nlp2.read_csv_chunk(fpath, ","):
         chunk = []
         for row in rows:
             source_text = row[0]
             target_text = row[1]
             if target_text not in task_label_dict[task]:
                 task_label_dict[task].append(target_text)
-            chunk.append({"task": task, "input": source_text, "target": task_label_dict[task].index(target_text)})
+            chunk.append({
+                "task": task,
+                "input": source_text,
+                "target": task_label_dict[task].index(target_text),
+            })
         yield chunk
     return task_label_dict
 
 
 def get_gen_data_from_file(fpath):
     task_label_dict = defaultdict(list)
-    task = 'gen'
+    task = "gen"
     task_label_dict[task] = []
     print("Reading data from file...")
-    for rows in nlp2.read_csv_chunk(fpath, ','):
+    for rows in nlp2.read_csv_chunk(fpath, ","):
         chunk = []
         for row in rows:
             source_text = str(row[0]).strip()
             target_text = str(row[1]).strip()
             negative_text = str(row[2]).strip() if len(row) > 2 else None
-            chunk.append({"task": task, "input": source_text, "target": target_text, "ntarget": negative_text})
+            chunk.append({
+                "task": task,
+                "input": source_text,
+                "target": target_text,
+                "ntarget": negative_text,
+            })
         yield chunk
     return task_label_dict
 
 
 def get_qa_data_from_file(fpath):
     task_label_dict = defaultdict(list)
-    task = 'qa'
+    task = "qa"
     task_label_dict[task] = []
-    for rows in nlp2.read_csv_chunk(fpath, ','):
+    for rows in nlp2.read_csv_chunk(fpath, ","):
         chunk = []
         for row in rows:
             context, start, end = row
-            chunk.append({"task": task, "input": context, "target": [start,end]})
+            chunk.append({
+                "task": task,
+                "input": context,
+                "target": [start, end]
+            })
         yield chunk
     return task_label_dict
 
 
-def get_tag_data_from_file(fpath, text_index: int = 0, label_index: int = 1, separator=" "):
+def get_tag_data_from_file(fpath,
+                           text_index: int = 0,
+                           label_index: int = 1,
+                           separator=" "):
     task_label_dict = defaultdict(list)
-    task = 'tag'
+    task = "tag"
     labels = []
-    for rows in nlp2.read_csv_chunk(fpath, ','):
+    for rows in nlp2.read_csv_chunk(fpath, ","):
         for row in rows:
             for i in row[1].split(separator):
                 if i not in labels and len(i.strip()) > 0:
@@ -118,29 +146,38 @@ def get_tag_data_from_file(fpath, text_index: int = 0, label_index: int = 1, sep
                     labels.sort()
     task_label_dict[task] = labels
 
-    for rows in nlp2.read_csv_chunk(fpath, ','):
+    for rows in nlp2.read_csv_chunk(fpath, ","):
         chunk = []
         for row in rows:
-            chunk.append({"task": task, "input": row[text_index].strip(), "target": row[label_index].strip(),
-                          'separator': separator})
+            chunk.append({
+                "task": task,
+                "input": row[text_index].strip(),
+                "target": row[label_index].strip(),
+                "separator": separator,
+            })
         yield chunk
     return task_label_dict
 
 
-def get_tag_data_from_file_col(fpath, text_index: int = 0, label_index: int = 1, separator=" ", **kwargs):
+def get_tag_data_from_file_col(fpath,
+                               text_index: int = 0,
+                               label_index: int = 1,
+                               separator=" ",
+                               **kwargs):
     tasks = defaultdict(list)
-    task = 'default'
+    task = "default"
     labels = []
-    with open(fpath, 'r', encoding='utf-8') as f:
+    with open(fpath, "r", encoding="utf-8") as f:
         lines = f.read().splitlines()
         for line in tqdm(lines):
             rows = line.split(separator)
             if len(rows) > 1:
-                if rows[label_index] not in labels and len(rows[label_index]) > 0:
+                if rows[label_index] not in labels and len(
+                        rows[label_index]) > 0:
                     labels.append(rows[label_index])
                     labels.sort()
     tasks[task] = labels
-    with open(fpath, 'r', encoding='utf-8') as f:
+    with open(fpath, "r", encoding="utf-8") as f:
         lines = f.read().splitlines()
         x, y = "", ""
         for line in tqdm(lines):
diff --git a/tfkit/utility/dataloader.py b/tfkit/utility/dataloader.py
index 2226ed9..1e67f13 100644
--- a/tfkit/utility/dataloader.py
+++ b/tfkit/utility/dataloader.py
@@ -6,7 +6,7 @@
 
 def index_of(in_list, val):
     """
-    get token index in list, return -1 when it is not in the list 
+    get token index in list, return -1 when it is not in the list
     :rtype: int
     :param in_list: query list
     :param val: query target
@@ -27,15 +27,21 @@ def pad_batch(batch):
     """
     keys = list(batch[0].keys())
     for k in keys:
-        batch_key_length = [len(i[k]) if not isinstance(i[k], int) else 1 for i in batch]
-        if len(set(batch_key_length)) > 1:  # is all value same? if no, it need to pad with max length
+        batch_key_length = [
+            len(i[k]) if not isinstance(i[k], int) else 1 for i in batch
+        ]
+        if (len(set(batch_key_length)) >
+                1):  # is all value same? if no, it need to pad with max length
             pad_length = max(batch_key_length)
             for idx, _ in enumerate(batch):
                 if f"{k}_pad" in batch[idx]:
-                    padded = nn.ConstantPad1d((0, pad_length - len(batch[idx][k])), batch[idx][f"{k}_pad"][0])
+                    padded = nn.ConstantPad1d(
+                        (0, pad_length - len(batch[idx][k])),
+                        batch[idx][f"{k}_pad"][0])
 
                 else:
-                    padded = nn.ConstantPad1d((0, pad_length - len(batch[idx][k])), 0)
+                    padded = nn.ConstantPad1d(
+                        (0, pad_length - len(batch[idx][k])), 0)
                 # batch[idx][k] = torch.unsqueeze(padded(batch[idx][k]), 0)
                 batch[idx][k] = padded(batch[idx][k])
     for ind, dat in enumerate(batch):
diff --git a/tfkit/utility/dataset.py b/tfkit/utility/dataset.py
index 0384002..8537c43 100644
--- a/tfkit/utility/dataset.py
+++ b/tfkit/utility/dataset.py
@@ -18,24 +18,30 @@ def get_dataset(file_path, task_class, tokenizer, parameter):
     #         panel.add_element(k=missarg, v=all_arg[missarg], msg=missarg, default=all_arg[missarg])
     #     filled_arg = panel.get_result_dict()
     #     parameter.update(filled_arg)
-    ds = TFKitDataset(fpath=file_path, tokenizer=tokenizer,
-                      preprocessor=task_class.Preprocessor,
-                      preprocessing_arg=parameter)
+    ds = TFKitDataset(
+        fpath=file_path,
+        tokenizer=tokenizer,
+        preprocessor=task_class.Preprocessor,
+        preprocessing_arg=parameter,
+    )
     return ds
 
 
 class TFKitDataset(data.Dataset):
+
     def __init__(self, fpath, tokenizer, preprocessor, preprocessing_arg={}):
-        cache_path = fpath + "_" + tokenizer.name_or_path.replace("/", "_") + ".cache"
+        cache_path = fpath + "_" + tokenizer.name_or_path.replace(
+            "/", "_") + ".cache"
         self.task_dict = {}
         self.preprocessor = preprocessor(tokenizer, kwargs=preprocessing_arg)
         self.tokenizer = tokenizer
-        if os.path.isfile(cache_path) and preprocessing_arg.get('cache', False):
+        if os.path.isfile(cache_path) and preprocessing_arg.get(
+                "cache", False):
             with open(cache_path, "rb") as fo:
                 outdata = joblib.load(fo)
-                sample = outdata['sample']
-                length = outdata['length']
-                self.task_dict = outdata['task']
+                sample = outdata["sample"]
+                length = outdata["length"]
+                self.task_dict = outdata["task"]
         else:
             print(f"Start preprocessing...")
             sample = defaultdict(list)
@@ -43,8 +49,11 @@ def __init__(self, fpath, tokenizer, preprocessor, preprocessing_arg={}):
             get_data_item = self.preprocessor.read_file_to_data(fpath)
             while True:
                 try:
-                    for items in process_map(self.preprocessor.preprocess, next(get_data_item),
-                                             chunksize=1000):
+                    for items in process_map(
+                            self.preprocessor.preprocess,
+                            next(get_data_item),
+                            chunksize=1000,
+                    ):
                         for i in items:
                             length += 1
                             for k, v in i.items():
@@ -55,9 +64,13 @@ def __init__(self, fpath, tokenizer, preprocessor, preprocessing_arg={}):
                     break
             self.task_dict = tasks
             print(f"There are {length} datas after preprocessing.")
-            if preprocessing_arg.get('cache', False):
-                with open(cache_path, 'wb') as fo:
-                    outdata = {'sample': sample, 'task': self.task_dict, 'length': length}
+            if preprocessing_arg.get("cache", False):
+                with open(cache_path, "wb") as fo:
+                    outdata = {
+                        "sample": sample,
+                        "task": self.task_dict,
+                        "length": length,
+                    }
                     joblib.dump(outdata, fo)
         self.length = length
         self.sample = sample
@@ -73,6 +86,13 @@ def __len__(self):
 
     def __getitem__(self, idx):
         return self.preprocessor.postprocess(
-            {**{'task_dict': self.task_dict}, **{key: self.sample[key][idx] for key in self.sample.keys()}},
+            {
+                **{
+                    "task_dict": self.task_dict
+                },
+                **{key: self.sample[key][idx]
+                   for key in self.sample.keys()},
+            },
             self.tokenizer,
-            maxlen=self.preprocessor.parameters['maxlen'])
+            maxlen=self.preprocessor.parameters["maxlen"],
+        )
diff --git a/tfkit/utility/eval_metric.py b/tfkit/utility/eval_metric.py
index 476fab3..920a0b4 100644
--- a/tfkit/utility/eval_metric.py
+++ b/tfkit/utility/eval_metric.py
@@ -4,30 +4,31 @@
 from collections import Counter
 from collections import defaultdict
 
-from tfkit.utility import tok
 from tqdm.auto import tqdm
 
+from tfkit.utility import tok
+
 
-def _normalize_answer(s, task='emf1'):
+def _normalize_answer(s, task="emf1"):
     """Lower text and remove punctuation, articles and extra whitespace."""
 
     def remove_articles(text):
         if len(text) > 1:
-            return re.sub(r'\b(a|an|the)\b', ' ', text)
+            return re.sub(r"\b(a|an|the)\b", " ", text)
         else:
             return text
 
     def white_space_fix(text):
-        return ' '.join(text.split())
+        return " ".join(text.split())
 
     def remove_punc(text):
         exclude = set(string.punctuation)
-        return ''.join(ch for ch in text if ch not in exclude)
+        return "".join(ch for ch in text if ch not in exclude)
 
     def lower(text):
         return text.lower()
 
-    if task == 'emf1':
+    if task == "emf1":
         return white_space_fix(remove_articles(remove_punc(lower(s))))
     else:
         return white_space_fix((remove_punc(lower(s))))
@@ -55,17 +56,19 @@ def __init__(self, tokenizer, normalize_text=True):
         self.normalize_text = normalize_text
 
     def tokenize_text(self, text):
-        text = self.tokenizer.decode(self.tokenizer.encode(text, add_special_tokens=False))
+        text = self.tokenizer.decode(
+            self.tokenizer.encode(text, add_special_tokens=False))
         if self.normalize_text:
             text = text.replace(tok.tok_sep(self.tokenizer), " ")
             # return  _normalize_answer(text, task='others')  # remove punctuation
             # keep punctuation
             text = "".join(
-                (char if char.isalpha() or char == " " else " " + char + " ") for char in text)  # separate punctuation
-            text = ' '.join(text.split()).lower().strip()  # remove extra blank
+                (char if char.isalpha() or char == " " else " " + char + " ")
+                for char in text)  # separate punctuation
+            text = " ".join(text.split()).lower().strip()  # remove extra blank
         return text
 
-    def add_record(self, ori_input, ori_predicted, ori_target, task='default'):
+    def add_record(self, ori_input, ori_predicted, ori_target, task="default"):
         input = predicted = target = ""
         input_list = predicted_list = ori_predicted_list = target_list = []
 
@@ -90,14 +93,17 @@ def add_record(self, ori_input, ori_predicted, ori_target, task='default'):
                     predicted_list[i] = self.tokenize_text(t.strip())
                     ori_predicted_list[i] = t
                 else:
-                    predicted_list[i] = ''
-                    ori_predicted_list[i] = ''
+                    predicted_list[i] = ""
+                    ori_predicted_list[i] = ""
             predicted = " ".join(predicted_list)
         if isinstance(ori_target, str):
             target_list = []
             if tok.UNIVERSAL_SEP in ori_target:
                 target = ori_target
-                target_list.extend([self.tokenize_text(st.strip()) for st in ori_target.split(tok.UNIVERSAL_SEP)])
+                target_list.extend([
+                    self.tokenize_text(st.strip())
+                    for st in ori_target.split(tok.UNIVERSAL_SEP)
+                ])
             else:
                 target = self.tokenize_text(ori_target.strip())
                 target_list.append(target)
@@ -111,18 +117,18 @@ def add_record(self, ori_input, ori_predicted, ori_target, task='default'):
         for t in target_list:
             self.target_list[task][t] += 1
 
-        self.tasks[task]['input'].append(input)
-        self.tasks[task]['input_list'].append(input_list)
-        self.tasks[task]['predicted'].append(predicted)
-        self.tasks[task]['predicted_list'].append(predicted_list)
-        self.tasks[task]['target'].append(target)
-        self.tasks[task]['target_list'].append(target_list)
-        self.tasks[task]['ori_input'].append(ori_input)
-        self.tasks[task]['ori_predicted'].append(ori_predicted)
-        self.tasks[task]['ori_predicted_list'].append(ori_predicted_list)
-        self.tasks[task]['ori_target'].append(ori_target)
-
-    def get_record(self, task='default'):
+        self.tasks[task]["input"].append(input)
+        self.tasks[task]["input_list"].append(input_list)
+        self.tasks[task]["predicted"].append(predicted)
+        self.tasks[task]["predicted_list"].append(predicted_list)
+        self.tasks[task]["target"].append(target)
+        self.tasks[task]["target_list"].append(target_list)
+        self.tasks[task]["ori_input"].append(ori_input)
+        self.tasks[task]["ori_predicted"].append(ori_predicted)
+        self.tasks[task]["ori_predicted_list"].append(ori_predicted_list)
+        self.tasks[task]["ori_target"].append(ori_target)
+
+    def get_record(self, task="default"):
         return self.tasks[task]
 
     def cal_score(self, metric):
@@ -133,12 +139,14 @@ def cal_score(self, metric):
                 em = 0
                 total = 0
                 f1 = 0
-                for pos, predict in enumerate(task['predicted']):
+                for pos, predict in enumerate(task["predicted"]):
                     em_list = []
                     f1_list = []
-                    for target in task['target_list'][pos]:
-                        if _normalize_answer(str(predict)) == _normalize_answer(str(target)) and len(
-                                _normalize_answer(str(predict))) > 0 or len(str(predict)) == len(str(target)) == 0:
+                    for target in task["target_list"][pos]:
+                        if (_normalize_answer(
+                                str(predict)) == _normalize_answer(str(target))
+                                and len(_normalize_answer(str(predict))) > 0
+                                or len(str(predict)) == len(str(target)) == 0):
                             em_score = 1
                             f1_score = 1
                         else:
@@ -148,89 +156,141 @@ def cal_score(self, metric):
                         f1_list.append(f1_score)
                     em += max(em_list)
                     f1 += max(f1_list)
-                    data_score.append([predict, task['target_list'][pos][em_list.index(max(em_list))],
-                                       {'em': max(em_list), 'f1': max(f1_list)}])
+                    data_score.append([
+                        predict,
+                        task["target_list"][pos][em_list.index(max(em_list))],
+                        {
+                            "em": max(em_list),
+                            "f1": max(f1_list)
+                        },
+                    ])
                     total += 1
-                result = {"EM": em / (total or not total), "F1": f1 / (total or not total)}
-                data_score = sorted(data_score, key=lambda i: i[2]['em'], reverse=True)
+                result = {
+                    "EM": em / (total or not total),
+                    "F1": f1 / (total or not total),
+                }
+                data_score = sorted(data_score,
+                                    key=lambda i: i[2]["em"],
+                                    reverse=True)
             if "er" in metric:
                 try:
                     import asrp
                 except ImportError:
                     print(
-                        "asrp package not install, plz install it: pip install asrp")
+                        "asrp package not install, plz install it: pip install asrp"
+                    )
                     raise
                 predicts = []
                 targets = []
-                for pos, predict in enumerate(task['predicted']):
+                for pos, predict in enumerate(task["predicted"]):
                     wer_list = []
                     cer_list = []
-                    for target in task['target_list'][pos]:
+                    for target in task["target_list"][pos]:
                         if len(target) > 0 and len(predict) > 0:
-                            wer_list.append(100 * asrp.wer([target], [predict]))
-                            cer_list.append(100 * asrp.cer([target], [predict]))
+                            wer_list.append(100 *
+                                            asrp.wer([target], [predict]))
+                            cer_list.append(100 *
+                                            asrp.cer([target], [predict]))
                         else:
                             wer_list.append(100)
                             cer_list.append(100)
                     wer = min(wer_list)
                     cer = min(cer_list)
-                    target = task['target_list'][pos][wer_list.index(wer)]
+                    target = task["target_list"][pos][wer_list.index(wer)]
                     predicts.append(predict)
                     targets.append(target)
-                    data_score.append([predict, target, {'wer': wer, 'cer': cer}])
+                    data_score.append(
+                        [predict, target, {
+                            "wer": wer,
+                            "cer": cer
+                        }])
 
-                wer = 100 * asrp.wer(targets, predicts) if len(target) > 0 else 100
-                cer = 100 * asrp.cer(targets, predicts) if len(target) > 0 else 100
+                wer = 100 * asrp.wer(targets,
+                                     predicts) if len(target) > 0 else 100
+                cer = 100 * asrp.cer(targets,
+                                     predicts) if len(target) > 0 else 100
                 result = {"WER": wer, "CER": cer}
-                data_score = sorted(data_score, key=lambda i: i[2]['wer'], reverse=False)
+                data_score = sorted(data_score,
+                                    key=lambda i: i[2]["wer"],
+                                    reverse=False)
             if "nlg" in metric:
                 try:
                     from nlgeval import NLGEval
                 except ImportError:
                     print(
-                        "nlg-eval package not install, plz install it: pip install git+https://github.com/voidful/nlg-eval.git ; nlg-eval --setup ./nlg-eval-data/")
+                        "nlg-eval package not install, plz install it: pip install git+https://github.com/voidful/nlg-eval.git ; nlg-eval --setup ./nlg-eval-data/"
+                    )
                     raise
-                nlgeval = NLGEval(no_skipthoughts=True, no_glove=True, metrics_to_omit=["METEOR"])
+                nlgeval = NLGEval(no_skipthoughts=True,
+                                  no_glove=True,
+                                  metrics_to_omit=["METEOR"])
 
-                target_list = task['target_list']
-                predicted = task['predicted']
+                target_list = task["target_list"]
+                predicted = task["predicted"]
                 for idx, tl in enumerate(target_list):
                     max_candidate = max([len(i) for i in target_list])
                     if max_candidate - len(tl) > 0:
-                        target_list[idx].extend([""] * (max_candidate - len(tl)))
+                        target_list[idx].extend([""] *
+                                                (max_candidate - len(tl)))
 
-                for t, p in tqdm(zip(target_list, predicted), total=len(target_list)):
-                    data_score.append([p, t, nlgeval.compute_metrics(ref_list=list(map(list, zip(t))), hyp_list=[p])])
-                result = nlgeval.compute_metrics(ref_list=list(map(list, zip(*task['target_list']))),  # transpose
-                                                 hyp_list=predicted)
-                data_score = sorted(data_score, key=lambda i: i[2]['ROUGE_L'])
+                for t, p in tqdm(zip(target_list, predicted),
+                                 total=len(target_list)):
+                    data_score.append([
+                        p,
+                        t,
+                        nlgeval.compute_metrics(ref_list=list(map(
+                            list, zip(t))),
+                                                hyp_list=[p]),
+                    ])
+                result = nlgeval.compute_metrics(
+                    ref_list=list(map(list,
+                                      zip(*task["target_list"]))),  # transpose
+                    hyp_list=predicted,
+                )
+                data_score = sorted(data_score, key=lambda i: i[2]["ROUGE_L"])
             if "clas" in metric:
-                from sklearn.metrics import classification_report
+                from sklearn.metrics import (classification_report,
+                                             precision_recall_fscore_support)
                 from sklearn.preprocessing import MultiLabelBinarizer
-                from sklearn.metrics import precision_recall_fscore_support
-                target_key = [t for t in self.target_list[task_name].keys() if len(t) > 0]
+
+                target_key = [
+                    t for t in self.target_list[task_name].keys() if len(t) > 0
+                ]
                 mlb = MultiLabelBinarizer().fit([target_key])
                 # remove all blank target
-                task['target_list'] = [[j for j in sub if len(j) > 0] for sub in task['target_list']]
+                task["target_list"] = [[j for j in sub if len(j) > 0]
+                                       for sub in task["target_list"]]
                 # modify for tagging result
-                if isinstance(task['ori_predicted_list'][0][0], list):
-                    target_list = sum([[[j] for j in sub] for sub in task['target_list']], [])
-                    predicted = sum([[[j] for j in sub] for sub in task['ori_predicted_list']], [])
+                if isinstance(task["ori_predicted_list"][0][0], list):
+                    target_list = sum([[[j] for j in sub]
+                                       for sub in task["target_list"]], [])
+                    predicted = sum([[[j] for j in sub]
+                                     for sub in task["ori_predicted_list"]],
+                                    [])
                     if len(target_list) != len(predicted):
-                        diff = len(task['target_list']) - len(task['ori_predicted_list'])
-                        predicted.extend([['']] * diff)
+                        diff = len(task["target_list"]) - len(
+                            task["ori_predicted_list"])
+                        predicted.extend([[""]] * diff)
                 else:
-                    target_list = task['target_list']
-                    predicted = task['ori_predicted_list']
+                    target_list = task["target_list"]
+                    predicted = task["ori_predicted_list"]
 
                 for p, t in zip(predicted, target_list):
-                    score = dict(zip(["precision", "recall", "fbeta_score", "support"],
-                                     precision_recall_fscore_support(mlb.transform([t]), mlb.transform([p]),
-                                                                     average='weighted')))
+                    score = dict(
+                        zip(
+                            ["precision", "recall", "fbeta_score", "support"],
+                            precision_recall_fscore_support(
+                                mlb.transform([t]),
+                                mlb.transform([p]),
+                                average="weighted",
+                            ),
+                        ))
                     data_score.append([p, t, score])
                 result = classification_report(
                     mlb.transform(target_list),
                     mlb.transform(predicted),
-                    target_names=list(mlb.classes_))
-                data_score = sorted(data_score, key=lambda i: i[2]['fbeta_score'])
+                    target_names=list(mlb.classes_),
+                )
+                data_score = sorted(data_score,
+                                    key=lambda i: i[2]["fbeta_score"])
             yield (task_name, result, data_score)
diff --git a/tfkit/utility/model.py b/tfkit/utility/model.py
index a77efe8..c545270 100644
--- a/tfkit/utility/model.py
+++ b/tfkit/utility/model.py
@@ -7,31 +7,38 @@
 import nlp2
 import torch
 from torch import nn
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoModel
+from transformers import AutoTokenizer
 
 
 def list_all_model(ignore_list=[]):
-    dataset_dir = os.path.abspath(__file__ + "/../../") + '/task'
-    return list(filter(
-        lambda x: os.path.isdir(os.path.join(dataset_dir, x)) and '__pycache__' not in x and x not in ignore_list,
-        os.listdir(dataset_dir)))
+    dataset_dir = os.path.abspath(__file__ + "/../../") + "/task"
+    return list(
+        filter(
+            lambda x: os.path.isdir(os.path.join(dataset_dir, x)) and
+            "__pycache__" not in x and x not in ignore_list,
+            os.listdir(dataset_dir),
+        ))
 
 
 def load_predict_parameter(model, model_arg={}, enable_arg_panel=False):
     """use inquirer panel to let user input task parameter or just use default value"""
-    return nlp2.function_argument_panel(model.predictor.wrap_input, model_arg,
-                                        disable_input_panel=(not enable_arg_panel),
-                                        func_parent=model,
-                                        ignore_empty=True)
+    return nlp2.function_argument_panel(
+        model.predictor.wrap_input,
+        model_arg,
+        disable_input_panel=(not enable_arg_panel),
+        func_parent=model,
+        ignore_empty=True,
+    )
 
 
 def load_model_class(model_name):
-    return importlib.import_module('.' + model_name, 'tfkit.task')
+    return importlib.import_module("." + model_name, "tfkit.task")
 
 
 def load_pretrained_model(pretrained_config, model_type):
     pretrained = AutoModel.from_pretrained(pretrained_config)
-    if 'clm' in model_type:
+    if "clm" in model_type:
         pretrained.config.is_decoder = True
     return pretrained
 
@@ -47,15 +54,18 @@ def resize_pretrain_tok(pretrained, tokenizer):
     return pretrained, tokenizer
 
 
-def add_tokens_to_pretrain(pretrained, tokenizer, add_tokens, sample_init=False):
+def add_tokens_to_pretrain(pretrained,
+                           tokenizer,
+                           add_tokens,
+                           sample_init=False):
     origin_vocab_size = tokenizer.vocab_size
     print("===ADD TOKEN===")
     num_added_toks = tokenizer.add_tokens(add_tokens)
-    print('We have added', num_added_toks, 'tokens')
+    print("We have added", num_added_toks, "tokens")
     pretrained.resize_token_embeddings(len(tokenizer))
     if sample_init:
         input_embedding = pretrained.get_input_embeddings()
-        state_dict_weight = input_embedding.state_dict()['weight']
+        state_dict_weight = input_embedding.state_dict()["weight"]
         state_dict_weight[origin_vocab_size:len(tokenizer)] = copy.copy(
             state_dict_weight[100:100 + num_added_toks])
         pretrained.set_input_embeddings(input_embedding)
@@ -66,50 +76,70 @@ def add_tokens_to_pretrain(pretrained, tokenizer, add_tokens, sample_init=False)
 def load_trained_model(model_path, pretrained_config=None, tag=None):
     """loading saved task"""
 
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     torchpack = torch.load(model_path, map_location=device)
 
-    model_info = {key: torchpack[key] for key in torchpack.keys() if 'state_dict' not in key and 'models' not in key}
+    model_info = {
+        key: torchpack[key]
+        for key in torchpack.keys()
+        if "state_dict" not in key and "models" not in key
+    }
     print("===task info===")
-    [print(k, v[:10], "...") if isinstance(v, list) and len(v) > 10 else print(k, v) for k, v in model_info.items()]
-    print('===============')
+    [
+        print(k, v[:10], "...")
+        if isinstance(v, list) and len(v) > 10 else print(k, v)
+        for k, v in model_info.items()
+    ]
+    print("===============")
 
-    if 'tags' in torchpack and len(torchpack['tags']) > 1:
+    if "tags" in torchpack and len(torchpack["tags"]) > 1:
         if tag is None:
             print("Pick which models to use in multi-task models")
-            inquirer_res = inquirer.prompt(
-                [inquirer.List('tag', message="Select task", choices=torchpack['tags'])])
-            tag = inquirer_res['tag']
-        type_ind = torchpack['tags'].index(tag)
+            inquirer_res = inquirer.prompt([
+                inquirer.List("tag",
+                              message="Select task",
+                              choices=torchpack["tags"])
+            ])
+            tag = inquirer_res["tag"]
+        type_ind = torchpack["tags"].index(tag)
     else:
         type_ind = 0
     print("loading saved task")
 
     # get all loading parameter
-    maxlen = torchpack['maxlen']
+    maxlen = torchpack["maxlen"]
     if pretrained_config is not None:
         config = pretrained_config
     else:
-        config = torchpack['model_config'] if 'model_config' in torchpack else torchpack['bert']
-    model_types = [torchpack['type']] if not isinstance(torchpack['type'], list) else torchpack['type']
-    models_state = torchpack['models'] if 'models' in torchpack else [torchpack['model_state_dict']]
+        config = (torchpack["model_config"]
+                  if "model_config" in torchpack else torchpack["bert"])
+    model_types = ([
+        torchpack["type"]
+    ] if not isinstance(torchpack["type"], list) else torchpack["type"])
+    models_state = (torchpack["models"] if "models" in torchpack else
+                    [torchpack["model_state_dict"]])
     type = model_types[type_ind]
-    add_tokens = torchpack['add_tokens'] if 'add_tokens' in torchpack else None
+    add_tokens = torchpack["add_tokens"] if "add_tokens" in torchpack else None
     # load task
     tokenizer = AutoTokenizer.from_pretrained(config)
     pretrained = AutoModel.from_pretrained(config)
 
-    pretrained, tokenizer = add_tokens_to_pretrain(pretrained, tokenizer, add_tokens)
+    pretrained, tokenizer = add_tokens_to_pretrain(pretrained, tokenizer,
+                                                   add_tokens)
 
     model_class = load_model_class(type)
     task_detail = {}
-    if 'task-label' in torchpack:
-        task_detail = torchpack['task-label']
-    elif 'label' in torchpack:
-        task_detail = {'label': torchpack['label']}
-
-    model = model_class.Model(tokenizer=tokenizer, pretrained=pretrained, tasks_detail=task_detail,
-                              maxlen=maxlen)
+    if "task-label" in torchpack:
+        task_detail = torchpack["task-label"]
+    elif "label" in torchpack:
+        task_detail = {"label": torchpack["label"]}
+
+    model = model_class.Model(
+        tokenizer=tokenizer,
+        pretrained=pretrained,
+        tasks_detail=task_detail,
+        maxlen=maxlen,
+    )
     model.load_state_dict(models_state[type_ind], strict=False)
     model = model.to(device)
 
@@ -119,24 +149,31 @@ def load_trained_model(model_path, pretrained_config=None, tag=None):
     return model, type, model_class, model_info, preprocessor
 
 
-def save_model(models, input_arg, models_tag, epoch, fname, logger, accelerator, add_tokens=None):
+def save_model(models,
+               input_arg,
+               models_tag,
+               epoch,
+               fname,
+               logger,
+               accelerator,
+               add_tokens=None):
     accelerator.wait_for_everyone()
     save_model = {
-        'models': [accelerator.get_state_dict(m) for m in models],
-        'model_config': input_arg.get('config'),
-        'add_tokens': add_tokens,
-        'tags': models_tag,
-        'type': input_arg.get('task'),
-        'maxlen': input_arg.get('maxlen'),
-        'epoch': epoch
+        "models": [accelerator.get_state_dict(m) for m in models],
+        "model_config": input_arg.get("config"),
+        "add_tokens": add_tokens,
+        "tags": models_tag,
+        "type": input_arg.get("task"),
+        "maxlen": input_arg.get("maxlen"),
+        "epoch": epoch,
     }
 
-    for ind, m in enumerate(input_arg.get('task')):
+    for ind, m in enumerate(input_arg.get("task")):
         m = accelerator.unwrap_model(m)
-        if 'tag' in m:
-            save_model['label'] = models[ind].labels
+        if "tag" in m:
+            save_model["label"] = models[ind].labels
         if "clas" in m:
-            save_model['task-label'] = models[ind].tasks_detail
+            save_model["task-label"] = models[ind].tasks_detail
 
     torch.save(save_model, f"{fname}.pt")
     logger.write_log(f"weights were saved to {fname}.pt")
@@ -150,11 +187,11 @@ def tie_encoder_decoder_weights(encoder, decoder, base_model_prefix):
         )
 
     def tie_encoder_to_decoder_recursively(
-            decoder_pointer: nn.Module,
-            encoder_pointer: nn.Module,
-            module_name: str,
-            uninitialized_encoder_weights: List[str],
-            depth=0,
+        decoder_pointer: nn.Module,
+        encoder_pointer: nn.Module,
+        module_name: str,
+        uninitialized_encoder_weights: List[str],
+        depth=0,
     ):
         assert isinstance(decoder_pointer, nn.Module) and isinstance(
             encoder_pointer, nn.Module
@@ -171,18 +208,22 @@ def tie_encoder_to_decoder_recursively(
         decoder_modules = decoder_pointer._modules
         if len(decoder_modules) > 0:
             assert (
-                    len(encoder_modules) > 0
+                len(encoder_modules) > 0
             ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"
 
-            all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
+            all_encoder_weights = set([
+                module_name + "/" + sub_name
+                for sub_name in encoder_modules.keys()
+            ])
             encoder_layer_pos = 0
             for name, module in decoder_modules.items():
                 if name.isdigit():
                     encoder_name = str(int(name) + encoder_layer_pos)
                     decoder_name = name
-                    if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
-                            encoder_modules
-                    ) != len(decoder_modules):
+                    if not isinstance(
+                            decoder_modules[decoder_name],
+                            type(encoder_modules[encoder_name]),
+                    ) and len(encoder_modules) != len(decoder_modules):
                         # this can happen if the name corresponds to the position in a list module list of layers
                         # in this case the decoder has added a cross-attention that the encoder does not have
                         # thus skip this step and subtract one layer pos from encoder
@@ -208,7 +249,8 @@ def tie_encoder_to_decoder_recursively(
             uninitialized_encoder_weights += list(all_encoder_weights)
 
     # tie weights recursively
-    tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights)
+    tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix,
+                                       uninitialized_encoder_weights)
     if len(uninitialized_encoder_weights) > 0:
         print(
             f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}"
diff --git a/tfkit/utility/predictor.py b/tfkit/utility/predictor.py
index f8074e5..401c5be 100644
--- a/tfkit/utility/predictor.py
+++ b/tfkit/utility/predictor.py
@@ -1,21 +1,22 @@
 import json
 from collections import defaultdict
+from math import log
 
 import nlp2
 import numpy as np
 import torch
-from math import log
 from torch.distributions import Categorical
 
 from tfkit.utility import tok
 
 
 class BasePredictor:
+
     def __init__(self, model, preprocessor):
         self.preprocessor = preprocessor
         self.model = model
 
-    def wrap_input(self, input=''):
+    def wrap_input(self, input=""):
         pass
 
     def processing(self, input_args):
@@ -26,7 +27,7 @@ def wrap_output(self, results, detail_dicts):
 
     def predict(self, **input_args):
         """
-        Wrap the input 
+        Wrap the input
         :param input_args:
         :return: [answer list], {result detail dict}
         """
@@ -35,127 +36,228 @@ def predict(self, **input_args):
 
 
 class BaseTextGeneratePredictor(BasePredictor):
-    def wrap_input(self, input='', topK=1, topP=0.85, mode=['greedy', 'topK', 'topP'], decodenum=1,
-                   filtersim=False, reserved_len=0, task=None, handle_exceed='noop', eos_num=1, decode_maxlen=0,
-                   merge_strategy=['prob', 'entropy', 'count'],
-                   no_repeat=False):
+
+    def wrap_input(
+        self,
+        input="",
+        topK=1,
+        topP=0.85,
+        mode=["greedy", "topK", "topP"],
+        decodenum=1,
+        filtersim=False,
+        reserved_len=0,
+        task=None,
+        handle_exceed="noop",
+        eos_num=1,
+        decode_maxlen=0,
+        merge_strategy=["prob", "entropy", "count"],
+        no_repeat=False,
+    ):
         return {
-            'filtersim': json.loads(str(filtersim).lower()),
-            'topK': int(topK),
-            'no_repeat': json.loads(str(no_repeat).lower()),
-            'eos_num': int(eos_num),
-            'topP': float(topP),
-            'decodenum': int(decodenum),
-            'decode_maxlen': int(decode_maxlen),
-            'reserved_len': int(reserved_len),
-            'mode': mode[0] if isinstance(mode, list) else mode.lower(),
-            'merge_strategy': merge_strategy[0] if isinstance(merge_strategy, list) else merge_strategy.lower(),
-            'handle_exceed': handle_exceed,
-            'task': task,
-            'input': input
+            "filtersim":
+            json.loads(str(filtersim).lower()),
+            "topK":
+            int(topK),
+            "no_repeat":
+            json.loads(str(no_repeat).lower()),
+            "eos_num":
+            int(eos_num),
+            "topP":
+            float(topP),
+            "decodenum":
+            int(decodenum),
+            "decode_maxlen":
+            int(decode_maxlen),
+            "reserved_len":
+            int(reserved_len),
+            "mode":
+            mode[0] if isinstance(mode, list) else mode.lower(),
+            "merge_strategy":
+            merge_strategy[0]
+            if isinstance(merge_strategy, list) else merge_strategy.lower(),
+            "handle_exceed":
+            handle_exceed,
+            "task":
+            task,
+            "input":
+            input,
         }
 
 
 class ClassificationPredictor(BasePredictor):
-    def wrap_input(self, input='', task='', topK=1, mode=['greedy', 'topK', 'topP'], handle_exceed='noop',
-                   merge_strategy=['prob', 'entropy', 'count'],
-                   reserved_len=0):
+
+    def wrap_input(
+        self,
+        input="",
+        task="",
+        topK=1,
+        mode=["greedy", "topK", "topP"],
+        handle_exceed="noop",
+        merge_strategy=["prob", "entropy", "count"],
+        reserved_len=0,
+    ):
         return {
-            'topK': int(topK),
-            'mode': mode[0] if isinstance(mode, list) else mode.lower(),
-            'task': task if len(task) > 0 else list(self.model.tasks_detail.keys())[0],
-            'input': input,
-            'reserved_len': int(reserved_len),
-            'merge_strategy': merge_strategy[0] if isinstance(merge_strategy, list) else merge_strategy.lower(),
-            'handle_exceed': handle_exceed,
+            "topK":
+            int(topK),
+            "mode":
+            mode[0] if isinstance(mode, list) else mode.lower(),
+            "task":
+            task if len(task) > 0 else list(self.model.tasks_detail.keys())[0],
+            "input":
+            input,
+            "reserved_len":
+            int(reserved_len),
+            "merge_strategy":
+            merge_strategy[0]
+            if isinstance(merge_strategy, list) else merge_strategy.lower(),
+            "handle_exceed":
+            handle_exceed,
         }
 
     def processing(self, input_args):
         with torch.no_grad():
             ret_result = []
             ret_detail = []
-            proc = self.preprocessor(self.model.tokenizer, maxlen=self.model.maxlen,
-                                     handle_exceed=input_args['handle_exceed'],
-                                     reserved_len=input_args['reserved_len'])
-
-            for items in proc.preprocess({"task": input_args['task'], "input": input_args['input']}):
-                feature = proc.postprocess({**items, **{'task_dict': self.model.tasks}}, self.model.tokenizer,
-                                           self.model.maxlen)
+            proc = self.preprocessor(
+                self.model.tokenizer,
+                maxlen=self.model.maxlen,
+                handle_exceed=input_args["handle_exceed"],
+                reserved_len=input_args["reserved_len"],
+            )
+
+            for items in proc.preprocess({
+                    "task": input_args["task"],
+                    "input": input_args["input"]
+            }):
+                feature = proc.postprocess(
+                    {
+                        **items,
+                        **{
+                            "task_dict": self.model.tasks
+                        }
+                    },
+                    self.model.tokenizer,
+                    self.model.maxlen,
+                )
             predictions = self.model.forward(feature, eval=True)
-            if input_args['topK'] < 2:
-                ret_result.append(
-                    [i[input_args['task']] for i in predictions['max_item'] if input_args['task'] in i][0])
+            if input_args["topK"] < 2:
+                ret_result.append([
+                    i[input_args["task"]] for i in predictions["max_item"]
+                    if input_args["task"] in i
+                ][0])
                 ret_detail.append(predictions)
             else:
-                task_map = [i[input_args['task']] for i in predictions['label_prob'] if input_args['task'] in i][0]
-                ret_result.append(sorted(task_map, key=task_map.get, reverse=True)[:input_args['topK']])
+                task_map = [
+                    i[input_args["task"]] for i in predictions["label_prob"]
+                    if input_args["task"] in i
+                ][0]
+                ret_result.append(
+                    sorted(task_map, key=task_map.get,
+                           reverse=True)[:input_args["topK"]])
                 ret_detail.append(predictions)
 
         # apply different strategy to merge result after sliding windows
-        if input_args['merge_strategy'] == 'count':
+        if input_args["merge_strategy"] == "count":
             ret_result = max(ret_result, key=ret_result.count)
         else:
             results_prob = []
             results_entropy = []
             for detail in ret_detail:
-                prob_map = detail['label_prob'][0][input_args['task']]
+                prob_map = detail["label_prob"][0][input_args["task"]]
                 result_value = [v for _, v in prob_map.items()]
-                results_entropy.append(Categorical(probs=torch.tensor(result_value)).entropy().data.tolist())
+                results_entropy.append(
+                    Categorical(probs=torch.tensor(
+                        result_value)).entropy().data.tolist())
                 results_prob.append(max(result_value))
             min_entropy_index = results_entropy.index(min(results_entropy))
             max_prob_index = results_prob.index(max(results_prob))
-            if input_args['merge_strategy'] == 'entropy':
+            if input_args["merge_strategy"] == "entropy":
                 ret_result = ret_result[min_entropy_index]
-            if input_args['merge_strategy'] == 'prob':
+            if input_args["merge_strategy"] == "prob":
                 ret_result = ret_result[max_prob_index]
         return ret_result, ret_detail
 
 
 class QuestionAnsweringPredictor(BasePredictor):
-    def wrap_input(self, input='', task='', topK=1, mode=['greedy', 'topK', 'topP'], handle_exceed='slide',
-                   merge_strategy=['prob', 'entropy', 'count'],
-                   reserved_len=0):
+
+    def wrap_input(
+        self,
+        input="",
+        task="",
+        topK=1,
+        mode=["greedy", "topK", "topP"],
+        handle_exceed="slide",
+        merge_strategy=["prob", "entropy", "count"],
+        reserved_len=0,
+    ):
         return {
-            'topK': int(topK),
-            'mode': mode[0] if isinstance(mode, list) else mode.lower(),
-            'input': input,
-            'task': task,
-            'reserved_len': int(reserved_len),
-            'merge_strategy': merge_strategy[0] if isinstance(merge_strategy, list) else merge_strategy.lower(),
-            'handle_exceed': handle_exceed,
+            "topK":
+            int(topK),
+            "mode":
+            mode[0] if isinstance(mode, list) else mode.lower(),
+            "input":
+            input,
+            "task":
+            task,
+            "reserved_len":
+            int(reserved_len),
+            "merge_strategy":
+            merge_strategy[0]
+            if isinstance(merge_strategy, list) else merge_strategy.lower(),
+            "handle_exceed":
+            handle_exceed,
         }
 
     def processing(self, input_args):
         with torch.no_grad():
             ret_result = []
             ret_detail = []
-            proc = self.preprocessor(self.model.tokenizer, maxlen=self.model.maxlen,
-                                     handle_exceed=input_args['handle_exceed'],
-                                     reserved_len=input_args['reserved_len'])
-
-            for items in proc.preprocess({"task": input_args['task'], "input": input_args['input']}):
-                raw_input = items['raw_input']
-                feature = proc.postprocess(items, self.model.tokenizer, self.model.maxlen)
+            proc = self.preprocessor(
+                self.model.tokenizer,
+                maxlen=self.model.maxlen,
+                handle_exceed=input_args["handle_exceed"],
+                reserved_len=input_args["reserved_len"],
+            )
+
+            for items in proc.preprocess({
+                    "task": input_args["task"],
+                    "input": input_args["input"]
+            }):
+                raw_input = items["raw_input"]
+                feature = proc.postprocess(items, self.model.tokenizer,
+                                           self.model.maxlen)
             for k, v in feature.items():
                 feature[k] = [v]
 
             result = self.model.forward(feature, eval=True)
-            start_dict = [i['start'] for i in result['label_prob_all'] if 'start' in i][0]
-            end_dict = [i['end'] for i in result['label_prob_all'] if 'end' in i][0]
+            start_dict = [
+                i["start"] for i in result["label_prob_all"] if "start" in i
+            ][0]
+            end_dict = [
+                i["end"] for i in result["label_prob_all"] if "end" in i
+            ][0]
 
             answers = []
-            sorted_start = sorted(start_dict.items(), key=lambda item: item[1], reverse=True)[:50]
-            sorted_end = sorted(end_dict.items(), key=lambda item: item[1], reverse=True)[:50]
+            sorted_start = sorted(start_dict.items(),
+                                  key=lambda item: item[1],
+                                  reverse=True)[:50]
+            sorted_end = sorted(end_dict.items(),
+                                key=lambda item: item[1],
+                                reverse=True)[:50]
             for start_index, start_prob in sorted_start:
                 for end_index, end_prob in sorted_end:
                     if start_index > end_index:
                         continue
-                    answers.append((start_index, end_index, start_prob + end_prob))
-            answer_results = sorted(answers, key=lambda answers: answers[2],
-                                    reverse=True)[:input_args['topK']]
-            ret_result.append(
-                ["".join(self.model.tokenizer.convert_tokens_to_string(raw_input[ans[0]:ans[1]])) for ans in
-                 answer_results])
+                    answers.append(
+                        (start_index, end_index, start_prob + end_prob))
+            answer_results = sorted(answers,
+                                    key=lambda answers: answers[2],
+                                    reverse=True)[:input_args["topK"]]
+            ret_result.append([
+                "".join(
+                    self.model.tokenizer.convert_tokens_to_string(
+                        raw_input[ans[0]:ans[1]])) for ans in answer_results
+            ])
             ret_detail.append(result)
 
             # apply different strategy to merge result after sliding windows
@@ -169,98 +271,151 @@ def processing(self, input_args):
             results_prob = []
             results_entropy = []
             for detail in non_empty_detail:
-                prob_start = detail['label_prob_all'][0]['start'][int(detail['label_map'][0]['start'])]
-                prob_end = detail['label_prob_all'][0]['end'][int(detail['label_map'][0]['end'])]
-                prob_sum = [sum(x) for x in zip(list(detail['label_prob_all'][0]['start'].values()),
-                                                list(detail['label_prob_all'][0]['end'].values()))]
-                results_entropy.append(Categorical(probs=torch.tensor(prob_sum)).entropy().data.tolist())
+                prob_start = detail["label_prob_all"][0]["start"][int(
+                    detail["label_map"][0]["start"])]
+                prob_end = detail["label_prob_all"][0]["end"][int(
+                    detail["label_map"][0]["end"])]
+                prob_sum = [
+                    sum(x) for x in zip(
+                        list(detail["label_prob_all"][0]["start"].values()),
+                        list(detail["label_prob_all"][0]["end"].values()),
+                    )
+                ]
+                results_entropy.append(
+                    Categorical(
+                        probs=torch.tensor(prob_sum)).entropy().data.tolist())
                 results_prob.append(-np.log(prob_start) + -np.log(prob_end))
         return non_empty_result, non_empty_detail
 
 
 class TaggingPredictor(BasePredictor):
-    def wrap_input(self, input='', task='', topK=1, mode=['greedy', 'topK', 'topP'], handle_exceed='slide',
-                   merge_strategy=['prob', 'entropy', 'count'],
-                   neg="O",
-                   start_contain="B",
-                   end_contain="I",
-                   minlen=1,
-                   reserved_len=0):
+
+    def wrap_input(
+        self,
+        input="",
+        task="",
+        topK=1,
+        mode=["greedy", "topK", "topP"],
+        handle_exceed="slide",
+        merge_strategy=["prob", "entropy", "count"],
+        neg="O",
+        start_contain="B",
+        end_contain="I",
+        minlen=1,
+        reserved_len=0,
+    ):
         return {
-            'topK': int(topK),
-            'mode': mode[0] if isinstance(mode, list) else mode.lower(),
-            'input': input,
-            'task': task,
-            'neg': neg,
-            'minlen': minlen,
-            'start_contain': start_contain,
-            'end_contain': end_contain,
-            'reserved_len': int(reserved_len),
-            'merge_strategy': merge_strategy[0] if isinstance(merge_strategy, list) else merge_strategy.lower(),
-            'handle_exceed': handle_exceed,
+            "topK":
+            int(topK),
+            "mode":
+            mode[0] if isinstance(mode, list) else mode.lower(),
+            "input":
+            input,
+            "task":
+            task,
+            "neg":
+            neg,
+            "minlen":
+            minlen,
+            "start_contain":
+            start_contain,
+            "end_contain":
+            end_contain,
+            "reserved_len":
+            int(reserved_len),
+            "merge_strategy":
+            merge_strategy[0]
+            if isinstance(merge_strategy, list) else merge_strategy.lower(),
+            "handle_exceed":
+            handle_exceed,
         }
 
     def processing(self, input_args):
         with torch.no_grad():
             ret_detail = []
             predicted_pos_prob = defaultdict(lambda: defaultdict(list))
-            proc = self.preprocessor(self.model.tokenizer, maxlen=self.model.maxlen,
-                                     handle_exceed=input_args['handle_exceed'],
-                                     reserved_len=input_args['reserved_len'])
-            for items in proc.preprocess({"task": input_args['task'], "input": input_args['input']}):
-                token_word_mapping = items['token_word_mapping']
-                feature = proc.postprocess({**items, **{'task_dict': self.model.labels}}, self.model.tokenizer,
-                                           self.model.maxlen)
+            proc = self.preprocessor(
+                self.model.tokenizer,
+                maxlen=self.model.maxlen,
+                handle_exceed=input_args["handle_exceed"],
+                reserved_len=input_args["reserved_len"],
+            )
+            for items in proc.preprocess({
+                    "task": input_args["task"],
+                    "input": input_args["input"]
+            }):
+                token_word_mapping = items["token_word_mapping"]
+                feature = proc.postprocess(
+                    {
+                        **items,
+                        **{
+                            "task_dict": self.model.labels
+                        }
+                    },
+                    self.model.tokenizer,
+                    self.model.maxlen,
+                )
                 for k, v in feature.items():
                     feature[k] = [v]
-                feature['token_word_mapping'] = token_word_mapping
+                feature["token_word_mapping"] = token_word_mapping
                 result = self.model.forward(feature, eval=True)
-                for token_pred, token_map in zip(result['label_prob_all'], result['token_word_mapping']):
+                for token_pred, token_map in zip(result["label_prob_all"],
+                                                 result["token_word_mapping"]):
                     token_prob = list(token_pred.values())[0]
                     max_label = max(token_prob, key=token_prob.get)
                     max_prob = token_prob[max_label]
-                    max_entropy = Categorical(probs=torch.tensor(list(token_prob.values()))).entropy().data.tolist()
-                    predicted_pos_prob[token_map['pos']]['char'] = token_map['word']
-                    predicted_pos_prob[token_map['pos']]['labels'].append(max_label)
-                    predicted_pos_prob[token_map['pos']]['prob'].append(max_prob)
-                predicted_pos_prob[token_map['pos']]['entropy'].append(max_entropy)
+                    max_entropy = (Categorical(probs=torch.tensor(
+                        list(token_prob.values()))).entropy().data.tolist())
+                    predicted_pos_prob[
+                        token_map["pos"]]["char"] = token_map["word"]
+                    predicted_pos_prob[token_map["pos"]]["labels"].append(
+                        max_label)
+                    predicted_pos_prob[token_map["pos"]]["prob"].append(
+                        max_prob)
+                predicted_pos_prob[token_map["pos"]]["entropy"].append(
+                    max_entropy)
 
             ret_detail.append(result)
 
             ret_result = []
             for key, value in predicted_pos_prob.items():
-                if input_args['merge_strategy'] == 'count':
-                    label = max(value['labels'], key=value['labels'].count)
-                if input_args['merge_strategy'] == 'entropy':
-                    min_entropy_index = value['entropy'].index(min(value['entropy']))
-                    label = value['labels'][min_entropy_index]
-                if input_args['merge_strategy'] == 'prob':
-                    max_prob_index = value['prob'].index(max(value['prob']))
-                    label = value['labels'][max_prob_index]
-                ret_result.append({value['char']: label})
+                if input_args["merge_strategy"] == "count":
+                    label = max(value["labels"], key=value["labels"].count)
+                if input_args["merge_strategy"] == "entropy":
+                    min_entropy_index = value["entropy"].index(
+                        min(value["entropy"]))
+                    label = value["labels"][min_entropy_index]
+                if input_args["merge_strategy"] == "prob":
+                    max_prob_index = value["prob"].index(max(value["prob"]))
+                    label = value["labels"][max_prob_index]
+                ret_result.append({value["char"]: label})
 
             output = []
             target_str = ["", ""]
             after_start = False
             for mapping in ret_result:
                 for k, y in mapping.items():
-                    if input_args['start_contain'] in y:
+                    if input_args["start_contain"] in y:
                         after_start = True
                         if len(target_str[0]) > 0:
-                            if len(target_str[0]) > input_args['minlen']:
+                            if len(target_str[0]) > input_args["minlen"]:
                                 output.append(target_str)
                             target_str = ["", ""]
                         target_str[0] += k
                         target_str[1] = y
-                    elif y is not input_args['neg'] and after_start:
+                    elif y is not input_args["neg"] and after_start:
                         target_str[0] += k
                         target_str[1] = y
                     else:
                         after_start = False
-            if len(target_str[0]) > input_args['minlen'] and target_str not in output:
+            if len(target_str[0]
+                   ) > input_args["minlen"] and target_str not in output:
                 output.append(target_str)
-            output = [[ner, tag.replace(input_args['start_contain'], "").replace(input_args['end_contain'], "")] for
-                      ner, tag in output]
+            output = [[
+                ner,
+                tag.replace(input_args["start_contain"],
+                            "").replace(input_args["end_contain"], ""),
+            ] for ner, tag in output]
 
         return output, ret_detail
 
@@ -269,25 +424,37 @@ class NonAutoRegressivePredictor(BaseTextGeneratePredictor):
 
     def processing(self, input_args):
         with torch.no_grad():
-            proc = self.preprocessor(self.model.tokenizer, maxlen=self.model.maxlen,
-                                     handle_exceed=input_args['handle_exceed'],
-                                     reserved_len=input_args['reserved_len'])
-            for items in proc.preprocess({"task": input_args['task'], "input": input_args['input']}):
-                feature = proc.postprocess(items, self.model.tokenizer, self.model.maxlen)
+            proc = self.preprocessor(
+                self.model.tokenizer,
+                maxlen=self.model.maxlen,
+                handle_exceed=input_args["handle_exceed"],
+                reserved_len=input_args["reserved_len"],
+            )
+            for items in proc.preprocess({
+                    "task": input_args["task"],
+                    "input": input_args["input"]
+            }):
+                feature = proc.postprocess(items, self.model.tokenizer,
+                                           self.model.maxlen)
             for k, v in feature.items():
                 feature[k] = [v]
-            predictions = self.model.forward(feature, eval=True, beamsearch=input_args['decodenum'] > 1,
-                                             max_return=input_args['decodenum'])
-            return predictions['max_item'], predictions
+            predictions = self.model.forward(
+                feature,
+                eval=True,
+                beamsearch=input_args["decodenum"] > 1,
+                max_return=input_args["decodenum"],
+            )
+            return predictions["max_item"], predictions
 
 
 class AutoRegressivePredictor(BaseTextGeneratePredictor):
 
     def processing(self, input_args):
         previous = []
-        if tok.UNIVERSAL_SEP in input_args['input']:
-            previous = self.model.tokenizer.tokenize(input_args['input'].split(tok.UNIVERSAL_SEP)[-1])
-            input_args['eos_num'] += 1
+        if tok.UNIVERSAL_SEP in input_args["input"]:
+            previous = self.model.tokenizer.tokenize(input_args["input"].split(
+                tok.UNIVERSAL_SEP)[-1])
+            input_args["eos_num"] += 1
         sep_tok = tok.tok_sep(self.model.tokenizer)
         sequences = [[[], 1.0]]
         with torch.no_grad():
@@ -295,53 +462,89 @@ def processing(self, input_args):
                 all_candidates = list()
                 exceed = False
                 for seq in sequences:
-                    if sep_tok not in seq[0] or seq[0].count(sep_tok) < input_args['eos_num']:
+                    if (sep_tok not in seq[0]
+                            or seq[0].count(sep_tok) < input_args["eos_num"]):
                         tokens, score = seq
                         if not tokens:
                             tokens = previous
 
-                        proc = self.preprocessor(self.model.tokenizer, maxlen=self.model.maxlen,
-                                                 handle_exceed=input_args['handle_exceed'],
-                                                 reserved_len=input_args['reserved_len'])
-                        for items in proc.preprocess(
-                                {"task": input_args['task'], "input": input_args['input'], 'previous': tokens}):
+                        proc = self.preprocessor(
+                            self.model.tokenizer,
+                            maxlen=self.model.maxlen,
+                            handle_exceed=input_args["handle_exceed"],
+                            reserved_len=input_args["reserved_len"],
+                        )
+                        for items in proc.preprocess({
+                                "task":
+                                input_args["task"],
+                                "input":
+                                input_args["input"],
+                                "previous":
+                                tokens,
+                        }):
                             feature_dict = proc.postprocess_batch(
-                                proc.postprocess(items, self.model.tokenizer, self.model.maxlen))
+                                proc.postprocess(items, self.model.tokenizer,
+                                                 self.model.maxlen))
 
                         # check input exceed
-                        if len(tokens) >= self.model.maxlen \
-                                or ('prev' in feature_dict and len(feature_dict['prev']) >= self.model.maxlen) \
-                                or ('input' in feature_dict and len(feature_dict['input']) >= self.model.maxlen):
+                        if (len(tokens) >= self.model.maxlen or
+                            ("prev" in feature_dict and
+                             len(feature_dict["prev"]) >= self.model.maxlen) or
+                            ("input" in feature_dict and
+                             len(feature_dict["input"]) >= self.model.maxlen)):
                             exceed = True
                             all_candidates.append(seq)
                             continue
-                        predictions = self.model.forward(feature_dict, eval=True, use_prev=False,
-                                                         output_hidden_states=True,
-                                                         beamsearch=input_args['decodenum'] > 1,
-                                                         max_return=max(input_args['decodenum'],
-                                                                        input_args['topK']))
-                        label_prob = predictions['label_prob'] if 'label_prob' in predictions else \
-                            [predictions['max_item']]
+                        predictions = self.model.forward(
+                            feature_dict,
+                            eval=True,
+                            use_prev=False,
+                            output_hidden_states=True,
+                            beamsearch=input_args["decodenum"] > 1,
+                            max_return=max(input_args["decodenum"],
+                                           input_args["topK"]),
+                        )
+                        label_prob = (predictions["label_prob"]
+                                      if "label_prob" in predictions else
+                                      [predictions["max_item"]])
                         # topK topP
-                        if 'top' in input_args['mode']:
-                            prob_list = predictions['prob_list']
-                            if 'topk' in input_args['mode']:
-                                sample_list = prob_list[:input_args['topK']]
-                                decode_range = max(input_args['decodenum'], input_args['topK'])
-                                prob_norm = [float(i) / sum(sample_list) for i in sample_list]
-                                choice_list = np.random.choice(sample_list, p=prob_norm,
-                                                               size=decode_range,
-                                                               replace=False)
+                        if "top" in input_args["mode"]:
+                            prob_list = predictions["prob_list"]
+                            if "topk" in input_args["mode"]:
+                                sample_list = prob_list[:input_args["topK"]]
+                                decode_range = max(input_args["decodenum"],
+                                                   input_args["topK"])
+                                prob_norm = [
+                                    float(i) / sum(sample_list)
+                                    for i in sample_list
+                                ]
+                                choice_list = np.random.choice(
+                                    sample_list,
+                                    p=prob_norm,
+                                    size=decode_range,
+                                    replace=False,
+                                )
                             else:
                                 topP_list = np.cumsum(prob_list)
-                                index_overP = [i for i, x in enumerate(topP_list) if x > input_args['topP']]
-                                index_overP = 0 if len(index_overP) < 1 else index_overP[0]
+                                index_overP = [
+                                    i for i, x in enumerate(topP_list)
+                                    if x > input_args["topP"]
+                                ]
+                                index_overP = (0 if len(index_overP) < 1 else
+                                               index_overP[0])
                                 sample_list = prob_list[:index_overP + 1]
-                                prob_norm = [float(i) / sum(sample_list) for i in sample_list]
-                                choice_list = np.random.choice(sample_list, p=prob_norm,
-                                                               size=input_args['decodenum'])
-                            for idx in range(input_args['decodenum']):
-                                sampling_index = prob_list.index(choice_list[idx])
+                                prob_norm = [
+                                    float(i) / sum(sample_list)
+                                    for i in sample_list
+                                ]
+                                choice_list = np.random.choice(
+                                    sample_list,
+                                    p=prob_norm,
+                                    size=input_args["decodenum"],
+                                )
+                            for idx in range(input_args["decodenum"]):
+                                sampling_index = prob_list.index(
+                                    choice_list[idx])
                                 k, v = label_prob[sampling_index]
                                 candidate = [tokens + [k], score + -log(v)]
                                 all_candidates.append(candidate)
@@ -349,23 +552,29 @@ def processing(self, input_args):
                         # greedy / beam search
                         else:
                             for k, v in label_prob[:50]:
-                                if (input_args['no_repeat'] and len(tokens) > 0 and tokens[-1] == k) or len(k) < 1:
+                                if (input_args["no_repeat"] and len(tokens) > 0
+                                        and tokens[-1] == k) or len(k) < 1:
                                     continue
-                                candidate = [tokens + [k], score + -log(v) if v > 0 else 0]
+                                candidate = [
+                                    tokens + [k],
+                                    score + -log(v) if v > 0 else 0,
+                                ]
                                 all_candidates.append(candidate)
                     else:
                         all_candidates.append(seq)
 
                 ordered = sorted(all_candidates, key=lambda tup: tup[1])
-                if input_args['filtersim']:  # default disable, it will took a lots of time on long sequence generation
-                    nlp2.filter_jaccard_similar_text_from_list(ordered, input_args['decodenum'])
-                sequences = ordered[:input_args['decodenum']]
+                # default disable, it will took a lots of time on long sequence generation
+                if input_args["filtersim"]:
+                    nlp2.filter_jaccard_similar_text_from_list(
+                        ordered, input_args["decodenum"])
+                sequences = ordered[:input_args["decodenum"]]
                 stop = 0
                 for i in sequences:
                     # i[0], i[1] - sequence, sequence score
-                    if (i[0].count(sep_tok) >= input_args['eos_num']) or \
-                            len(i[0]) > self.model.maxlen or \
-                            0 < input_args['decode_maxlen'] < len(i[0]):
+                    if ((i[0].count(sep_tok) >= input_args["eos_num"])
+                            or len(i[0]) > self.model.maxlen
+                            or 0 < input_args["decode_maxlen"] < len(i[0])):
                         stop += 1
                 if stop == len(sequences) or exceed:
                     break
@@ -375,7 +584,8 @@ def processing(self, input_args):
                     sequences[i][0] = sequences[i][0][:-1]
                 slide_len = len(previous) if len(previous) > 0 else 0
                 sequences[i][0] = self.model.tokenizer.decode(
-                    self.model.tokenizer.convert_tokens_to_ids(sequences[i][0][slide_len:]))
+                    self.model.tokenizer.convert_tokens_to_ids(
+                        sequences[i][0][slide_len:]))
 
             self.model.clean_cache()
-            return [i[0] for i in sequences], {'label_map': sequences}
+            return [i[0] for i in sequences], {"label_map": sequences}
diff --git a/tfkit/utility/preprocess.py b/tfkit/utility/preprocess.py
index 1b583eb..a519ad4 100644
--- a/tfkit/utility/preprocess.py
+++ b/tfkit/utility/preprocess.py
@@ -15,13 +15,27 @@ class GeneralNLPPreprocessor:
     it will call get feature from data when it's converting to model input
     """
 
-    def __init__(self, tokenizer, maxlen=512, handle_exceed='slide', reserved_len=0, uint16_save=False,
-                 kwargs={}):
+    def __init__(
+        self,
+        tokenizer,
+        maxlen=512,
+        handle_exceed="slide",
+        reserved_len=0,
+        uint16_save=False,
+        kwargs={},
+    ):
         self.tokenizer = tokenizer
         self.uint16_save = uint16_save
-        self.parameters = {**{'tokenizer': tokenizer, 'maxlen': maxlen, 'handle_exceed': handle_exceed,
-                              'reserved_len': reserved_len}, **kwargs}
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.parameters = {
+            **{
+                "tokenizer": tokenizer,
+                "maxlen": maxlen,
+                "handle_exceed": handle_exceed,
+                "reserved_len": reserved_len,
+            },
+            **kwargs,
+        }
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         # item = {key: value.tolist() for key, value in item.items()}
         self.tok_pad_id = tok.tok_pad_id(tokenizer)
         self.tok_bos_id = tok.tok_begin_id(tokenizer)
@@ -29,7 +43,7 @@ def __init__(self, tokenizer, maxlen=512, handle_exceed='slide', reserved_len=0,
         self.tok_mask_id = tok.tok_mask_id(tokenizer)
 
     def read_file_to_data(self, filepath):
-        assert 'plz override this funciton'
+        assert "plz override this funciton"
 
     def set_global_parameters(self):
         self.tokenize_target = False
@@ -39,67 +53,101 @@ def preprocess(self, item):
         preprocessed_data = []
         item = self.preprocess_component_prepare_input(item)
         # target may be none in eval
-        t_input_list, t_target_list, t_input_index, t_target_index = self.preprocess_component_split_into_list(
-            item['input'],
-            item.get('target'))
-        for t_input, t_target, t_input_index, t_target_index in zip(t_input_list,
-                                                                    t_target_list,
-                                                                    t_input_index,
-                                                                    t_target_index):
-            slice_length = self.parameters['maxlen'] - self.parameters.get('reserved_len') - 3
-            item['input'] = [tok.tok_begin(self.tokenizer)] + t_input[:slice_length]
-            item['input_index'] = t_input_index
-            item['target_index'] = t_target_index
+        (
+            t_input_list,
+            t_target_list,
+            t_input_index,
+            t_target_index,
+        ) = self.preprocess_component_split_into_list(item["input"],
+                                                      item.get("target"))
+        for t_input, t_target, t_input_index, t_target_index in zip(
+                t_input_list, t_target_list, t_input_index, t_target_index):
+            slice_length = (self.parameters["maxlen"] -
+                            self.parameters.get("reserved_len") - 3)
+            item["input"] = [tok.tok_begin(self.tokenizer)
+                             ] + t_input[:slice_length]
+            item["input_index"] = t_input_index
+            item["target_index"] = t_target_index
             if len(t_target) > 0:
-                item['target'] = t_target
-            for convert_feature_input_dict in self.preprocess_component_convert_to_id(item):
+                item["target"] = t_target
+            for convert_feature_input_dict in self.preprocess_component_convert_to_id(
+                    item):
                 if self.uint16_save:
-                    data_item = {k: np.array(v, dtype=uint16) if isinstance(v, list) else v for k, v in
-                                 convert_feature_input_dict.items()}
+                    data_item = {
+                        k:
+                        np.array(v, dtype=uint16) if isinstance(v, list) else v
+                        for k, v in convert_feature_input_dict.items()
+                    }
                 else:
                     data_item = convert_feature_input_dict
                 preprocessed_data.append(data_item)
         return preprocessed_data
 
     def preprocess_component_prepare_input(self, item):
-        if tok.UNIVERSAL_SEP in item['input']:
-            part = item['input'].split(tok.UNIVERSAL_SEP)
-            item['previous'] = self.tokenizer.tokenize(part[-1])
-            item['input'] = "".join(part[:-1])
+        if tok.UNIVERSAL_SEP in item["input"]:
+            part = item["input"].split(tok.UNIVERSAL_SEP)
+            item["previous"] = self.tokenizer.tokenize(part[-1])
+            item["input"] = "".join(part[:-1])
         return item
 
-    def preprocess_component_split_into_list(self, input_text, target_text=None):
-        t_input_list, t_input_index = tok.handle_exceed(self.tokenizer, input_text,
-                                                        maxlen=self.parameters['maxlen'] - 3,
-                                                        mode=self.parameters.get('handle_exceed'))
+    def preprocess_component_split_into_list(self,
+                                             input_text,
+                                             target_text=None):
+        t_input_list, t_input_index = tok.handle_exceed(
+            self.tokenizer,
+            input_text,
+            maxlen=self.parameters["maxlen"] - 3,
+            mode=self.parameters.get("handle_exceed"),
+        )
         if self.tokenize_target and target_text:
-            t_target_list, t_target_index = tok.handle_exceed(self.tokenizer, target_text,
-                                                              maxlen=self.parameters['maxlen'] - 3,
-                                                              mode=self.parameters.get('handle_exceed'))
+            t_target_list, t_target_index = tok.handle_exceed(
+                self.tokenizer,
+                target_text,
+                maxlen=self.parameters["maxlen"] - 3,
+                mode=self.parameters.get("handle_exceed"),
+            )
         elif target_text:
-            t_target_list, t_target_index = [target_text * len(t_input_list)], [[0] * len(t_input_list)]
+            t_target_list, t_target_index = [target_text * len(t_input_list)
+                                             ], [[0] * len(t_input_list)]
         else:
-            t_target_list, t_target_index = ['' * len(t_input_list)], [[0] * len(t_input_list)]
+            t_target_list, t_target_index = ["" * len(t_input_list)
+                                             ], [[0] * len(t_input_list)]
         return t_input_list, t_target_list, t_input_index, t_target_index
 
     def preprocess_component_convert_to_id(self, item):
-        yield {k: self.tokenizer.convert_tokens_to_ids(v) if isinstance(v, list) else v for k, v in item.items()}
+        yield {
+            k: self.tokenizer.convert_tokens_to_ids(v)
+            if isinstance(v, list) else v
+            for k, v in item.items()
+        }
 
     def postprocess(self, item, tokenizer, maxlen, **kwargs):
-        return {key: torch.tensor(value).to(self.device) for key, value in item.items() if isinstance(value, list)}
+        return {
+            key: torch.tensor(value).to(self.device)
+            for key, value in item.items() if isinstance(value, list)
+        }
 
     def postprocess_batch(self, feature_dict, **kwargs):
-        return {key: torch.unsqueeze(torch.tensor(value), 0).to(self.device) for key, value in feature_dict.items()}
+        return {
+            key: torch.unsqueeze(torch.tensor(value), 0).to(self.device)
+            for key, value in feature_dict.items()
+        }
 
 
 class GeneralCVPreprocessor:
+
     def __init__(self, feature_extractor, kwargs={}):
         self.feature_extractor = feature_extractor
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.parameters = {**{'feature_extractor': feature_extractor}, **kwargs}
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.parameters = {
+            **{
+                "feature_extractor": feature_extractor
+            },
+            **kwargs
+        }
 
     def read_file_to_data(self, filepath):
-        assert 'plz override this funciton'
+        assert "plz override this funciton"
 
     def preprocess(self, item):
         preprocessed_data = []
@@ -107,18 +155,27 @@ def preprocess(self, item):
         return preprocessed_data
 
     def postprocess(self, item, **kwargs):
-        item['input'] = self.feature_extractor(item['input'])
-        return {key: torch.tensor(value).to(self.device) for key, value in item.items()}
+        item["input"] = self.feature_extractor(item["input"])
+        return {
+            key: torch.tensor(value).to(self.device)
+            for key, value in item.items()
+        }
 
 
 class GeneralSpeechPreprocessor:
+
     def __init__(self, feature_extractor, kwargs={}):
         self.feature_extractor = feature_extractor
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.parameters = {**{'feature_extractor': feature_extractor}, **kwargs}
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.parameters = {
+            **{
+                "feature_extractor": feature_extractor
+            },
+            **kwargs
+        }
 
     def read_file_to_data(self, filepath):
-        assert 'plz override this funciton'
+        assert "plz override this funciton"
 
     def preprocess(self, item):
         preprocessed_data = []
@@ -126,5 +183,8 @@ def preprocess(self, item):
         return preprocessed_data
 
     def postprocess(self, item, **kwargs):
-        item['input'] = self.feature_extractor(item['input'])
-        return {key: torch.tensor(value).to(self.device) for key, value in item.items()}
+        item["input"] = self.feature_extractor(item["input"])
+        return {
+            key: torch.tensor(value).to(self.device)
+            for key, value in item.items()
+        }
diff --git a/tfkit/utility/tok.py b/tfkit/utility/tok.py
index ad4942f..546ab8c 100644
--- a/tfkit/utility/tok.py
+++ b/tfkit/utility/tok.py
@@ -8,11 +8,11 @@
 
 
 def tok_begin(tokenizer):
-    if tokenizer.special_tokens_map.get('bos_token'):
-        return tokenizer.special_tokens_map.get('bos_token')
-    elif tokenizer.special_tokens_map.get('cls_token'):
-        tokenizer.special_tokens_map.get('cls_token')
-    return 'cls'
+    if tokenizer.special_tokens_map.get("bos_token"):
+        return tokenizer.special_tokens_map.get("bos_token")
+    elif tokenizer.special_tokens_map.get("cls_token"):
+        tokenizer.special_tokens_map.get("cls_token")
+    return "cls"
 
 
 def tok_begin_id(tokenizer):
@@ -20,11 +20,11 @@ def tok_begin_id(tokenizer):
 
 
 def tok_sep(tokenizer):
-    if tokenizer.special_tokens_map.get('sep_token'):
-        return tokenizer.special_tokens_map.get('sep_token')
-    elif tokenizer.special_tokens_map.get('eos_token'):
-        tokenizer.special_tokens_map.get('eos_token')
-    return 'sep'
+    if tokenizer.special_tokens_map.get("sep_token"):
+        return tokenizer.special_tokens_map.get("sep_token")
+    elif tokenizer.special_tokens_map.get("eos_token"):
+        tokenizer.special_tokens_map.get("eos_token")
+    return "sep"
 
 
 def tok_sep_id(tokenizer):
@@ -32,9 +32,9 @@ def tok_sep_id(tokenizer):
 
 
 def tok_mask(tokenizer):
-    if tokenizer.special_tokens_map.get('mask_token'):
-        return tokenizer.special_tokens_map.get('mask_token')
-    return 'msk'
+    if tokenizer.special_tokens_map.get("mask_token"):
+        return tokenizer.special_tokens_map.get("mask_token")
+    return "msk"
 
 
 def tok_mask_id(tokenizer):
@@ -42,9 +42,9 @@ def tok_mask_id(tokenizer):
 
 
 def tok_pad(tokenizer):
-    if tokenizer.special_tokens_map.get('pad_token'):
-        return tokenizer.special_tokens_map.get('pad_token')
-    return 'pad'
+    if tokenizer.special_tokens_map.get("pad_token"):
+        return tokenizer.special_tokens_map.get("pad_token")
+    return "pad"
 
 
 def tok_pad_id(tokenizer):
@@ -56,30 +56,37 @@ def get_all_tok_from_config(config):
     return list(tokenizer.get_vocab().keys())
 
 
-def handle_exceed(tokenizer, seq, maxlen, mode=['noop', 'remove', 'slide', 'start_slice', 'end_slice'],
-                  keep_after_sep=True):
+def handle_exceed(
+    tokenizer,
+    seq,
+    maxlen,
+    mode=["noop", "remove", "slide", "start_slice", "end_slice"],
+    keep_after_sep=True,
+):
     if isinstance(seq, list):
         return seq, [[len(seq)]]
     mode = mode[0] if isinstance(mode, list) else mode
     sep_tok = tok_sep(tokenizer)
     sep_split = seq.split(sep_tok)
-    ext_seq = [sep_tok] + tokenizer.tokenize(sep_tok.join(sep_split[1:])) \
-        if len(sep_split) > 1 and keep_after_sep else []
+    ext_seq = ([sep_tok] + tokenizer.tokenize(sep_tok.join(sep_split[1:]))
+               if len(sep_split) > 1 and keep_after_sep else [])
     t_seq = tokenizer.tokenize(sep_split[0])
-    if mode == 'noop':
+    if mode == "noop":
         return [t_seq + ext_seq], [[0, len(t_seq + ext_seq)]]
-    if mode == 'remove':
+    if mode == "remove":
         if len(t_seq + ext_seq) <= maxlen:
             return [t_seq + ext_seq], [[0, len(t_seq + ext_seq)]]
         else:
             return [], [[0, 0]]
-    if mode == 'slide':
-        return nlp2.sliding_windows(t_seq, maxlen - len(ext_seq), append_seq=ext_seq)
-    if mode == 'start_slice':
+    if mode == "slide":
+        return nlp2.sliding_windows(t_seq,
+                                    maxlen - len(ext_seq),
+                                    append_seq=ext_seq)
+    if mode == "start_slice":
         slices = t_seq[:maxlen - len(ext_seq)]
         slices.extend(ext_seq)
         return [slices], [[0, maxlen - len(ext_seq)]]
-    if mode == 'end_slice':
+    if mode == "end_slice":
         start_pos = len(t_seq) + len(ext_seq) - maxlen
         slices = t_seq[start_pos:]
         slices.extend(ext_seq)


Test patch:
diff --git a/tfkit/test/__init__.py b/tfkit/test/__init__.py
index 8957dbb..e1dda87 100644
--- a/tfkit/test/__init__.py
+++ b/tfkit/test/__init__.py
@@ -2,41 +2,41 @@
 
 ROOT_DIR = os.path.dirname(os.path.abspath(__file__ + "/../../"))
 
-DATASET_DIR = os.path.join(ROOT_DIR, 'demo_data')
-TAG_DATASET = os.path.join(DATASET_DIR, 'tag.csv')
-CLAS_DATASET = os.path.join(DATASET_DIR, 'classification.csv')
-GEN_DATASET = os.path.join(DATASET_DIR, 'generation.csv')
-MASK_DATASET = os.path.join(DATASET_DIR, 'mask.csv')
-MCQ_DATASET = os.path.join(DATASET_DIR, 'mcq.csv')
-QA_DATASET = os.path.join(DATASET_DIR, 'qa.csv')
-ADDTOK_DATASET = os.path.join(DATASET_DIR, 'unk_tok.csv')
-NEWTOKEN_FILE = os.path.join(DATASET_DIR, 'tok_list.txt')
+DATASET_DIR = os.path.join(ROOT_DIR, "demo_data")
+TAG_DATASET = os.path.join(DATASET_DIR, "tag.csv")
+CLAS_DATASET = os.path.join(DATASET_DIR, "classification.csv")
+GEN_DATASET = os.path.join(DATASET_DIR, "generation.csv")
+MASK_DATASET = os.path.join(DATASET_DIR, "mask.csv")
+MCQ_DATASET = os.path.join(DATASET_DIR, "mcq.csv")
+QA_DATASET = os.path.join(DATASET_DIR, "qa.csv")
+ADDTOK_DATASET = os.path.join(DATASET_DIR, "unk_tok.csv")
+NEWTOKEN_FILE = os.path.join(DATASET_DIR, "tok_list.txt")
 
-MODEL_SAVE_DIR = os.path.join(ROOT_DIR, 'tfkit/test/cache/')
-ADDTOKFREQ_SAVE_DIR = os.path.join(MODEL_SAVE_DIR, 'addtokfreq/')
-ADDTOKFILE_SAVE_DIR = os.path.join(MODEL_SAVE_DIR, 'addtokfile/')
-CLAS_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'clas/')
-TAG_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'tag/')
-TAGCRF_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'tagcrf/')
-ONEBYONE_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'onebyone/')
-CLM_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'clm/')
-SEQ2SEQ_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'seq2seq/')
-ONCE_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'once/')
-ONCECTC_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'oncectc/')
-MASK_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'mask/')
-MCQ_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'mcq/')
-QA_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'qa/')
-MTTASK_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, 'mttask/')
+MODEL_SAVE_DIR = os.path.join(ROOT_DIR, "tfkit/test/cache/")
+ADDTOKFREQ_SAVE_DIR = os.path.join(MODEL_SAVE_DIR, "addtokfreq/")
+ADDTOKFILE_SAVE_DIR = os.path.join(MODEL_SAVE_DIR, "addtokfile/")
+CLAS_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "clas/")
+TAG_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "tag/")
+TAGCRF_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "tagcrf/")
+ONEBYONE_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "onebyone/")
+CLM_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "clm/")
+SEQ2SEQ_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "seq2seq/")
+ONCE_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "once/")
+ONCECTC_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "oncectc/")
+MASK_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "mask/")
+MCQ_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "mcq/")
+QA_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "qa/")
+MTTASK_MODEL_DIR = os.path.join(MODEL_SAVE_DIR, "mttask/")
 
-ONEBYONE_MODEL_PATH = os.path.join(ONEBYONE_MODEL_DIR, '2.pt')
-ONCE_MODEL_PATH = os.path.join(ONCE_MODEL_DIR, '2.pt')
-ONCECTC_MODEL_PATH = os.path.join(ONCECTC_MODEL_DIR, '1.pt')
-SEQ2SEQ_MODEL_PATH = os.path.join(SEQ2SEQ_MODEL_DIR, '2.pt')
-CLM_MODEL_PATH = os.path.join(CLM_MODEL_DIR, '2.pt')
-CLAS_MODEL_PATH = os.path.join(CLAS_MODEL_DIR, '2.pt')
-MASK_MODEL_PATH = os.path.join(MASK_MODEL_DIR, '2.pt')
-MCQ_MODEL_PATH = os.path.join(MCQ_MODEL_DIR, '2.pt')
-TAG_MODEL_PATH = os.path.join(TAG_MODEL_DIR, '2.pt')
-QA_MODEL_PATH = os.path.join(QA_MODEL_DIR, '2.pt')
-ADDTOKFREQ_MODEL_PATH = os.path.join(ADDTOKFREQ_SAVE_DIR, '2.pt')
-ADDTOKFILE_MODEL_PATH = os.path.join(ADDTOKFILE_SAVE_DIR, '2.pt')
+ONEBYONE_MODEL_PATH = os.path.join(ONEBYONE_MODEL_DIR, "2.pt")
+ONCE_MODEL_PATH = os.path.join(ONCE_MODEL_DIR, "2.pt")
+ONCECTC_MODEL_PATH = os.path.join(ONCECTC_MODEL_DIR, "1.pt")
+SEQ2SEQ_MODEL_PATH = os.path.join(SEQ2SEQ_MODEL_DIR, "2.pt")
+CLM_MODEL_PATH = os.path.join(CLM_MODEL_DIR, "2.pt")
+CLAS_MODEL_PATH = os.path.join(CLAS_MODEL_DIR, "2.pt")
+MASK_MODEL_PATH = os.path.join(MASK_MODEL_DIR, "2.pt")
+MCQ_MODEL_PATH = os.path.join(MCQ_MODEL_DIR, "2.pt")
+TAG_MODEL_PATH = os.path.join(TAG_MODEL_DIR, "2.pt")
+QA_MODEL_PATH = os.path.join(QA_MODEL_DIR, "2.pt")
+ADDTOKFREQ_MODEL_PATH = os.path.join(ADDTOKFREQ_SAVE_DIR, "2.pt")
+ADDTOKFILE_MODEL_PATH = os.path.join(ADDTOKFILE_SAVE_DIR, "2.pt")
diff --git a/tfkit/test/task/test_task_model.py b/tfkit/test/task/test_task_model.py
index 7220a8b..febeea2 100644
--- a/tfkit/test/task/test_task_model.py
+++ b/tfkit/test/task/test_task_model.py
@@ -1,33 +1,47 @@
 import unittest
 
+from torch import Tensor
+from transformers import AutoModel
+from transformers import AutoTokenizer
+from transformers import BertTokenizer
+
 import tfkit
 from tfkit.utility.dataloader import pad_batch
-from torch import Tensor
-from transformers import BertTokenizer, AutoModel, AutoTokenizer
 
 
 class TestModel(unittest.TestCase):
 
-
     def testGenerationModel(self):
         input = "See you next time"
         maxlen = 32
-        tokenizer = AutoTokenizer.from_pretrained('sshleifer/bart-tiny-random')
-        pretrained = AutoModel.from_pretrained('sshleifer/bart-tiny-random')
+        tokenizer = AutoTokenizer.from_pretrained("sshleifer/bart-tiny-random")
+        pretrained = AutoModel.from_pretrained("sshleifer/bart-tiny-random")
         # tfkit.task.seq2seq, tfkit.task.once, tfkit.task.oncectc, tfkit.task.clm
-        for gmodel in [tfkit.task.seq2seq, tfkit.task.once, tfkit.task.oncectc, tfkit.task.clm]:
+        for gmodel in [
+                tfkit.task.seq2seq,
+                tfkit.task.once,
+                tfkit.task.oncectc,
+                tfkit.task.clm,
+        ]:
             print(str(gmodel))
             model = gmodel.Model(tokenizer, pretrained, maxlen=maxlen)
-            preprocessor = gmodel.Preprocessor(tokenizer, maxlen=maxlen, handle_exceed='start_slice', reserved_len=0)
-            for preprocessed_item in preprocessor.preprocess(
-                    {'task': 'taskA', 'input': input}):
+            preprocessor = gmodel.Preprocessor(tokenizer,
+                                               maxlen=maxlen,
+                                               handle_exceed="start_slice",
+                                               reserved_len=0)
+            for preprocessed_item in preprocessor.preprocess({
+                    "task": "taskA",
+                    "input": input
+            }):
                 print("preprocessed_item", preprocessed_item)
-                feature = preprocessor.postprocess(preprocessed_item, tokenizer, maxlen=maxlen)
+                feature = preprocessor.postprocess(preprocessed_item,
+                                                   tokenizer,
+                                                   maxlen=maxlen)
                 feature = preprocessor.postprocess_batch(feature)
                 print(model(feature, eval=True))
                 self.assertTrue(isinstance(model(feature, eval=True), dict))
                 model_dict = model(feature, eval=True)
-                self.assertTrue('max_item' in model_dict)
+                self.assertTrue("max_item" in model_dict)
 
             # greedy
             print("greedy")
@@ -38,7 +52,11 @@ def testGenerationModel(self):
             self.assertTrue(isinstance(detail, dict))
 
             # TopK
-            result, detail = model.predict(input=input, decodenum=3, mode='topK', topK=3, filtersim=False)
+            result, detail = model.predict(input=input,
+                                           decodenum=3,
+                                           mode="topK",
+                                           topK=3,
+                                           filtersim=False)
             print("topK", result)
             self.assertTrue(len(result) == 3)
             self.assertTrue(isinstance(result, list))
@@ -52,7 +70,10 @@ def testGenerationModel(self):
             self.assertTrue(isinstance(detail, dict))
 
             # TopP
-            result, detail = model.predict(input=input, decodenum=3, mode='topP', topP=0.8)
+            result, detail = model.predict(input=input,
+                                           decodenum=3,
+                                           mode="topP",
+                                           topP=0.8)
             print("TopP", len(result), result, model_dict)
             self.assertTrue(len(result) == 3)
             self.assertTrue(isinstance(result, list))
@@ -65,7 +86,8 @@ def testGenerationModel(self):
             self.assertTrue(isinstance(detail, dict))
             print("exceed max len", result)
 
-            result, model_dict = model.predict(input="T " * 550, reserved_len=10)
+            result, model_dict = model.predict(input="T " * 550,
+                                               reserved_len=10)
             print(result)
             self.assertTrue(isinstance(result, list))
             self.assertTrue(isinstance(detail, dict))
@@ -175,31 +197,31 @@ def testGenerationModel(self):
     #         print(result)
     #         self.assertTrue(isinstance(result, list))
 
-        # proc = tfkit.task.tag.Preprocessor(tokenizer, maxlen=512, handle_exceed='start_slice',
-        #                                   reserved_len=0)
-        # for items in proc.prepare_data({"input": input}):
-        #     raw_input = items['raw_input']
-        #     feature = proc.postprocess(items, tokenizer, 512)
-        #     for k, v in feature.items():
-        #         feature[k] = [v]
-        #     self.assertTrue(isinstance(model(feature), Tensor))
-        #     print(model(feature))
-        #     # test eval
-        #     model_dict = model(feature, eval=True)
-        #     self.assertTrue('label_prob_all' in model_dict)
-        #     self.assertTrue('label_map' in model_dict)
-        #     self.assertEqual(len(model_dict['label_map']), len(input.split(" ")))
-        #
-        # # test predict
-        # result, model_dict = model.predict(input=input, start_contain="A", end_contain="A")
-        # self.assertTrue('label_prob_all' in model_dict[0])
-        # self.assertTrue('label_map' in model_dict[0])
-        # print("result", result, len(result))
-        # self.assertTrue(isinstance(result, list))
-        #
-        # # test exceed 512
-        # for merge_strategy in ['minentropy', 'maxcount', 'maxprob']:
-        #     result, model_dict = model.predict(input=" ".join([str(i) for i in range(1000)]),
-        #                                        merge_strategy=merge_strategy, start_contain="A", end_contain="A")
-        #     print(result)
-        #     self.assertTrue(isinstance(result, list))
+    # proc = tfkit.task.tag.Preprocessor(tokenizer, maxlen=512, handle_exceed='start_slice',
+    #                                   reserved_len=0)
+    # for items in proc.prepare_data({"input": input}):
+    #     raw_input = items['raw_input']
+    #     feature = proc.postprocess(items, tokenizer, 512)
+    #     for k, v in feature.items():
+    #         feature[k] = [v]
+    #     self.assertTrue(isinstance(model(feature), Tensor))
+    #     print(model(feature))
+    #     # test eval
+    #     model_dict = model(feature, eval=True)
+    #     self.assertTrue('label_prob_all' in model_dict)
+    #     self.assertTrue('label_map' in model_dict)
+    #     self.assertEqual(len(model_dict['label_map']), len(input.split(" ")))
+    #
+    # # test predict
+    # result, model_dict = model.predict(input=input, start_contain="A", end_contain="A")
+    # self.assertTrue('label_prob_all' in model_dict[0])
+    # self.assertTrue('label_map' in model_dict[0])
+    # print("result", result, len(result))
+    # self.assertTrue(isinstance(result, list))
+    #
+    # # test exceed 512
+    # for merge_strategy in ['minentropy', 'maxcount', 'maxprob']:
+    #     result, model_dict = model.predict(input=" ".join([str(i) for i in range(1000)]),
+    #                                        merge_strategy=merge_strategy, start_contain="A", end_contain="A")
+    #     print(result)
+    #     self.assertTrue(isinstance(result, list))
diff --git a/tfkit/test/test_atrain.py b/tfkit/test/test_atrain.py
index d7689ff..bcbc811 100644
--- a/tfkit/test/test_atrain.py
+++ b/tfkit/test/test_atrain.py
@@ -1,7 +1,8 @@
 import unittest
 
 import pytest
-from transformers import BertTokenizer, AutoModel
+from transformers import AutoModel
+from transformers import BertTokenizer
 
 import tfkit
 from tfkit.test import *
@@ -11,128 +12,396 @@
 class TestTrain(unittest.TestCase):
 
     def testHelp(self):
-        result = os.system('tfkit-train -h')
-        assert (result == 0)
+        result = os.system("tfkit-train -h")
+        assert result == 0
 
     def test_parser(self):
-        input_arg, model_arg = tfkit.train.parse_train_args(
-            ['--task', 'once', '--train', 'train.csv', '--test', 'test.csv', '--config',
-             'voidful/albert_chinese_tiny'])
+        input_arg, model_arg = tfkit.train.parse_train_args([
+            "--task",
+            "once",
+            "--train",
+            "train.csv",
+            "--test",
+            "test.csv",
+            "--config",
+            "voidful/albert_chinese_tiny",
+        ])
         print(input_arg, model_arg)
-        self.assertTrue(input_arg.get('task') == ['once'])
-        self.assertTrue(isinstance(input_arg.get('train'), list))
+        self.assertTrue(input_arg.get("task") == ["once"])
+        self.assertTrue(isinstance(input_arg.get("train"), list))
 
-        input_arg, model_arg = tfkit.train.parse_train_args(
-            ['--task', 'once', '--train', 'train.csv', '--test', 'test.csv', '--config',
-             'voidful/albert_chinese_tiny', '--likelihood', 'pos'])
+        input_arg, model_arg = tfkit.train.parse_train_args([
+            "--task",
+            "once",
+            "--train",
+            "train.csv",
+            "--test",
+            "test.csv",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--likelihood",
+            "pos",
+        ])
         print(input_arg, model_arg)
-        self.assertTrue(model_arg.get('likelihood') == 'pos')
-        self.assertTrue(isinstance(input_arg.get('train'), list))
+        self.assertTrue(model_arg.get("likelihood") == "pos")
+        self.assertTrue(isinstance(input_arg.get("train"), list))
 
     def test_optimizer(self):
-        model_class = load_model_class('clas')
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
-        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')
-        model = model_class.Model(tokenizer=tokenizer, pretrained=pretrained, tasks_detail={"taskA": ["a", "b"]},
-                                  maxlen=128)
+        model_class = load_model_class("clas")
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
+        pretrained = AutoModel.from_pretrained("voidful/albert_chinese_tiny")
+        model = model_class.Model(
+            tokenizer=tokenizer,
+            pretrained=pretrained,
+            tasks_detail={"taskA": ["a", "b"]},
+            maxlen=128,
+        )
         optim, scheduler = tfkit.train.optimizer(model, lr=0.1, total_step=10)
         print(optim, scheduler)
         optim.zero_grad()
         scheduler.step()
 
     def testMultiTask(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', MTTASK_MODEL_DIR, '--train', CLAS_DATASET, GEN_DATASET,
-             '--lr', '5e-5', '--test', CLAS_DATASET, GEN_DATASET, '--task', 'once', 'clm', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            MTTASK_MODEL_DIR,
+            "--train",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--task",
+            "once",
+            "clm",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+        ])
         result = os.system(
-            'tfkit-train --batch 2 --epoch 2 --savedir ' + MTTASK_MODEL_DIR + ' --train ' + CLAS_DATASET + ' ' + GEN_DATASET + ' --lr 5e-5 --test ' + CLAS_DATASET + ' ' + GEN_DATASET + ' --task once clm --config voidful/albert_chinese_tiny --maxlen 50')
+            "tfkit-train --batch 2 --epoch 2 --savedir " + MTTASK_MODEL_DIR +
+            " --train " + CLAS_DATASET + " " + GEN_DATASET +
+            " --lr 5e-5 --test " + CLAS_DATASET + " " + GEN_DATASET +
+            " --task once clm --config voidful/albert_chinese_tiny --maxlen 50"
+        )
         self.assertTrue(result == 0)
 
     def testGenOnce(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ONCE_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', GEN_DATASET, '--task', 'once', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ONCE_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "once",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+        ])
         result = os.system(
-            'tfkit-train --batch 2 --epoch 2 --savedir ' + ONCE_MODEL_DIR + ' --train ' + GEN_DATASET + ' --test ' + GEN_DATASET + ' --task once --config voidful/albert_chinese_tiny --maxlen 50')
+            "tfkit-train --batch 2 --epoch 2 --savedir " + ONCE_MODEL_DIR +
+            " --train " + GEN_DATASET + " --test " + GEN_DATASET +
+            " --task once --config voidful/albert_chinese_tiny --maxlen 50")
         self.assertTrue(result == 0)
 
     def testGenOnceCTC(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ONCECTC_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '3e-4', '--test', GEN_DATASET, '--task', 'oncectc', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ONCECTC_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "3e-4",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "oncectc",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+        ])
         result = os.system(
-            'tfkit-train --batch 2 --epoch 2 --savedir ' + ONCE_MODEL_DIR + ' --train ' + GEN_DATASET + ' --test ' + GEN_DATASET + ' --task oncectc --config voidful/albert_chinese_tiny --maxlen 50')
+            "tfkit-train --batch 2 --epoch 2 --savedir " + ONCE_MODEL_DIR +
+            " --train " + GEN_DATASET + " --test " + GEN_DATASET +
+            " --task oncectc --config voidful/albert_chinese_tiny --maxlen 50")
         self.assertTrue(result == 0)
 
     def testGenSeq2Seq(self):
         # result = os.system(
         #     'tfkit-train --batch 2 --epoch 1 --savedir ' + SEQ2SEQ_MODEL_DIR + ' --train ' + GEN_DATASET + ' --test ' + GEN_DATASET + ' --task seq2seq --config prajjwal1/bert-small --maxlen 50 --selfkd True')
         # self.assertTrue(result == 0)
-        tfkit.train.main(
-            ['--batch', '1', '--epoch', '1', '--savedir', SEQ2SEQ_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-4', '--test', GEN_DATASET, '--task', 'seq2seq', '--config',
-             'prajjwal1/bert-small', '--maxlen', '20'])
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '2', '--savedir', SEQ2SEQ_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-4', '--test', GEN_DATASET, '--task', 'seq2seq', '--config',
-             'prajjwal1/bert-small', '--maxlen', '20', '--likelihood', 'pos'])
+        tfkit.train.main([
+            "--batch",
+            "1",
+            "--epoch",
+            "1",
+            "--savedir",
+            SEQ2SEQ_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-4",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "seq2seq",
+            "--config",
+            "prajjwal1/bert-small",
+            "--maxlen",
+            "20",
+        ])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "2",
+            "--savedir",
+            SEQ2SEQ_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-4",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "seq2seq",
+            "--config",
+            "prajjwal1/bert-small",
+            "--maxlen",
+            "20",
+            "--likelihood",
+            "pos",
+        ])
 
     def testGenCLM(self):
         result = os.system(
-            'tfkit-train --batch 2 --epoch 1 --savedir ' + CLM_MODEL_DIR + ' --train ' + GEN_DATASET + ' --test ' + GEN_DATASET + ' --task clm --config prajjwal1/bert-small --maxlen 50')
+            "tfkit-train --batch 2 --epoch 1 --savedir " + CLM_MODEL_DIR +
+            " --train " + GEN_DATASET + " --test " + GEN_DATASET +
+            " --task clm --config prajjwal1/bert-small --maxlen 50")
         self.assertTrue(result == 0)
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '2', '--savedir', CLM_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-4', '--test', GEN_DATASET, '--task', 'clm', '--config',
-             'prajjwal1/bert-small', '--maxlen', '20'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "2",
+            "--savedir",
+            CLM_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-4",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "clm",
+            "--config",
+            "prajjwal1/bert-small",
+            "--maxlen",
+            "20",
+        ])
 
     def testAddTokenFile(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ADDTOKFILE_SAVE_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', ADDTOK_DATASET, '--task', 'clm', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '100', '--add_tokens_file', NEWTOKEN_FILE])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ADDTOKFILE_SAVE_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            ADDTOK_DATASET,
+            "--task",
+            "clm",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "100",
+            "--add_tokens_file",
+            NEWTOKEN_FILE,
+        ])
         result = os.system(
-            f'tfkit-train --batch 2 --add_tokens_file {NEWTOKEN_FILE}  --savedir {ADDTOKFILE_SAVE_DIR} --epoch 2  --train {ADDTOK_DATASET}  --test {ADDTOK_DATASET} --task clm --config voidful/albert_chinese_tiny --maxlen 50')
+            f"tfkit-train --batch 2 --add_tokens_file {NEWTOKEN_FILE}  --savedir {ADDTOKFILE_SAVE_DIR} --epoch 2  --train {ADDTOK_DATASET}  --test {ADDTOK_DATASET} --task clm --config voidful/albert_chinese_tiny --maxlen 50"
+        )
         self.assertTrue(result == 0)
 
     def testResume(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ONCE_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', GEN_DATASET, '--task', 'once', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--tag', 'testresume'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ONCE_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "once",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--tag",
+            "testresume",
+        ])
 
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ONCE_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', GEN_DATASET, '--task', 'once', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--resume', os.path.join(ONCE_MODEL_DIR, "1.pt")])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ONCE_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "once",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--resume",
+            os.path.join(ONCE_MODEL_DIR, "1.pt"),
+        ])
 
     def testResumeMultiModel(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', MTTASK_MODEL_DIR, '--train', CLAS_DATASET, GEN_DATASET,
-             '--lr', '5e-5', '--test', CLAS_DATASET, GEN_DATASET, '--task', 'once', 'clm', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--tag', 'once', 'clm'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            MTTASK_MODEL_DIR,
+            "--train",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--task",
+            "once",
+            "clm",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--tag",
+            "once",
+            "clm",
+        ])
         # resume to train all task
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', MTTASK_MODEL_DIR, '--train', CLAS_DATASET, GEN_DATASET,
-             '--lr', '5e-5', '--test', CLAS_DATASET, GEN_DATASET, '--task', 'once', 'clm', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--tag', 'once', 'clm', '--resume',
-             os.path.join(MTTASK_MODEL_DIR, "1.pt")])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            MTTASK_MODEL_DIR,
+            "--train",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            CLAS_DATASET,
+            GEN_DATASET,
+            "--task",
+            "once",
+            "clm",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--tag",
+            "once",
+            "clm",
+            "--resume",
+            os.path.join(MTTASK_MODEL_DIR, "1.pt"),
+        ])
         # resume to train only one task
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', MTTASK_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', GEN_DATASET, '--task', 'clm', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--resume', os.path.join(MTTASK_MODEL_DIR, "1.pt"),
-             '--tag', 'clm'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            MTTASK_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "clm",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--resume",
+            os.path.join(MTTASK_MODEL_DIR, "1.pt"),
+            "--tag",
+            "clm",
+        ])
 
     @pytest.mark.skip()
     def testLoggerwandb(self):
-        tfkit.train.main(
-            ['--batch', '2', '--epoch', '1', '--savedir', ONCE_MODEL_DIR, '--train',
-             GEN_DATASET, '--lr', '5e-5', '--test', GEN_DATASET, '--task', 'once', '--config',
-             'voidful/albert_chinese_tiny', '--maxlen', '50', '--wandb'])
+        tfkit.train.main([
+            "--batch",
+            "2",
+            "--epoch",
+            "1",
+            "--savedir",
+            ONCE_MODEL_DIR,
+            "--train",
+            GEN_DATASET,
+            "--lr",
+            "5e-5",
+            "--test",
+            GEN_DATASET,
+            "--task",
+            "once",
+            "--config",
+            "voidful/albert_chinese_tiny",
+            "--maxlen",
+            "50",
+            "--wandb",
+        ])
 
     # def testClas(self):
     #     tfkit.train.main(
diff --git a/tfkit/test/test_package.py b/tfkit/test/test_package.py
index 0e9390f..5544e76 100644
--- a/tfkit/test/test_package.py
+++ b/tfkit/test/test_package.py
@@ -1,9 +1,10 @@
+import os
 import unittest
 
 from transformers import AutoTokenizer
 
 import tfkit
-import os
+
 
 class TestPackage(unittest.TestCase):
 
diff --git a/tfkit/test/test_zeval.py b/tfkit/test/test_zeval.py
index d8ecce6..59a2727 100644
--- a/tfkit/test/test_zeval.py
+++ b/tfkit/test/test_zeval.py
@@ -7,70 +7,133 @@
 class TestEval(unittest.TestCase):
 
     def testHelp(self):
-        result = os.system('tfkit-eval -h')
+        result = os.system("tfkit-eval -h")
         self.assertTrue(result == 0)
 
     def test_parser(self):
-        parser, _ = tfkit.eval.parse_eval_args(
-            ['--model', 'once', '--metric', 'emf1', '--valid', 'test.csv', '--print'])
+        parser, _ = tfkit.eval.parse_eval_args([
+            "--model", "once", "--metric", "emf1", "--valid", "test.csv",
+            "--print"
+        ])
         print(parser)
-        self.assertTrue(parser.get('model') == ['once'])
+        self.assertTrue(parser.get("model") == ["once"])
 
-        eval_parser, model_parser = tfkit.eval.parse_eval_args(
-            ['--model', 'once', '--metric', 'emf1', '--valid', 'test.csv', '--print', '--decodenum', '2'])
-        self.assertTrue(eval_parser.get('model') == ['once'])
-        self.assertTrue(model_parser.get('decodenum') == '2')
+        eval_parser, model_parser = tfkit.eval.parse_eval_args([
+            "--model",
+            "once",
+            "--metric",
+            "emf1",
+            "--valid",
+            "test.csv",
+            "--print",
+            "--decodenum",
+            "2",
+        ])
+        self.assertTrue(eval_parser.get("model") == ["once"])
+        self.assertTrue(model_parser.get("decodenum") == "2")
 
     def testEvalGen(self):
-        tfkit.eval.main(
-            ['--model', ONCE_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print'])
-        result = os.system(
-            'tfkit-eval --model ' + ONCE_MODEL_PATH + ' --valid ' + GEN_DATASET + ' --metric emf1 --print')
+        tfkit.eval.main([
+            "--model",
+            ONCE_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+        ])
+        result = os.system("tfkit-eval --model " + ONCE_MODEL_PATH +
+                           " --valid " + GEN_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     def testEvalGenOnce(self):
-        tfkit.eval.main(
-            ['--model', ONCE_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print'])
-        result = os.system(
-            'tfkit-eval --model ' + ONCE_MODEL_PATH + ' --valid ' + GEN_DATASET + ' --metric emf1 --print')
+        tfkit.eval.main([
+            "--model",
+            ONCE_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+        ])
+        result = os.system("tfkit-eval --model " + ONCE_MODEL_PATH +
+                           " --valid " + GEN_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     def testEvalGenOnceCTC(self):
-        tfkit.eval.main(
-            ['--model', ONCECTC_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print'])
-        result = os.system(
-            'tfkit-eval --model ' + ONCECTC_MODEL_PATH + ' --valid ' + GEN_DATASET + ' --metric emf1 --print')
+        tfkit.eval.main([
+            "--model",
+            ONCECTC_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+        ])
+        result = os.system("tfkit-eval --model " + ONCECTC_MODEL_PATH +
+                           " --valid " + GEN_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     def testEvalSeq2Seq(self):
-        tfkit.eval.main(
-            ['--model', SEQ2SEQ_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print',
-             '--decodenum', '2'])
-        tfkit.eval.main(
-            ['--model', SEQ2SEQ_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print'])
-        result = os.system(
-            'tfkit-eval --model ' + SEQ2SEQ_MODEL_PATH + ' --valid ' + GEN_DATASET + ' --metric emf1 --print')
+        tfkit.eval.main([
+            "--model",
+            SEQ2SEQ_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+            "--decodenum",
+            "2",
+        ])
+        tfkit.eval.main([
+            "--model",
+            SEQ2SEQ_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+        ])
+        result = os.system("tfkit-eval --model " + SEQ2SEQ_MODEL_PATH +
+                           " --valid " + GEN_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     def testEvalCLM(self):
-        tfkit.eval.main(
-            ['--model', CLM_MODEL_PATH, '--valid', GEN_DATASET, '--metric', 'emf1', '--print'])
-        result = os.system(
-            'tfkit-eval --model ' + CLM_MODEL_PATH + ' --valid ' + GEN_DATASET + ' --metric emf1 --print')
+        tfkit.eval.main([
+            "--model",
+            CLM_MODEL_PATH,
+            "--valid",
+            GEN_DATASET,
+            "--metric",
+            "emf1",
+            "--print",
+        ])
+        result = os.system("tfkit-eval --model " + CLM_MODEL_PATH +
+                           " --valid " + GEN_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     def testEvalAddedTokenModel(self):
-        result = os.system(
-            'tfkit-eval --model ' + ADDTOKFILE_MODEL_PATH + ' --valid ' + ADDTOK_DATASET + ' --metric emf1 --print')
+        result = os.system("tfkit-eval --model " + ADDTOKFILE_MODEL_PATH +
+                           " --valid " + ADDTOK_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
-        result = os.system(
-            'tfkit-eval --model ' + ADDTOKFILE_MODEL_PATH + ' --valid ' + ADDTOK_DATASET + ' --metric emf1 --print')
+        result = os.system("tfkit-eval --model " + ADDTOKFILE_MODEL_PATH +
+                           " --valid " + ADDTOK_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
-        result = os.system(
-            'tfkit-eval --model ' + ADDTOKFREQ_MODEL_PATH + ' --valid ' + ADDTOK_DATASET + ' --metric emf1 --print')
+        result = os.system("tfkit-eval --model " + ADDTOKFREQ_MODEL_PATH +
+                           " --valid " + ADDTOK_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
-        result = os.system(
-            'tfkit-eval --model ' + ADDTOKFREQ_MODEL_PATH + ' --valid ' + ADDTOK_DATASET + ' --metric emf1 --print')
+        result = os.system("tfkit-eval --model " + ADDTOKFREQ_MODEL_PATH +
+                           " --valid " + ADDTOK_DATASET +
+                           " --metric emf1 --print")
         self.assertTrue(result == 0)
 
     # def testEvalQA(self):
@@ -92,4 +155,4 @@ def testEvalAddedTokenModel(self):
     #         ['--model', TAG_MODEL_PATH, '--valid', TAG_DATASET, '--metric', 'clas', '--print'])
     #     result = os.system(
     #         'tfkit-eval --model ' + TAG_MODEL_PATH + ' --valid ' + TAG_DATASET + ' --metric clas --print')
-    #     self.assertTrue(result == 0)
\ No newline at end of file
+    #     self.assertTrue(result == 0)
diff --git a/tfkit/test/utility/test_utility_datafile.py b/tfkit/test/utility/test_utility_datafile.py
index 6e0fa30..835aa34 100644
--- a/tfkit/test/utility/test_utility_datafile.py
+++ b/tfkit/test/utility/test_utility_datafile.py
@@ -7,11 +7,13 @@
 class TestDataFile(unittest.TestCase):
 
     def test_get_x_data_from_file(self):
-        for get_x_iter in [get_gen_data_from_file(GEN_DATASET),
-                           get_qa_data_from_file(QA_DATASET),
-                           get_tag_data_from_file(TAG_DATASET),
-                           get_clas_data_from_file(CLAS_DATASET),
-                           get_multiclas_data_from_file(CLAS_DATASET)]:
+        for get_x_iter in [
+                get_gen_data_from_file(GEN_DATASET),
+                get_qa_data_from_file(QA_DATASET),
+                get_tag_data_from_file(TAG_DATASET),
+                get_clas_data_from_file(CLAS_DATASET),
+                get_multiclas_data_from_file(CLAS_DATASET),
+        ]:
             while True:
                 try:
                     print(next(get_x_iter))
diff --git a/tfkit/test/utility/test_utility_dataloader.py b/tfkit/test/utility/test_utility_dataloader.py
index b7f0369..3f1380b 100644
--- a/tfkit/test/utility/test_utility_dataloader.py
+++ b/tfkit/test/utility/test_utility_dataloader.py
@@ -8,11 +8,20 @@
 class TestUtilityDataLoader(unittest.TestCase):
 
     def test_batch_reduce_pad(self):
-        k = [{'input': torch.tensor([1, 2, 3])},
-             {'input': torch.tensor([3, 4])},
-             {'input': torch.tensor([5])}]
+        k = [
+            {
+                "input": torch.tensor([1, 2, 3])
+            },
+            {
+                "input": torch.tensor([3, 4])
+            },
+            {
+                "input": torch.tensor([5])
+            },
+        ]
         reduced_batch = pad_batch(k)
-        self.assertEqual(len(reduced_batch[0]['input']), len(reduced_batch[1]['input']))
+        self.assertEqual(len(reduced_batch[0]["input"]),
+                         len(reduced_batch[1]["input"]))
         print(reduced_batch)
-        self.assertCountEqual(reduced_batch[0]['input'], [1, 2, 3])
-        self.assertCountEqual(reduced_batch[1]['input'], [3, 4, 0])
+        self.assertCountEqual(reduced_batch[0]["input"], [1, 2, 3])
+        self.assertCountEqual(reduced_batch[1]["input"], [3, 4, 0])
diff --git a/tfkit/test/utility/test_utility_datapreprocess.py b/tfkit/test/utility/test_utility_datapreprocess.py
index 419da38..e12caa2 100644
--- a/tfkit/test/utility/test_utility_datapreprocess.py
+++ b/tfkit/test/utility/test_utility_datapreprocess.py
@@ -7,11 +7,13 @@
 class TestDataPreprocess(unittest.TestCase):
 
     def test_get_x_data_from_file(self):
-        for get_x_iter in [get_gen_data_from_file(GEN_DATASET),
-                           get_qa_data_from_file(QA_DATASET),
-                           get_tag_data_from_file(TAG_DATASET),
-                           get_clas_data_from_file(CLAS_DATASET),
-                           get_multiclas_data_from_file(CLAS_DATASET)]:
+        for get_x_iter in [
+                get_gen_data_from_file(GEN_DATASET),
+                get_qa_data_from_file(QA_DATASET),
+                get_tag_data_from_file(TAG_DATASET),
+                get_clas_data_from_file(CLAS_DATASET),
+                get_multiclas_data_from_file(CLAS_DATASET),
+        ]:
             while True:
                 try:
                     print(next(get_x_iter))
diff --git a/tfkit/test/utility/test_utility_dataset.py b/tfkit/test/utility/test_utility_dataset.py
index 91e384e..c9d513d 100644
--- a/tfkit/test/utility/test_utility_dataset.py
+++ b/tfkit/test/utility/test_utility_dataset.py
@@ -1,49 +1,55 @@
 import os
 import sys
+import unittest
 
 import numpy
 from transformers import AutoTokenizer
 
-from tfkit.utility.dataset import get_dataset, TFKitDataset
-from tfkit.utility.model import load_pretrained_tokenizer, load_model_class
+import tfkit
+from tfkit.test import *
+from tfkit.utility.dataset import get_dataset
+from tfkit.utility.dataset import TFKitDataset
+from tfkit.utility.model import load_model_class
+from tfkit.utility.model import load_pretrained_tokenizer
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import unittest
-import tfkit
-from tfkit.test import *
-
 
 class TestDataset(unittest.TestCase):
 
     def check_type_for_dataloader(self, data_item):
-        if (isinstance(data_item, list) and not isinstance(data_item[-1], str) and self.check_type_for_dataloader(
-                data_item[-1])) or \
-                isinstance(data_item, numpy.ndarray) or \
-                isinstance(data_item, int):
+        if ((isinstance(data_item, list)
+             and not isinstance(data_item[-1], str)
+             and self.check_type_for_dataloader(data_item[-1]))
+                or isinstance(data_item, numpy.ndarray)
+                or isinstance(data_item, int)):
             return True
         else:
             return False
 
     def test_check_type_for_dataloader(self):
         self.assertTrue(self.check_type_for_dataloader(123))
-        self.assertFalse(self.check_type_for_dataloader('a'))
-        self.assertFalse(self.check_type_for_dataloader(['a']))
-        self.assertFalse(self.check_type_for_dataloader([{'a'}]))
+        self.assertFalse(self.check_type_for_dataloader("a"))
+        self.assertFalse(self.check_type_for_dataloader(["a"]))
+        self.assertFalse(self.check_type_for_dataloader([{"a"}]))
 
     def test_get_dataset(self):
-        tokenizer = load_pretrained_tokenizer('voidful/albert_chinese_tiny')
-        model_class = load_model_class('clas')
+        tokenizer = load_pretrained_tokenizer("voidful/albert_chinese_tiny")
+        model_class = load_model_class("clas")
         file_path = CLAS_DATASET
         dataset_arg = {
-            'maxlen': 123,
-            'handle_exceed': 'slide',
-            'config': 'voidful/albert_chinese_tiny',
-            'cache': False
+            "maxlen": 123,
+            "handle_exceed": "slide",
+            "config": "voidful/albert_chinese_tiny",
+            "cache": False,
         }
-        ds = get_dataset(file_path=file_path, task_class=model_class, tokenizer=tokenizer,
-                                       parameter=dataset_arg)
+        ds = get_dataset(
+            file_path=file_path,
+            task_class=model_class,
+            tokenizer=tokenizer,
+            parameter=dataset_arg,
+        )
         print(ds, ds[0])
 
     # def testClassifier(self):
diff --git a/tfkit/test/utility/test_utility_eval_metric.py b/tfkit/test/utility/test_utility_eval_metric.py
index e1a229c..ad60708 100644
--- a/tfkit/test/utility/test_utility_eval_metric.py
+++ b/tfkit/test/utility/test_utility_eval_metric.py
@@ -1,335 +1,525 @@
-import sys
 import os
+import sys
+import unittest
 
 import pytest
+from transformers import AutoTokenizer
+from transformers import BertTokenizer
+
+import tfkit
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import unittest
-import tfkit
-from transformers import BertTokenizer, AutoTokenizer
-
 
 class TestEval(unittest.TestCase):
+
     @pytest.mark.skip()
     def testER(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "abc", "abb///acc///abc", task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "abc", "abb///acc///abc", task="default")
+        for s in eval.cal_score("er"):
             print(s)
-            self.assertTrue(s[1]['WER'] == 0)
-            self.assertTrue(s[1]['CER'] == 0)
+            self.assertTrue(s[1]["WER"] == 0)
+            self.assertTrue(s[1]["CER"] == 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "ab", "abb///acc///ab c", task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "ab", "abb///acc///ab c", task="default")
+        for s in eval.cal_score("er"):
             print(s)
-            self.assertTrue(s[1]['WER'] == 50)
+            self.assertTrue(s[1]["WER"] == 50)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a b c", "a b b///a c c///", task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "a b c", "a b b///a c c///", task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "", "a b b///a c c", task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "", "a b b///a c c", task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a", ["a"], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "a", ["a"], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a", ["b"], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "a", ["b"], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "b", ["a"], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "b", ["a"], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "b", ["b"], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "b", ["b"], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "", [""], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "", [""], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "梵 語", ["梵 語"], task='default')
-        eval.add_record("input", "梵 語", ["梵 語d"], task='default')
-        eval.add_record("input", "梵 語", ["梵 語c"], task='default')
-        for s in eval.cal_score('er'):
+        eval.add_record("input", "梵 語", ["梵 語"], task="default")
+        eval.add_record("input", "梵 語", ["梵 語d"], task="default")
+        eval.add_record("input", "梵 語", ["梵 語c"], task="default")
+        for s in eval.cal_score("er"):
             print(s)
 
     def testEMF1(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "abc", "abb///acc///abc", task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "abc", "abb///acc///abc", task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 1)
-            self.assertTrue(s[1]['F1'] == 1)
+            self.assertTrue(s[1]["EM"] == 1)
+            self.assertTrue(s[1]["F1"] == 1)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "ab", "abb///acc///ab c", task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "ab", "abb///acc///ab c", task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 0)
-            self.assertTrue(s[1]['F1'] > 0)
+            self.assertTrue(s[1]["EM"] == 0)
+            self.assertTrue(s[1]["F1"] > 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a b c", "a b b///a c c///", task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "a b c", "a b b///a c c///", task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 0)
-            self.assertTrue(s[1]['F1'] > 0)
+            self.assertTrue(s[1]["EM"] == 0)
+            self.assertTrue(s[1]["F1"] > 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "", "a b b///a c c", task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "", "a b b///a c c", task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 0)
-            self.assertTrue(s[1]['F1'] == 0)
+            self.assertTrue(s[1]["EM"] == 0)
+            self.assertTrue(s[1]["F1"] == 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a", ["a"], task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "a", ["a"], task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 1)
-            self.assertTrue(s[1]['F1'] == 1)
+            self.assertTrue(s[1]["EM"] == 1)
+            self.assertTrue(s[1]["F1"] == 1)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a", ["b"], task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "a", ["b"], task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 0)
-            self.assertTrue(s[1]['F1'] == 0)
+            self.assertTrue(s[1]["EM"] == 0)
+            self.assertTrue(s[1]["F1"] == 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "b", ["a"], task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "b", ["a"], task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 0)
-            self.assertTrue(s[1]['F1'] == 0)
+            self.assertTrue(s[1]["EM"] == 0)
+            self.assertTrue(s[1]["F1"] == 0)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "b", ["b"], task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "b", ["b"], task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 1)
-            self.assertTrue(s[1]['F1'] == 1)
+            self.assertTrue(s[1]["EM"] == 1)
+            self.assertTrue(s[1]["F1"] == 1)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "", [""], task='default')
-        for s in eval.cal_score('emf1'):
+        eval.add_record("input", "", [""], task="default")
+        for s in eval.cal_score("emf1"):
             print(s)
-            self.assertTrue(s[1]['EM'] == 1)
-            self.assertTrue(s[1]['F1'] == 1)
+            self.assertTrue(s[1]["EM"] == 1)
+            self.assertTrue(s[1]["F1"] == 1)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "梵 語", ["梵 語"], task='default')
-        eval.add_record("input", "梵 語", ["梵 語d"], task='default')
-        eval.add_record("input", "梵 語", ["梵 語c"], task='default')
-        for s in eval.cal_score('emf1'):
-            self.assertAlmostEqual(s[1]['EM'], 0.3333333333)
+        eval.add_record("input", "梵 語", ["梵 語"], task="default")
+        eval.add_record("input", "梵 語", ["梵 語d"], task="default")
+        eval.add_record("input", "梵 語", ["梵 語c"], task="default")
+        for s in eval.cal_score("emf1"):
+            self.assertAlmostEqual(s[1]["EM"], 0.3333333333)
 
     def test_tokenize_text(self):
-        tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
-        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer, normalize_text=False)
-        self.assertEqual(eval.tokenize_text("How's this work"), "How's this work")
-        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer, normalize_text=True)
-        self.assertEqual(eval.tokenize_text("How's this work"), "how ' s this work")
+        tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer,
+                                                    normalize_text=False)
+        self.assertEqual(eval.tokenize_text("How's this work"),
+                         "How's this work")
+        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer,
+                                                    normalize_text=True)
+        self.assertEqual(eval.tokenize_text("How's this work"),
+                         "how ' s this work")
 
     @pytest.mark.skip()
     def testNLGWithPAD(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a b", "a b ///a c c", task='default')
-        eval.add_record("input", "d c", "a b c///a c c///d c", task='default')
-        for s in eval.cal_score('nlg'):
+        eval.add_record("input", "a b", "a b ///a c c", task="default")
+        eval.add_record("input", "d c", "a b c///a c c///d c", task="default")
+        for s in eval.cal_score("nlg"):
             s1 = s[1]
             print(s1)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a b", "a b ///a c c", task='nlg')
-        eval.add_record("input", "d c", "a b c///a c c///d c", task='nlg')
-        for s in eval.cal_score('nlg'):
+        eval.add_record("input", "a b", "a b ///a c c", task="nlg")
+        eval.add_record("input", "d c", "a b c///a c c///d c", task="nlg")
+        for s in eval.cal_score("nlg"):
             s2 = s[1]
             print(s2)
 
         self.assertEqual(s1, s2)
-        self.assertAlmostEqual(s2['Bleu_1'], 1)
+        self.assertAlmostEqual(s2["Bleu_1"], 1)
 
     @pytest.mark.skip()
     def testNLG(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "a b c", "a b c///a c c///", task='default')
-        for s in eval.cal_score('nlg'):
-            self.assertAlmostEqual(s[1]['Bleu_1'], 1)
+        eval.add_record("input", "a b c", "a b c///a c c///", task="default")
+        for s in eval.cal_score("nlg"):
+            self.assertAlmostEqual(s[1]["Bleu_1"], 1)
             print(s)
 
         eval1 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval1.add_record("input", "abc", " abc ", task='default')
-        for s1 in eval1.cal_score('nlg'):
-            self.assertAlmostEqual(s1[1]['Bleu_1'], 1)
+        eval1.add_record("input", "abc", " abc ", task="default")
+        for s1 in eval1.cal_score("nlg"):
+            self.assertAlmostEqual(s1[1]["Bleu_1"], 1)
             print(s1)
 
         eval1 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval1.add_record("input", "abb", " abc ", task='default')
-        eval1.add_record("input", "abc", " abc ", task='default')
-        eval1.add_record("input", "abd", " abc ", task='default')
-        for s1 in eval1.cal_score('nlg'):
-            self.assertAlmostEqual(s1[1]['Bleu_1'], 0.3333333)
+        eval1.add_record("input", "abb", " abc ", task="default")
+        eval1.add_record("input", "abc", " abc ", task="default")
+        eval1.add_record("input", "abd", " abc ", task="default")
+        for s1 in eval1.cal_score("nlg"):
+            self.assertAlmostEqual(s1[1]["Bleu_1"], 0.3333333)
             print(s1)
 
         eval3 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval3.add_record("input", "abc ", "abb ///acc/// abc ", task='default')
-        for s3 in eval3.cal_score('nlg'):
-            self.assertAlmostEqual(s3[1]['Bleu_1'], 1)
+        eval3.add_record("input", "abc ", "abb ///acc/// abc ", task="default")
+        for s3 in eval3.cal_score("nlg"):
+            self.assertAlmostEqual(s3[1]["Bleu_1"], 1)
             print(s3)
 
         eval6 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval6.add_record("input", "abc", "abb /// acc ///abc", task='default')
-        for s6 in eval6.cal_score('nlg'):
-            self.assertAlmostEqual(s6[1]['Bleu_1'], 1)
+        eval6.add_record("input", "abc", "abb /// acc ///abc", task="default")
+        for s6 in eval6.cal_score("nlg"):
+            self.assertAlmostEqual(s6[1]["Bleu_1"], 1)
             print(s6)
         self.assertTrue(s1[0] == s3[0] == s6[0])
 
         eval1 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval1.add_record("input", "opq", "abc", task='default')
-        for s1 in eval1.cal_score('nlg'):
-            self.assertAlmostEqual(s1[1]['Bleu_1'], 0)
+        eval1.add_record("input", "opq", "abc", task="default")
+        for s1 in eval1.cal_score("nlg"):
+            self.assertAlmostEqual(s1[1]["Bleu_1"], 0)
             print(s1)
 
         eval3 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval3.add_record("input", "opq", "abb///acc///abc", task='default')
-        for s3 in eval3.cal_score('nlg'):
-            self.assertAlmostEqual(s3[1]['Bleu_1'], 0)
+        eval3.add_record("input", "opq", "abb///acc///abc", task="default")
+        for s3 in eval3.cal_score("nlg"):
+            self.assertAlmostEqual(s3[1]["Bleu_1"], 0)
             print(s3)
 
         eval6 = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval6.add_record("input", "opq", "abb /// acc///abc", task='default')
-        for s6 in eval6.cal_score('nlg'):
-            self.assertAlmostEqual(s6[1]['Bleu_1'], 0)
+        eval6.add_record("input", "opq", "abb /// acc///abc", task="default")
+        for s6 in eval6.cal_score("nlg"):
+            self.assertAlmostEqual(s6[1]["Bleu_1"], 0)
             print(s6)
         self.assertTrue(s1[0] == s3[0] == s6[0])
 
     def testTag(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
-        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        for s in eval.cal_score('classification'):
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
+        eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
     def testClassify(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "abc", "abb///acc///abc", task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record("input", "abc", "abb///acc///abc", task="default")
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "你 好", "我 好///你 好 嗎///好 嗎", task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record("input", "你 好", "我 好///你 好 嗎///好 嗎", task="default")
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "1 3 2", "1 2 3", task='default')
-        eval.add_record("input", "1 3 2", "1 3 3", task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record("input", "1 3 2", "1 2 3", task="default")
+        eval.add_record("input", "1 3 2", "1 3 3", task="default")
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", ["1", "3", "2"], ["1", "2", "3"], task='default')
-        eval.add_record("input", ["1", "3", "2"], ["1", "3", "3"], task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record("input", ["1", "3", "2"], ["1", "2", "3"],
+                        task="default")
+        eval.add_record("input", ["1", "3", "2"], ["1", "3", "3"],
+                        task="default")
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", ['O', 'B_Location', 'I_Location', 'I_Location', 'I_Location', 'I_Location', 'O', 'O'],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", [""] * 7,
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        eval.add_record("input", [],
-                        ['O', 'B_Location', 'I_Location', 'B_Location', 'I_Thing', 'I_Location', 'O'],
-                        task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [
+                "O", "B_Location", "I_Location", "I_Location", "I_Location",
+                "I_Location"
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "I_Location",
+                "O",
+                "O",
+            ],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [""] * 7,
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        eval.add_record(
+            "input",
+            [],
+            [
+                "O",
+                "B_Location",
+                "I_Location",
+                "B_Location",
+                "I_Thing",
+                "I_Location",
+                "O",
+            ],
+            task="default",
+        )
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
             print(s)
 
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", '攝影台', ['攝影台'], task='default')
-        for s in eval.cal_score('classification'):
+        eval.add_record("input", "攝影台", ["攝影台"], task="default")
+        for s in eval.cal_score("classification"):
             print(s[0])
             print(s[1])
 
     @pytest.mark.skip()
     def testNLGOnModel(self):
 
-        tokenizer = AutoTokenizer.from_pretrained('facebook/bart-base')
+        tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "An interesting research into chocolate",
-                        "More Chocolate , Less Health
Chocolate and Blood Pressure
Advice on Eating Chocolate",
-                        task='nlg')
-        for s in eval.cal_score('nlg'):
+        eval.add_record(
+            "input",
+            "An interesting research into chocolate",
+            "More Chocolate , Less Health
Chocolate and Blood Pressure
Advice on Eating Chocolate",
+            task="nlg",
+        )
+        for s in eval.cal_score("nlg"):
             print(s)
 
-        tokenizer = AutoTokenizer.from_pretrained('bert-base-cased')
+        tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
         eval = tfkit.utility.eval_metric.EvalMetric(tokenizer)
-        eval.add_record("input", "An interesting research into chocolate",
-                        "More Chocolate , Less Health///Chocolate and Blood Pressure///Advice on Eating Chocolate",
-                        task='nlg')
-        for s in eval.cal_score('nlg'):
+        eval.add_record(
+            "input",
+            "An interesting research into chocolate",
+            "More Chocolate , Less Health///Chocolate and Blood Pressure///Advice on Eating Chocolate",
+            task="nlg",
+        )
+        for s in eval.cal_score("nlg"):
             print(s)
diff --git a/tfkit/test/utility/test_utility_logger.py b/tfkit/test/utility/test_utility_logger.py
index 92ca3bd..ba0399f 100644
--- a/tfkit/test/utility/test_utility_logger.py
+++ b/tfkit/test/utility/test_utility_logger.py
@@ -1,24 +1,23 @@
 import csv
-import sys
 import os
+import sys
+import unittest
 
+import tfkit
 from tfkit.utility.logger import Logger
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import unittest
-import tfkit
-
 
 class TestLogger(unittest.TestCase):
     ROOT_DIR = os.path.dirname(os.path.abspath(__file__ + "/../../"))
-    MODEL_SAVE_PATH = os.path.join(ROOT_DIR, './test/cache/')
+    MODEL_SAVE_PATH = os.path.join(ROOT_DIR, "./test/cache/")
 
     def test_write_log(self):
         logger = Logger(savedir=self.MODEL_SAVE_PATH)
         logger.write_log("test")
-        with open(logger.logfilepath, 'r') as f:
+        with open(logger.logfilepath, "r") as f:
             lines = f.read().splitlines()
             last_line = lines[-1]
             print(last_line)
@@ -27,6 +26,6 @@ def test_write_log(self):
     def test_write_metric(self):
         logger = Logger(savedir=self.MODEL_SAVE_PATH)
         logger.write_metric("test", 1, 0)
-        with open(logger.metricfilepath, 'r') as f:
+        with open(logger.metricfilepath, "r") as f:
             last_row = list(csv.reader(f))[-1]
-            self.assertEqual(last_row, ["test", '1', '0'])
+            self.assertEqual(last_row, ["test", "1", "0"])
diff --git a/tfkit/test/utility/test_utility_loss.py b/tfkit/test/utility/test_utility_loss.py
index 2f0c348..a0f3843 100644
--- a/tfkit/test/utility/test_utility_loss.py
+++ b/tfkit/test/utility/test_utility_loss.py
@@ -1,63 +1,81 @@
 import os
 import sys
+import unittest
 
 import torch
 from torch import nn
 from torch.autograd import Variable
 
+import tfkit
+
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import unittest
-import tfkit
-
 
 class TestLoss(unittest.TestCase):
-    outputs = Variable(torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]]), requires_grad=False)
+    outputs = Variable(
+        torch.Tensor([[0.00000000000009, 5, 0.5],
+                      [0.00000000000000000001, 69, 9]]),
+        requires_grad=False,
+    )
     targets = Variable(torch.Tensor([1, 1]).long(), requires_grad=False)
     alln_targets = Variable(torch.Tensor([-1, -1]).long(), requires_grad=False)
     onen_targets = Variable(torch.Tensor([1, -1]).long(), requires_grad=False)
 
     def testLabelSmoothingCrossEntropy(self):
-        outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
+        outputs = torch.Tensor([[0.00000000000009, 5, 0.5],
+                                [0.00000000000000000001, 69, 9]])
         targets = torch.Tensor([1, 1]).long()
         alln_targets = torch.Tensor([0, -1]).long()
         onen_targets = torch.Tensor([1, -1]).long()
 
         criterion = nn.CrossEntropyLoss(ignore_index=-1)
-        custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3, ignore_index=-1)
+        custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(
+            3, ignore_index=-1)
 
-        self.assertTrue(criterion(outputs, targets).item() <
-                        custom_criterion(outputs, targets).item())
-        self.assertTrue(criterion(outputs, onen_targets).item() <
-                        custom_criterion(outputs, onen_targets).item())
+        self.assertTrue(
+            criterion(outputs, targets).item() < custom_criterion(
+                outputs, targets).item())
+        self.assertTrue(
+            criterion(outputs, onen_targets).item() < custom_criterion(
+                outputs, onen_targets).item())
 
         criterion = nn.CrossEntropyLoss()
         custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3)
-        self.assertTrue(criterion(outputs, targets).item() <
-                        custom_criterion(outputs, targets).item())
+        self.assertTrue(
+            criterion(outputs, targets).item() < custom_criterion(
+                outputs, targets).item())
 
-        custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(3, reduction='none')
+        custom_criterion = tfkit.utility.loss.LabelSmoothingLoss(
+            3, reduction="none")
         print(custom_criterion(self.outputs, self.targets))
-        self.assertTrue(list(custom_criterion(self.outputs, self.targets).shape) == [2])
+        self.assertTrue(
+            list(custom_criterion(self.outputs, self.targets).shape) == [2])
 
     def testDiceLoss(self):
         custom_criterion = tfkit.utility.loss.DiceLoss(ignore_index=-1)
-        self.assertTrue(0.8 < custom_criterion(self.outputs, self.targets).item() < 1)
-        self.assertTrue(0.99 < custom_criterion(self.outputs, self.alln_targets).item() <= 1)
-        self.assertTrue(0.8 < custom_criterion(self.outputs, self.onen_targets).item() < 1)
+        self.assertTrue(
+            0.8 < custom_criterion(self.outputs, self.targets).item() < 1)
+        self.assertTrue(0.99 < custom_criterion(self.outputs,
+                                                self.alln_targets).item() <= 1)
+        self.assertTrue(
+            0.8 < custom_criterion(self.outputs, self.onen_targets).item() < 1)
 
-        custom_criterion = tfkit.utility.loss.DiceLoss(reduction='none')
+        custom_criterion = tfkit.utility.loss.DiceLoss(reduction="none")
         print(custom_criterion(self.outputs, self.targets))
-        self.assertTrue(list(custom_criterion(self.outputs, self.targets).shape) == [2])
+        self.assertTrue(
+            list(custom_criterion(self.outputs, self.targets).shape) == [2])
 
     def testLossDrop(self):
-        outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
+        outputs = torch.Tensor([[0.00000000000009, 5, 0.5],
+                                [0.00000000000000000001, 69, 9]])
         targets = torch.Tensor([1, 1]).long()
         norm_loss_fct = nn.CrossEntropyLoss(ignore_index=-1)
-        loss_fct = nn.CrossEntropyLoss(reduction='none', ignore_index=-1)  # -1 index = padding token
+        loss_fct = nn.CrossEntropyLoss(
+            reduction="none", ignore_index=-1)  # -1 index = padding token
         masked_lm_loss = loss_fct(outputs, targets)
-        masked_lm_loss = masked_lm_loss.view(-1, len(targets))  # view by batch size
+        # view by batch size
+        masked_lm_loss = masked_lm_loss.view(-1, len(targets))
         masked_lm_loss = masked_lm_loss.sum(dim=0)
         masked_lm_loss = masked_lm_loss.mean()
         print(masked_lm_loss.mean(), norm_loss_fct(outputs, targets).mean())
@@ -67,11 +85,13 @@ def testBCEFocalLoss(self):
         targets = torch.Tensor([[0, 1, 0], [1, 0, 0]])
         criterion = nn.BCELoss()
         custom_criterion = tfkit.utility.loss.BCEFocalLoss()
-        self.assertTrue(criterion(outputs, targets).item() >
-                        custom_criterion(outputs, targets).item())
+        self.assertTrue(
+            criterion(outputs, targets).item() > custom_criterion(
+                outputs, targets).item())
 
     def testNegativeCElLoss(self):
-        outputs = torch.Tensor([[0.00000000000009, 5, 0.5], [0.00000000000000000001, 69, 9]])
+        outputs = torch.Tensor([[0.00000000000009, 5, 0.5],
+                                [0.00000000000000000001, 69, 9]])
         targets = torch.Tensor([1, 1]).long()
         alln_targets = torch.Tensor([-1, -1]).long()
         onen_targets = torch.Tensor([1, -1]).long()
@@ -79,21 +99,35 @@ def testNegativeCElLoss(self):
         criterion = nn.CrossEntropyLoss(ignore_index=-1)
         custom_criterion = tfkit.utility.loss.NegativeCElLoss()
         self.assertTrue(
-            criterion(outputs, targets).item() < custom_criterion(outputs, self.targets).item())
-        self.assertTrue(criterion(outputs, onen_targets).item() < custom_criterion(outputs, onen_targets).item())
+            criterion(outputs, targets).item() < custom_criterion(
+                outputs, self.targets).item())
+        self.assertTrue(
+            criterion(outputs, onen_targets).item() < custom_criterion(
+                outputs, onen_targets).item())
 
     def testFocalLoss(self):
         criterion = nn.CrossEntropyLoss(ignore_index=-1)
         custom_criterion = tfkit.utility.loss.FocalLoss(gamma=0)
-        self.assertAlmostEqual(criterion(self.outputs, self.targets).item(),
-                               custom_criterion(self.outputs, self.targets).item())
-        self.assertAlmostEqual(criterion(self.outputs, self.alln_targets).item(),
-                               custom_criterion(self.outputs, self.alln_targets).item())
-        self.assertAlmostEqual(criterion(self.outputs, self.onen_targets).item(),
-                               custom_criterion(self.outputs, self.onen_targets).item())
+        self.assertAlmostEqual(
+            criterion(self.outputs, self.targets).item(),
+            custom_criterion(self.outputs, self.targets).item(),
+        )
+        self.assertAlmostEqual(
+            criterion(self.outputs, self.alln_targets).item(),
+            custom_criterion(self.outputs, self.alln_targets).item(),
+        )
+        self.assertAlmostEqual(
+            criterion(self.outputs, self.onen_targets).item(),
+            custom_criterion(self.outputs, self.onen_targets).item(),
+        )
 
         custom_criterion = tfkit.utility.loss.FocalLoss(gamma=1)
-        self.assertTrue(criterion(self.outputs, self.targets) > custom_criterion(self.outputs, self.targets))
-        self.assertTrue(criterion(self.outputs, self.alln_targets).item() - custom_criterion(self.outputs,
-                                                                                             self.alln_targets).item() < 1)
-        self.assertTrue(criterion(self.outputs, self.onen_targets) > custom_criterion(self.outputs, self.onen_targets))
+        self.assertTrue(
+            criterion(self.outputs, self.targets) > custom_criterion(
+                self.outputs, self.targets))
+        self.assertTrue(
+            criterion(self.outputs, self.alln_targets).item() -
+            custom_criterion(self.outputs, self.alln_targets).item() < 1)
+        self.assertTrue(
+            criterion(self.outputs, self.onen_targets) > custom_criterion(
+                self.outputs, self.onen_targets))
diff --git a/tfkit/test/utility/test_utility_model.py b/tfkit/test/utility/test_utility_model.py
index 3b33ebf..b4e6fe1 100644
--- a/tfkit/test/utility/test_utility_model.py
+++ b/tfkit/test/utility/test_utility_model.py
@@ -1,41 +1,50 @@
 import os
 import sys
+import unittest
+
+from transformers import AutoModel
+from transformers import BertTokenizer
 
-from tfkit.utility.model import list_all_model, load_model_class, load_predict_parameter, load_trained_model
+from tfkit.utility.model import list_all_model
+from tfkit.utility.model import load_model_class
+from tfkit.utility.model import load_predict_parameter
+from tfkit.utility.model import load_trained_model
 
 dir_path = os.path.dirname(os.path.realpath(__file__))
 sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
-import unittest
-from transformers import BertTokenizer, AutoModel
-
 
 class TestModelLoader(unittest.TestCase):
     ROOT_DIR = os.path.dirname(os.path.abspath(__file__ + "/../../../"))
-    MODEL_SAVE_PATH = os.path.join(ROOT_DIR, 'tfkit/test/cache/')
+    MODEL_SAVE_PATH = os.path.join(ROOT_DIR, "tfkit/test/cache/")
 
     def test_list_all_model(self):
         models = list_all_model()
         self.assertTrue(isinstance(models, list))
 
     def test_load_model_class(self):
-        load_model_class('clas')
-        load_model_class('once')
+        load_model_class("clas")
+        load_model_class("once")
 
     def test_load_predict_parameter(self):
-        model_class = load_model_class('clas')
+        model_class = load_model_class("clas")
         # load pre-train task
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
-        pretrained = AutoModel.from_pretrained('voidful/albert_chinese_tiny')
-        model = model_class.Model(tokenizer=tokenizer, pretrained=pretrained, tasks_detail={"taskA": ["a", "b"]},
-                                  maxlen=128)
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
+        pretrained = AutoModel.from_pretrained("voidful/albert_chinese_tiny")
+        model = model_class.Model(
+            tokenizer=tokenizer,
+            pretrained=pretrained,
+            tasks_detail={"taskA": ["a", "b"]},
+            maxlen=128,
+        )
         clas_param = load_predict_parameter(model)
         print("clas_param", clas_param)
-        self.assertTrue('input' in clas_param)
-        self.assertTrue('topK' in clas_param)
-        self.assertTrue('task' in clas_param)
-        self.assertTrue('handle_exceed' in clas_param)
-        self.assertTrue(isinstance(clas_param['handle_exceed'], str))
+        self.assertTrue("input" in clas_param)
+        self.assertTrue("topK" in clas_param)
+        self.assertTrue("task" in clas_param)
+        self.assertTrue("handle_exceed" in clas_param)
+        self.assertTrue(isinstance(clas_param["handle_exceed"], str))
 
     # def test_load_trained_model(self):
     #     model_path = os.path.join(self.MODEL_SAVE_PATH, '1.pt')
diff --git a/tfkit/test/utility/test_utility_tok.py b/tfkit/test/utility/test_utility_tok.py
index 65cee2a..8acdab0 100644
--- a/tfkit/test/utility/test_utility_tok.py
+++ b/tfkit/test/utility/test_utility_tok.py
@@ -1,20 +1,23 @@
-import sys
 import os
+import sys
+import unittest
 
-dir_path = os.path.dirname(os.path.realpath(__file__))
-sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
+from transformers import AutoTokenizer
+from transformers import BertTokenizer
 
-import unittest
 import tfkit
-from transformers import AutoTokenizer, BertTokenizer
+
+dir_path = os.path.dirname(os.path.realpath(__file__))
+sys.path.append(os.path.abspath(os.path.join(dir_path, os.pardir)))
 
 
 class TestTok(unittest.TestCase):
     ROOT_DIR = os.path.dirname(os.path.abspath(__file__ + "/../../../"))
-    DATASET_DIR = os.path.join(ROOT_DIR, 'demo_data')
+    DATASET_DIR = os.path.join(ROOT_DIR, "demo_data")
 
     def testTok(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
         begin = tfkit.utility.tok.tok_begin(tokenizer)
         self.assertEqual(begin, "[CLS]")
         sep = tfkit.utility.tok.tok_sep(tokenizer)
@@ -25,7 +28,7 @@ def testTok(self):
         self.assertEqual(pad, "[PAD]")
 
     def testTok(self):
-        tokenizer = AutoTokenizer.from_pretrained('distilroberta-base')
+        tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
         begin = tfkit.utility.tok.tok_begin(tokenizer)
         self.assertEqual(begin, "
")
         sep = tfkit.utility.tok.tok_sep(tokenizer)
@@ -36,29 +39,40 @@ def testTok(self):
         self.assertEqual(pad, "<pad>")
 
     def testGetXUnkToken(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
-        result = tfkit.utility.tok.get_topP_unk_token(tokenizer, file_paths=[], topP=0.5)
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
+        result = tfkit.utility.tok.get_topP_unk_token(tokenizer,
+                                                      file_paths=[],
+                                                      topP=0.5)
         self.assertFalse(result)
-        result = tfkit.utility.tok.get_freqK_unk_token(tokenizer, file_paths=[], freqK=10)
+        result = tfkit.utility.tok.get_freqK_unk_token(tokenizer,
+                                                       file_paths=[],
+                                                       freqK=10)
         self.assertFalse(result)
-        result = tfkit.utility.tok.get_freqK_unk_token(tokenizer, file_paths=[self.DATASET_DIR + '/unk_tok.csv'],
-                                                       freqK=1)
+        result = tfkit.utility.tok.get_freqK_unk_token(
+            tokenizer, file_paths=[self.DATASET_DIR + "/unk_tok.csv"], freqK=1)
         self.assertTrue(len(result) > 0)
-        result = tfkit.utility.tok.get_topP_unk_token(tokenizer, file_paths=[self.DATASET_DIR + '/unk_tok.csv'],
-                                                      topP=0.9)
+        result = tfkit.utility.tok.get_topP_unk_token(
+            tokenizer,
+            file_paths=[self.DATASET_DIR + "/unk_tok.csv"],
+            topP=0.9)
         self.assertTrue(len(result) > 0)
 
     def testHandleExceed(self):
-        tokenizer = BertTokenizer.from_pretrained('voidful/albert_chinese_tiny')
+        tokenizer = BertTokenizer.from_pretrained(
+            "voidful/albert_chinese_tiny")
         seq = " ".join([str(_) for _ in range(100)])
         maxlen = 50
-        for mode in ['noop', 'remove', 'slide', 'start_slice', 'end_slice']:
-            rlt, _ = tfkit.utility.tok.handle_exceed(tokenizer, seq, maxlen, mode=mode)
-            if mode == 'remove':
+        for mode in ["noop", "remove", "slide", "start_slice", "end_slice"]:
+            rlt, _ = tfkit.utility.tok.handle_exceed(tokenizer,
+                                                     seq,
+                                                     maxlen,
+                                                     mode=mode)
+            if mode == "remove":
                 self.assertTrue(len(rlt) == 0)
-            if mode == 'slide':
+            if mode == "slide":
                 self.assertTrue(len(rlt) > 1)
             for i in rlt:
                 print(i)
-                if mode != 'noop':
+                if mode != "noop":
                     self.assertTrue(len(i) == 50)