Skip to main content

Transformers kit - Multi-task QA/Tagging/Multi-label Multi-Class Classification/Generation with BERT/ALBERT/T5/BERT

Project description

PyPI Download Build Last Commit CodeFactor Visitor

TFKit lets everyone make use of transformer architecture on many tasks and models in small change of config.
At the same time, it can do multi-task multi-model learning, and can introduce its own data sets and tasks through simple modifications.


  • One-click replacement of different pre-trained models
  • Support multi-model and multi-task
  • Classifier with multiple labels and multiple classifications
  • Unify input formats for different tasks
  • Separation of data reading and model architecture
  • Support various loss function and indicators


  • Model list: Support Bert/GPT/GPT2/XLM/XLNet/RoBERTa/CTRL/ALBert/...
  • NLPrep: download and preprocessing data in one line
  • nlp2go: create demo api as quickly as possible.


Learn more from the docs.

Quick Start

Installing via pip

pip install tfkit

Running TFKit to train a ner model

install nlprep and nlp2go

pip install nlprep  nlp2go -U

download dataset using nlprep

nlprep --dataset tag_clner  --outdir ./clner_row --util s2t

train model with albert

tfkit-train --batch 20 \
--epoch 5 \
--lr 5e-5 \
--train ./clner_row/clner-train.csv \
--test ./clner_row/clner-test.csv \
--maxlen 512 \
--model tagRow \
--savedir ./albert_ner \
--config voidful/albert_chinese_small

eval model

tfkit-eval --model ./albert_ner/ --valid ./clner_row/validation.csv --metric clas


Task : default report 
TASK:  default 0
                precision    recall  f1-score   support

    B_Abstract       0.00      0.00      0.00         1
    B_Location       1.00      1.00      1.00         1
      B_Metric       1.00      1.00      1.00         1
B_Organization       0.00      0.00      0.00         1
      B_Person       1.00      1.00      1.00         1
    B_Physical       0.00      0.00      0.00         1
       B_Thing       1.00      1.00      1.00         1
        B_Time       1.00      1.00      1.00         1
    I_Abstract       1.00      1.00      1.00         1
    I_Location       1.00      1.00      1.00         1
      I_Metric       1.00      1.00      1.00         1
I_Organization       0.00      0.00      0.00         1
      I_Person       1.00      1.00      1.00         1
    I_Physical       0.00      0.00      0.00         1
       I_Thing       1.00      1.00      1.00         1
        I_Time       1.00      1.00      1.00         1
             O       1.00      1.00      1.00         1

     micro avg       1.00      0.71      0.83        17
     macro avg       0.71      0.71      0.71        17
  weighted avg       0.71      0.71      0.71        17
   samples avg       1.00      0.71      0.83        17

host prediction service

nlp2go --model ./albert_ner/ --api_path ner

You can also try tfkit in Google Colab: Google Colab



$ tfkit-train
Run training

  --train TRAIN [TRAIN ...]     train dataset path
  --test TEST [TEST ...]        test dataset path
  --config CONFIG               distilbert-base-multilingual-cased/bert-base-multilingual-cased/voidful/albert_chinese_small
  --model {once,twice,onebyone,clas,tagRow,tagCol,qa,onebyone-neg,onebyone-pos,onebyone-both} [{once,twice,onebyone,clas,tagRow,tagCol,qa,onebyone-neg,onebyone-pos,onebyone-both} ...]
                                model task
  --savedir SAVEDIR     model saving dir, default /checkpoints
optional arguments:
  -h, --help            show this help message and exit
  --batch BATCH         batch size, default 20
  --lr LR [LR ...]      learning rate, default 5e-5
  --epoch EPOCH         epoch, default 10
  --maxlen MAXLEN       max tokenized sequence length, default 368
  --lossdrop            loss dropping for text generation
  --tag TAG [TAG ...]   tag to identity task in multi-task
  --seed SEED           random seed, default 609
  --worker WORKER       number of worker on pre-processing, default 8
  --grad_accum          gradient accumulation, default 1
  --tensorboard         Turn on tensorboard graphing
  --resume RESUME       resume training
  --cache               cache training data


$ tfkit-eval
Run evaluation on different benchmark
  --model MODEL             model path
  --metric {emf1,nlg,clas}  evaluate metric
  --valid VALID             evaluate data path

optional arguments:
  -h, --help            show this help message and exit
  --print               print each pair of evaluate data
  --enable_arg_panel    enable panel to input argument


Thanks for your interest.There are many ways to contribute to this project. Get started here.

License PyPI - License

Icons reference

Icons modify from Freepik from
Icons modify from Nikita Golubev from

Project details

Release history Release notifications | RSS feed

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for tfkit, version 0.6.24
Filename, size File type Python version Upload date Hashes
Filename, size tfkit-0.6.24-py3.7.egg (178.3 kB) File type Egg Python version 3.7 Upload date Hashes View
Filename, size tfkit-0.6.24-py3-none-any.whl (80.0 kB) File type Wheel Python version py3 Upload date Hashes View
Filename, size tfkit-0.6.24.tar.gz (228.2 kB) File type Source Python version None Upload date Hashes View

Supported by

Pingdom Pingdom Monitoring Google Google Object Storage and Download Analytics Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN DigiCert DigiCert EV certificate StatusPage StatusPage Status page