Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/metrics/retrieval/retrieval_metric.py (Python)

Code before:
        self.target = torch.cat([self.target, target])

    def compute(self) -> torch.Tensor:
        res = []
        for group in get_mini_groups(self.idx):
            if self.target[group].sum() == 0:
                if self.query_without_relevant_docs == 'error':
                    raise ValueError(
                        f"`{self.__class__.__name__}.compute()` was provided with "
                        f"a query without positive targets, indexes: {group}"
                    )
                if self.query_without_relevant_docs == 'pos':
                    res.append(torch.tensor(1.0))
                elif self.query_without_relevant_docs == 'neg':
                    res.append(torch.tensor(0.0))
            else:
                res.append(
                    self._metric(group)
                )
        return torch.stack(res).mean()

    def _metric(self, group: List[int]) -> torch.Tensor:
        r"""
        Compute a metric over a single group.
        """
        raise NotImplementedError("This method must be overridden by subclasses")


Review comment (refactor):
If the intention is that the user should not be using this class directly, but only subclass it, then it should be an abstract class and this should be an abstract method.

Code after:

        self.idx.append(idx)
        self.preds.append(preds)
        self.target.append(target)

    def compute(self) -> torch.Tensor:
        r"""
        First concat state `idx`, `preds` and `target` since they were stored as lists. After that,
        compute list of groups that will help in keeping together predictions about the same query.
        Finally, for each group compute the `_metric` if the number of positive targets is at least
        1, otherwise behave as specified by `self.query_without_relevant_docs`.
        """

        idx = torch.cat(self.idx, dim=0)
        preds = torch.cat(self.preds, dim=0)
        target = torch.cat(self.target, dim=0)

        res = []
        kwargs = {'device': idx.device, 'dtype': torch.float32}

        groups = get_group_indexes(idx)
        for group in groups:

            mini_preds = preds[group]
            mini_target = target[group]

            if not mini_target.sum():
                if self.query_without_relevant_docs == 'error':
                    raise ValueError(
                        f"`{self.__class__.__name__}.compute()` was provided with "
                        f"a query without positive targets, indexes: {group}"
                    )
                if self.query_without_relevant_docs == 'pos':
                    res.append(torch.tensor(1.0, **kwargs))
                elif self.query_without_relevant_docs == 'neg':
                    res.append(torch.tensor(0.0, **kwargs))
            else:
                res.append(self._metric(mini_preds, mini_target))

        if len(res) > 0:
            return torch.stack(res).mean()
        return torch.tensor(0.0, **kwargs)

    @abstractmethod
    def _metric(self, preds: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        r"""
        Compute a metric over a predictions and target of a single group.
        This method should be overridden by subclasses.
        """