tirank.Model.RiskscorePredictor

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

Bases: Module

Prediction head for ‘Cox’ survival analysis.

Predicts a single risk score, applying a sigmoid activation to constrain the output between 0 and 1.

Parameters:

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

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

  • nhout (int, optional) – Output dimension. Defaults to 1.

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

forward(embedding)[source]

Forward pass for the risk score predictor.

Parameters:

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

Returns:

Predicted risk score (scalar tensor).

Return type:

torch.Tensor