Wrapping everything in a single class for the ease of deployment

[Ajith-Shenoy]

Hey team,
I am very new to lightning , so far I'm very impressed. But the parts that are troublesome for me are dealing with hugging face transformer based models.
The environment I am using is heavily space limited a hugging face model.bin is around 3x more compact than a checkpoint object while saving , and there is no straightforward way of saving the model.bin directly as far as I've seen.

Coming to my core query , can I encapsulate the perks of Lightning module in my own custom classes ?
From a deployment point of view I have a python class created that helps me train , save , and predict on new text data using the hugging face library for "text classification".
It implements methods to :

• encode and decode class labels from integers from a class_label to integer mapping file
• load model and tokenizers
• Implement a linear scheduler with warmup if required, if not without warmup.
• fine tune a model , and display validation metrics after each epoch (needs zero grad ,loss computation ,optimiser step)
• save the model object with the config file , the way huggingface expects, to the directory specified.
• a prediction method that can take external text inputs and give predictions (single examples), if model path is specified.

Inspired from : BERT Fine-Tuning by Chris McCormick and Nick Ryan

Can I use all these utilities methods along with the perks that Lightning provides ? especially the logging and in the training_step of the model ? The mixed Precision training is something I'm a huge fan of.

ideally I would want a class:

Class SuperClassifier(L.Lightningmodule):
# All my utility methods with the goodness of pytorch lightning