<p>The <code>distilbert</code> model in <strong>ktrain</strong> is created using Hugging Face <strong>transformers</strong>, which means you can use that library to prune the model.  See <a href="https://huggingface.co/transformers/bertology.html" rel="nofollow noreferrer">this link</a> for more information and <a href="https://github.com/huggingface/transformers/blob/master/examples/research_projects/bertology/run_bertology.py" rel="nofollow noreferrer">the example script</a>. You may need to convert the model to PyTorch before using the script (in addition to making some modifications to the script itself). The approach is based on the paper <a href="https://arxiv.org/abs/1905.10650" rel="nofollow noreferrer">Are Sixteen Heads Really Better Than One?</a>.</p>
How can I apply pruning on a BERT model?
<p>I have trained a <a href="https://en.wikipedia.org/wiki/BERT_(language_model)" rel="nofollow noreferrer">BERT</a> model using ktrain (TensorFlow wrapper) to recognize emotion on text. It works, but it suffers from really slow inference. That makes my model not suitable for a production environment. I have done some research, and it seems pruning could help.</p>
<p>TensorFlow provides some options for pruning, e.g., <em>tf.contrib.model_pruning</em>. The problem is that it is not a not a widely used technique. What would be a simple enough example that could help me to understand how to use it?</p>
<p>I provide my working code below for reference.</p>
<pre><code>import pandas as pd
import numpy as np
import preprocessor as p
import emoji
import re
import ktrain
from ktrain import text
from unidecode import unidecode
import nltk

# Text preprocessing class
class TextPreprocessing:
    def __init__(self):
        p.set_options(p.OPT.MENTION, p.OPT.URL)

    def _punctuation(self, val):
        val = re.sub(r'[^\w\s]', ' ', val)
        val = re.sub('_', ' ', val)
        return val

    def _whitespace(self, val):
        return &quot; &quot;.join(val.split())

    def _removenumbers(self, val):
        val = re.sub('[0-9] + ', '', val)
        return val

    def _remove_unicode(self, text):
        text = unidecode(text).encode(&quot;ascii&quot;)
        text = str(text, &quot;ascii&quot;)
        return text

    def _split_to_sentences(self, body_text):
        sentences = re.split(r&quot;(?&lt;!\w\.\w.)(?&lt;![A-Z][a-z]\.)(?&lt;=\.|\?)\s&quot;, body_text)
        return sentences

    def _clean_text(self, val):
        val = val.lower()
        val = self._removenumbers(val)
        val = p.clean(val)
        val = ' '.join(self._punctuation(emoji.demojize(val)).split())
        val = self._remove_unicode(val)
        val = self._whitespace(val)
        return val

    def text_preprocessor(self, body_text):

        body_text_df = pd.DataFrame({&quot;body_text&quot;: body_text}, index=[1])

        sentence_split_df = body_text_df.copy()

        sentence_split_df[&quot;body_text&quot;] = sentence_split_df[&quot;body_text&quot;].apply(
            self._split_to_sentences)

        lst_col = &quot;body_text&quot;
        sentence_split_df = pd.DataFrame(
            {
                col: np.repeat(
                    sentence_split_df[col].values, sentence_split_df[lst_col].str.len(
                    )
                )
                for col in sentence_split_df.columns.drop(lst_col)
            }
        ).assign(**{lst_col: np.concatenate(sentence_split_df[lst_col].values)})[
            sentence_split_df.columns
        ]

        body_text_df[&quot;body_text&quot;] = body_text_df[&quot;body_text&quot;].apply(self._clean_text)

        final_df = (
            pd.concat([sentence_split_df, body_text_df])
            .reset_index()
            .drop(columns=[&quot;index&quot;])
        )

        return final_df[&quot;body_text&quot;]

# Instantiate data preprocessing object
text1 = TextPreprocessing()

# Import data
data_train = pd.read_csv('data_train_v5.csv', encoding='utf8', engine='python')
data_test = pd.read_csv('data_test_v5.csv', encoding='utf8', engine='python')

# Clean the data
data_train['Text'] = data_train['Text'].apply(text1._clean_text)
data_test['Text'] = data_test['Text'].apply(text1._clean_text)

X_train = data_train.Text.tolist()
X_test = data_test.Text.tolist()

y_train = data_train.Emotion.tolist()
y_test = data_test.Emotion.tolist()

data = data_train.append(data_test, ignore_index=True)

class_names = ['joy', 'sadness', 'fear', 'anger', 'neutral']

encoding = {
    'joy': 0,
    'sadness': 1,
    'fear': 2,
    'anger': 3,
    'neutral': 4
}

# Integer values for each class
y_train = [encoding[x] for x in y_train]
y_test = [encoding[x] for x in y_test]

trn, val, preproc = text.texts_from_array(x_train=X_train, y_train=y_train,
                                          x_test=X_test, y_test=y_test,
                                          class_names=class_names,
                                          preprocess_mode='distilbert',
                                          maxlen=350)

model = text.text_classifier('distilbert', train_data=trn, preproc=preproc)

learner = ktrain.get_learner(model, train_data=trn, val_data=val, batch_size=6)

predictor = ktrain.get_predictor(learner.model, preproc)

# Save the model on a file for later use
predictor.save(&quot;models/bert_model&quot;)

message = &quot;This is a happy message&quot;

# Cleaning - takes 5 ms to run
clean = text1._clean_text(message)

# Prediction - takes 325 ms to run
predictor.predict_proba(clean)
</code></pre>