FastNN Modules

Question Answering Modules

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

Source code in fastnn/nn/question_answering.py

class BertQAModule(torch.nn.Module):
    """Module for Bert-based Question Answering models from Transformers

    Usage:
    ```python
    >>> 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
    """

    def __init__(
        self, model_name_or_path: str = "distilbert-base-cased-distilled-squad"
    ):
        super(BertQAModule, self).__init__()
        self.model = AutoModelForQuestionAnswering.from_pretrained(
            model_name_or_path, torchscript=True
        )

    def forward(
        self, input_ids, attention_mask, token_type_ids, example_indices, *args
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """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
        """
        start_logits, end_logits = self.model(
            input_ids=input_ids, attention_mask=attention_mask
        )
        return start_logits, end_logits, example_indices

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

Source code in fastnn/nn/question_answering.py

def forward(
    self, input_ids, attention_mask, token_type_ids, example_indices, *args
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """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
    """
    start_logits, end_logits = self.model(
        input_ids=input_ids, attention_mask=attention_mask
    )
    return start_logits, end_logits, example_indices

Token Tagging Modules

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

Source code in fastnn/nn/token_tagging.py

class NERModule(torch.nn.Module):
    """Module for Token Classification models from Transformers

    Usage:
    ```python
    >>> 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
    """

    def __init__(
        self,
        model_name_or_path: str = "dbmdz/bert-large-cased-finetuned-conll03-english",
    ):
        super(NERModule, self).__init__()
        self.model = AutoModelForTokenClassification.from_pretrained(
            model_name_or_path, torchscript=True
        )

    def forward(
        self, input_ids, attention_mask
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """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
        """
        logits = self.model(input_ids=input_ids, attention_mask=attention_mask)
        return logits[0], input_ids

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

Source code in fastnn/nn/token_tagging.py

def forward(
    self, input_ids, attention_mask
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """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
    """
    logits = self.model(input_ids=input_ids, attention_mask=attention_mask)
    return logits[0], input_ids

Object Detection Modules

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

Usage:

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

Parameters:

Source code in fastnn/nn/object_detection.py

class FasterRCNNModule(torch.nn.Module):
    """Module for Faster R-CNN Model with ResNet-50 backbone pre-trained on Coco dataset from PyTorch

    Usage:
    ```python
    >>> model = FasterRCNNModule()
    >>> output = model()
    ```

    **Parameters:**

    """

    def __init__(self):
        super(FasterRCNNModule, self).__init__()
        self.model = fasterrcnn_resnet50_fpn(pretrained=True)

    def forward(
        self, images: Union[torch.Tensor, List[torch.Tensor]]
    ) -> Union[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], List[torch.Tensor]]:
        """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]
        """
        if isinstance(images, torch.Tensor):
            predictions = self.model([images])
            return (
                predictions[0]["boxes"],
                predictions[0]["labels"],
                predictions[0]["scores"],
            )
        else:
            predictions = self.model(images)
            return self._dict2tensor(predictions)

    def _dict2tensor(self, predictions):
        prediction_outputs = []
        for p in predictions:
            if "masks" in p.keys():
                prediction_outputs += [p["boxes"], p["labels"], p["scores"], p["masks"]]
            prediction_outputs += [p["boxes"], p["labels"], p["scores"]]
        return prediction_outputs

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]

Source code in fastnn/nn/object_detection.py

def forward(
    self, images: Union[torch.Tensor, List[torch.Tensor]]
) -> Union[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], List[torch.Tensor]]:
    """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]
    """
    if isinstance(images, torch.Tensor):
        predictions = self.model([images])
        return (
            predictions[0]["boxes"],
            predictions[0]["labels"],
            predictions[0]["scores"],
        )
    else:
        predictions = self.model(images)
        return self._dict2tensor(predictions)