This model is a fine-tuned version of
prajjwal1/bert-tiny on the imdb dataset.
It achieves the following results on the evaluation set:
This is the smallest version of BERT model suggested by Google in this
GitHub Repo, this model contains 2 transformer layers and an a hidden layer output length of 128, ie
(L=2, H=128). There are a total 4.39 million paramteres in the model.
This model should be used for text classification tasks specifically on movie reviews or other such text data. Also you can use this model for other downstream tasks like:
This model should not be used for any tasks other than the above mentioned or any language other than English.
1from transformers import pipeline
2
3# load pipeline
4tiny_bert = pipeline("text-classification", "arnabdhar/tinybert-imdb")
5
6# perform inference
7results = pipeline(input_text, truncation=True, max_length=128)
1from transformers import AutoTokenizer, pipeline
2from optimum.onnxruntime import ORTModelForSequenceClassification
3
4# load tokenizer & model
5model_name = "arnabdhar/tinybert-imdb"
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7onnx_model = ORTModelForSequenceClassification.from_pretrained(model_name)
8
9# build pipeline
10tiny_bert_onnx = pipeline(
11 task = "text-classification",
12 tokenizer = tokenizer,
13 model = onnx_model
14)
15
16# perform inference
17results = tiny_bert_onnx(input_text, truncation=True, max_length=128)
The model was finetuned on Google Colab using the NVIDIA V100 GPU and was trained for 9 epochs, it took around 12 minutes to finish finetuning.
This model has been trained on the
imdb dataset which has 25,000 data text data for each training set and testing set, but I have combined both the partitions and then split the dataset in 80:20 ratio and used it for finetuning. This approach gave me a larger dataset to finetune the model.
