composer.models.huggingface#

A wrapper class that converts ๐Ÿค— Transformers models to composer models

Classes

HuggingFaceModel

A wrapper class that converts ๐Ÿค— Transformers models to composer models.
class composer.models.huggingface.HuggingFaceModel(model, metrics=None)[source]#

Bases: composer.models.base.ComposerModel

A wrapper class that converts ๐Ÿค— Transformers models to composer models.

Parameters

  • model (PreTrainedModel) โ€“ A ๐Ÿค— Transformers model.

  • metrics (list[Metric], optional) โ€“ list of torchmetrics to apply to the output of validate. Default: None.

Warning

This wrapper is designed to work with ๐Ÿค— datasets that define a labels column.

Example:

import transformers
from composer.models import HuggingFaceModel

hf_model = transformers.AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
model = HuggingFaceModel(hf_model)