ray.train.huggingface.TransformersPredictor#

class ray.train.huggingface.TransformersPredictor(pipeline: Optional[Pipeline] = None, preprocessor: Optional[Preprocessor] = None, use_gpu: bool = False)[source]#

Bases: ray.train.predictor.Predictor

A predictor for HuggingFace Transformers PyTorch models.

This predictor uses Transformers Pipelines for inference.

Parameters

  • pipeline – The Transformers pipeline to use for inference.

  • preprocessor – A preprocessor used to transform data batches prior to prediction.

  • use_gpu – If set, the model will be moved to GPU on instantiation and prediction happens on GPU.

PublicAPI (alpha): This API is in alpha and may change before becoming stable.

Methods

from_checkpoint(checkpoint, *[, ...])

Instantiate the predictor from a Checkpoint.

from_pandas_udf(pandas_udf)

Create a Predictor from a Pandas UDF.

get_preprocessor()

Get the preprocessor to use prior to executing predictions.

predict(data[, feature_columns])

Run inference on data batch.

preferred_batch_format()

Batch format hint for upstream producers to try yielding best block format.

set_preprocessor(preprocessor)

Set the preprocessor to use prior to executing predictions.