Skip to content

FastNN Modules

Question Answering Modules

class fastnn.nn.question_answering.BertQAModule(model_name_or_path='distilbert-base-cased-distilled-squad')

Module for Bert-based Question Answering models from Transformers

Usage:

>>> model = BertQAModule(model_name_or_path="distilbert-base-cased-distilled-squad")
>>> output = model(*batch)

Parameters:

  • model_name_or_path - String key of Transformer's pre-trained model hosted in Hugging Face's Model Repository
forward(self, input_ids, attention_mask, token_type_ids, example_indices, *args)

Forward pass on this QA module requires input_ids, attention_mask, and example_indices. The example_indices parameter is required specifically for span-based extractive question anwering model outputs in order to retain the structure of inputs within the final start and end logit calculations.

  • input_ids - Tensor generated from FastNN Processor class
  • attention_mask - Tensor generated from FastNN Processor class
  • token_type_ids - Tensor generated from FastNN Processor class
  • example_indices - Tensor generated from FastNN Processor class
  • *args - args available to make abstraction easier for variant Processor classes

Token Tagging Modules

class fastnn.nn.token_tagging.NERModule(model_name_or_path='dbmdz/bert-large-cased-finetuned-conll03-english')

Module for Token Classification models from Transformers

Usage:

>>> model = NERModule(model_name_or_path="dbmdz/bert-large-cased-finetuned-conll03-english")
>>> output = model(*batch)

Parameters:

  • model_name_or_path - String key of Transformer's pre-trained model hosted in Hugging Face's Model Repository
forward(self, input_ids, attention_mask)

Forward pass on this NER module requires input_ids and attention_mask.

  • input_ids - Tensor generated from FastNN Processor class
  • attention_mask - Tensor generated from FastNN Processor class
  • *args - args available to make abstraction easier for variant Processor classes

Object Detection Modules

class fastnn.nn.object_detection.FasterRCNNModule()

Module for Faster R-CNN Model with ResNet-50 backbone pre-trained on Coco dataset from PyTorch

Usage:

>>> model = FasterRCNNModule()
>>> output = model()

Parameters:

forward(self, images)

Forward pass on this Module only requires a list of dynamic images of shape [C, H, W]. To work with a Triton Server which does not encode Python list objects, we enable a torch.Tensor input that we can trace a model compatible with Triton inference server.

Both inputs are traceable with this model.

If input is list of torch.Tensor then the output will be a List of torch Tensors moddable by 3 where every 3 entries corresponds to one image.

  • images - List of tensor images of shape [C, H, W]
_dict2tensor(self, predictions)