## SelectiveInferenceTask


Wrap a trained `LightningModule` to attach selection results during inference.


Usage

``` python
SelectiveInferenceTask(
    task,
    score,
    acc_test_outputs=False,
    input_key=None,
    target_key=None,
    rc_metric=None
)
```


The wrapper calls the wrapped model in inference mode and combines its predictions with selection outputs produced by a provided [UncertaintyScore](scores.UncertaintyScore.md#seapig.scores.UncertaintyScore).

Key behavior:

- The wrapped task must provide an `.embed(x)` method. The wrapper calls `task.embed(x)` to produce embeddings used by the score.
- The wrapped task is copied and set to `eval()` during initialization to avoid accidental training side effects.
- If the wrapped task defines `test_metrics` (a `Metric` or `MetricCollection`), it will be wrapped by [SelectiveMetric](SelectiveMetric.md#seapig.SelectiveMetric) so metrics are computed only on selected examples.
- If `rc_metric` (a [RiskCoverageMetric](RiskCoverageMetric.md#seapig.RiskCoverageMetric)) is provided, the wrapper will update it during test steps; the final risk-coverage values are available via [get_risk_coverage_curve()](SelectiveInferenceTask.md#seapig.SelectiveInferenceTask.get_risk_coverage_curve).


## Parameters


`task: LightningModule`  
A trained `LightningModule` whose `forward(x)` returns predictions. The module must implement `embed(x)` to produce embeddings for scoring.

`score: UncertaintyScore`  
A seapig [UncertaintyScore](scores.UncertaintyScore.md#seapig.scores.UncertaintyScore) instance providing the [UncertaintyScore.select](scores.UncertaintyScore.md#seapig.scores.UncertaintyScore.select) method.

`input_key: INPUT_KEYS | None = None`  
Key used to extract inputs from an incoming batch. If `None` (default), the first element of the batch is used (positional index 0). When a string is given it must be one of: `'image'`, `'input'`, `'images'`, `'inputs'`, `'x'`.

`target_key: TARGET_KEYS | None = None`  
Key used to extract targets from an incoming batch. If `None` (default), the second element of the batch is used (positional index 1). When a string is given it must be one of: `'mask'`, `'label'`, `'masks'`, `'labels'`, `'targets'`, `'target'`, `'y'`, `'y_true'`.

`acc_test_outputs: bool = ``False`  
If `True`, per-batch outputs (predictions merged with selection results) are accumulated in the `test_outputs` list for later inspection. If `False` (default), outputs are not accumulated and metrics are logged as usual.

`rc_metric: RiskCoverageMetric | None = None`  
Optional [RiskCoverageMetric](RiskCoverageMetric.md#seapig.RiskCoverageMetric) that will be updated during testing.


## Examples

``` python
from seapig import SelectiveInferenceTask
from seapig.scores import EuclideanScore

score = EuclideanScore()
# score.fit(X=train_embeddings)  # fit before wrapping
selective_task = SelectiveInferenceTask(task=model, score=score)
```


## Methods

| Name | Description |
|----|----|
| [forward()](#forward) | Run the wrapped model and attach selection results. |
| [get_risk_coverage_curve()](#get_risk_coverage_curve) | Return the latest computed risk-coverage curve(s), or None if not available. |
| [on_test_epoch_end()](#on_test_epoch_end) | Log final computed test metrics once at the end of testing. |
| [predict_step()](#predict_step) | Perform prediction and return predictions with selection outputs. |
| [test_step()](#test_step) | Perform a test step and include selection outputs. |

------------------------------------------------------------------------


#### forward()


Run the wrapped model and attach selection results.


Usage

``` python
forward(x)
```


Steps performed:

- Calls the wrapped model. If a `torch.Tensor` is returned it is placed under the key `'predictions'`.
- Computes embeddings with `task.embed(x)` and call `score.select(embs)`.
- Merges prediction mapping and selection mapping and return the result.


##### Returns


`dict[str, torch.Tensor]`  
A `dict` containing the model predictions and the selection outputs returned by the score (`'score'` and `'selected'`).


------------------------------------------------------------------------


#### get_risk_coverage_curve()


Return the latest computed risk-coverage curve(s), or None if not available.


Usage

``` python
get_risk_coverage_curve()
```


------------------------------------------------------------------------


#### on_test_epoch_end()


Log final computed test metrics once at the end of testing.


Usage

``` python
on_test_epoch_end()
```


------------------------------------------------------------------------


#### predict_step()


Perform prediction and return predictions with selection outputs.


Usage

``` python
predict_step(batch, batch_idx, dataloader_idx=0)
```


The wrapper calls `forward(x)` and returns the combined mapping produced by the wrapped model and the score. This mapping typically contains the model's predictions and the selection outputs (e.g. [score](scores.PCAScore.md#seapig.scores.PCAScore.score) and `selected`).


------------------------------------------------------------------------


#### test_step()


Perform a test step and include selection outputs.


Usage

``` python
test_step(batch, batch_idx, dataloader_idx=0)
```


Behavior:

- Extracts inputs and targets from the batch using `input_key`/`target_key`.
- Calls `forward(x)` to get predictions augmented with selection results.
- If a [SelectiveMetric](SelectiveMetric.md#seapig.SelectiveMetric) was created (from the wrapped task's `test_metrics`), it is updated with (`predictions`, `targets`, `selected_mask`) and logged.
- If `rc_metric` is provided, it is updated with (`predictions`, `targets`, [score](scores.PCAScore.md#seapig.scores.PCAScore.score)) and its values are logged; final rc values are logged at [on_test_epoch_end](SelectiveInferenceTask.md#seapig.SelectiveInferenceTask.on_test_epoch_end) step.
- If `test_outputs` was enabled at construction, the per-batch outputs are appended to the `test_outputs` list for later inspection.


##### Notes

This method does not return a value; metrics are updated and logged via `Lightning`'s logging utilities.