How to add custom model into LIT?

[abhibisht89]

how to add your own NER model and dataset into lit.

Any help ?

Thanks

[jameswex]

See https://github.com/PAIR-code/lit/wiki/api.md#adding-models-and-data for defining your own model and dataset for use in LIT. You just need to write a Model and Dataset wrapper for your model and dataset.

You can see and run code for a demo of the Stanza NER model in LIT here: https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/stanza_demo.py, with the model defined at https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/models/stanza_models.py

[pratikchhapolika]

Hi @jameswex any end to end code to add custom model and dataset for Roberta classification task. Or, are you planning to have one notebook(.ipynb) to get started with custom models and datasets. For beginners like me its little confusing.
Apologies!

[SpamDude15]

Hi @jameswex, how would I add an image classification model? I have a .h5 file of the model, but I'm having a hard time loading it.

I get this error:
AttributeError: 'Sequential' object has no attribute 'init_spec'

Although I load the model in this way which is not similar to other examples since this is an image classification example.

model = tf.keras.models.load_model('model/model.h5')
models = {'model': model}

[jameswex]

@SpamDude15, as per #599, you should create a class for your model that inherits from Model. Inside that class, you can load up your model in whichever way you like in the class constructor, and then define the predict method to use that model to make predictions on the inputs passed to it and return the results. You will also have to define the input_spec and output_spec methods on the class to specify what the inputs to the model are and what the outputs from the model are.

One image model example can be found at https://github.com/PAIR-code/lit/blob/main/lit_nlp/examples/models/mobilenet.py

[RyanMullins]

Hi @jameswex any end to end code to add custom model and dataset for Roberta classification task. Or, are you planning to have one notebook(.ipynb) to get started with custom models and datasets. For beginners like me its little confusing. Apologies!

@pratikchhapolika The LIT team does not have time to implement base classes for all models. If you're interested in RoBERTa, you can access the model from Hugging Face and implement your own Model class. The key things to implement

As for data, LIT provides 2 of the 4 datasets (MNLI and SST-2) used in the RoBERTa paper out of the box. We don't have wrappers for SQuAD (on TFDS or HF Datasets) or RACE (on TFDS or HF Datasets) at this time.

Once you have the Model and Dataset classes of interest, it's relatively simple to initialize a LIT app instance either as a Python module (note I referenced our Penguin demo because it's among the simplest of the built-ins) or in Colab.

For all commenters on this issue (@pratikchhapolika @SpamDude15 @abhibisht89), please remember that if you do implement these classes or demo modules, we would welcome contributions via PRs to incorporate them into the core LIT example.

[SpamDude15]

@jameswex I was wondering if there was a dynamic option to load models and datasets into LIT without having to write three different python files for loading, model, and dataset.

[RyanMullins]

@SpamDude15 not for models. There are too many modeling frameworks, and increasingly serving frameworks, for our research team to keep up with. Instead we've chosen to go the self-service route.

Datasets are slightly easier. We are working on a generic dataset wrapper for TFDS, though I don't have an ETA for its release. That will likely act as a reference for creating a generic wrapper for the HF Datasets library, too, but again, no ETA on release.

There is a Blank Slate demo that's easily forkable. Once you have the model and dataset classes, it should be as simple as importing them into the demo file and launching.

[SpamDude15]

@RyanMullins That's interesting, is there a way I can load in .ppm images? I looked through the datatypes on LIT, and there seems to be no option for ppm, so does that mean it isn't compatible?

[RyanMullins]

@SpamDude15 cant say I'm familiar with PPM images, but they seem like Pillow can read them. LIT passes around all images in base64 encoding (makes it easier to visualize in the UI, and allows flexibility for different input formats to models), so you can use Pillow to read and convert the PPM to base64 and then pass that to your model class or one of the existing image classification models in our examples.

You can also add a new type PPMImage type as a subclass from ImageBytes in Python and TypeScript to further specify that.

If you work up a draft PR (in your own repo or a fork of LIT), I'm happy to review and provide some tips to get this working (in a way we could bring stuff upstream).

[SpamDude15]

@RyanMullins I got PPM images to work with LIT using base64 encoding, which was great, I was wondering if PNG could be loaded in the same way without having to use the defined datatype by LIT.
This is how I load PPM images, this should work with any image type as well right?

image_file_path = Path(base_path+path_name) / row['Filename']
pil_image = PILImage.open(image_file_path)
image_str = image_utils.convert_pil_to_image_str(pil_image)

[RyanMullins]

Yup, that'll work for pretty much any image Pillow supports. The more specificImageBytes subclasses are really just there for people who do pre- or post-processing on images and want a more type-safe way of differentiating code paths.