tirank.Model.ClassscorePredictor

class tirank.Model.ClassscorePredictor(*args: Any, **kwargs: Any)[source]

Bases: Module

Prediction head for ‘Classification’ analysis.

Predicts class probabilities using a Softmax activation.

Parameters:

  • n_features (int) – Input embedding dimension (from encoder).

  • nhid (int) – Hidden dimension of the predictor MLP.

  • nhout (int, optional) – Output dimension (number of classes). Defaults to 2.

  • dropout (float, optional) – Dropout value. Defaults to 0.5.

forward(embedding)[source]

Forward pass for the classification score predictor.

Parameters:

embedding (torch.Tensor) – Input embedding tensor from the encoder.

Returns:

Predicted class probabilities.

Return type:

torch.Tensor