extract_predictions

probly.integrations.transformers.extract_predictions(output: Any) → Tensor

Extract the prediction tensor from a transformers model output.

Predictions are looked up under the known output fields in PREDICTION_FIELDS, covering classification and regression heads (logits), depth estimation (predicted_depth), time series forecasting (prediction_outputs), and image reconstruction (reconstruction).

Parameters:

output – A transformers ModelOutput or a plain tensor.

Returns:

The prediction tensor.

Raises:

TypeError – If the output has none of the known prediction fields.