ronantakizawa/codereview-bench huggingface dataset:
Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/plugins/ddp_plugin.py (Python)

Code before:
from typing import List

from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel


class DDPPlugin(object):
    """
    Plugin to link a custom ddp implementation to any arbitrary accelerator.

    This plugin forwards all constructor arguments to `LightningDistributedDataParallel`,
    which in turn forwards all args to `DistributedDataParallel`.

    Example::

        class MyDDP(DDPPlugin):

            def configure_ddp(self, model, device_ids):
                model = MyDDPWrapper(model, device_ids)
                return model

        my_ddp = MyDDP()
        trainer = Trainer(accelerator='ddp_x', plugins=[my_ddp])
    """

    def __init__(self, **kwargs):

Review comment (suggestion):
```suggestion
from typing import List, Dict, Any
```

Code after:
from typing import List, Dict, Any

from pytorch_lightning.core.lightning import LightningModule
from pytorch_lightning.overrides.data_parallel import LightningDistributedDataParallel


class DDPPlugin(object):
    """
    Plugin to link a custom ddp implementation to any arbitrary accelerator.

    This plugin forwards all constructor arguments to `LightningDistributedDataParallel`,
    which in turn forwards all args to `DistributedDataParallel`.

    Example::

        class MyDDP(DDPPlugin):

            def configure_ddp(self, model, device_ids):
                model = MyDDPWrapper(model, device_ids)
                return model

        my_ddp = MyDDP()
        trainer = Trainer(accelerator='ddp_x', plugins=[my_ddp])
    """

    def __init__(self, **kwargs):
Repository: Lightning-AI/pytorch-lightning
File: tests/utilities/test_seed.py (Python)

Code before:
import os

import pytest

import pytorch_lightning.utilities.seed as seed_utils


def test_seed_stays_same_with_multiple_seed_everything_calls():
    """
    Ensure that after the initial seed everything,
    the seed stays the same for the same run.
    """
    if "PL_GLOBAL_SEED" in os.environ:
        del os.environ["PL_GLOBAL_SEED"]

    with pytest.warns(UserWarning, match="No correct seed found"):
        seed_utils.seed_everything()
    initial_seed = os.environ.get("PL_GLOBAL_SEED")

    with pytest.warns(None) as record:
        seed_utils.seed_everything()
    assert not record  # does not warn
    seed = os.environ.get("PL_GLOBAL_SEED")

    assert initial_seed == seed
    del os.environ["PL_GLOBAL_SEED"]


def test_correct_seed_with_environment_variable(monkeypatch):
    """
    Ensure that the PL_GLOBAL_SEED environment is read
    """
    if "PL_GLOBAL_SEED" in os.environ:
        del os.environ["PL_GLOBAL_SEED"]
    expected = 2020

Review comment (suggestion):
shall we rather mock the 
```suggestion
@mock.patch.dict(os.environ, {})
def test_seed_stays_same_with_multiple_seed_everything_calls():
```

Code after:
import os

from unittest import mock
import pytest

import pytorch_lightning.utilities.seed as seed_utils


@mock.patch.dict(os.environ, {}, clear=True)
def test_seed_stays_same_with_multiple_seed_everything_calls():
    """
    Ensure that after the initial seed everything,
    the seed stays the same for the same run.
    """
    with pytest.warns(UserWarning, match="No correct seed found"):
        seed_utils.seed_everything()
    initial_seed = os.environ.get("PL_GLOBAL_SEED")

    with pytest.warns(None) as record:
        seed_utils.seed_everything()
    assert not record  # does not warn
    seed = os.environ.get("PL_GLOBAL_SEED")

    assert initial_seed == seed


@mock.patch.dict(os.environ, {"PL_GLOBAL_SEED": "2020"}, clear=True)
def test_correct_seed_with_environment_variable():
    """
    Ensure that the PL_GLOBAL_SEED environment is read
    """
    assert seed_utils.seed_everything() == 2020


@mock.patch.dict(os.environ, {"PL_GLOBAL_SEED": "invalid"}, clear=True)
Repository: Lightning-AI/pytorch-lightning
File: tests/base/boring_model.py (Python)

Code before:
        batch_size: int = 1,
        in_features: int = 32,
        learning_rate: float = 0.1,
        optimizer_name: str = "SGD",
        out_features: int = 2,
    ):
        """
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()
        self.layer = torch.nn.Linear(in_features, out_features)
        self.batch_size = batch_size
        self.in_features = in_features
        self.learning_rate = learning_rate
        self.optimizer_name = optimizer_name
        self.out_features = out_features

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def step(self, x):
        x = self(x)
        out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
        return out

    def training_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def training_epoch_end(self, outputs) -> None:
        torch.stack([x["loss"] for x in outputs]).mean()


Review comment (suggestion):
we shall enforce the laters API, so use `save_hyperparameters()`

Code after:
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))

    def step(self, x):
        x = self(x)
        out = torch.nn.functional.mse_loss(x, torch.ones_like(x))
        return out

    def training_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def training_step_end(self, training_step_outputs):
        return training_step_outputs

    def training_epoch_end(self, outputs) -> None:
        torch.stack([x["loss"] for x in outputs]).mean()

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def validation_epoch_end(self, outputs) -> None:
        torch.stack([x['x'] for x in outputs]).mean()

    def test_step(self, batch, batch_idx):
Repository: Lightning-AI/pytorch-lightning
File: tests/base/boring_model.py (Python)

Code before:
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"loss": loss}

    def training_epoch_end(self, outputs) -> None:
        torch.stack([x["loss"] for x in outputs]).mean()

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def validation_epoch_end(self, outputs) -> None:
        torch.stack([x['x'] for x in outputs]).mean()

    def test_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"y": loss}

    def test_epoch_end(self, outputs) -> None:
        torch.stack([x["y"] for x in outputs]).mean()

    def configure_optimizers(self):
        optimizer_class = getattr(torch.optim, self.optimizer_name)
        optimizer = optimizer_class(self.layer.parameters(), lr=self.learning_rate)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]

    def train_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64), batch_size=self.batch_size)

    def val_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64), batch_size=self.batch_size)

    def test_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64), batch_size=self.batch_size)


Review comment (performance):
```suggestion
        optimizer = optimizer_class(self.parameters(), lr=self.learning_rate)
```

Code after:
        return training_step_outputs

    def training_epoch_end(self, outputs) -> None:
        torch.stack([x["loss"] for x in outputs]).mean()

    def validation_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"x": loss}

    def validation_epoch_end(self, outputs) -> None:
        torch.stack([x['x'] for x in outputs]).mean()

    def test_step(self, batch, batch_idx):
        output = self.layer(batch)
        loss = self.loss(batch, output)
        return {"y": loss}

    def test_epoch_end(self, outputs) -> None:
        torch.stack([x["y"] for x in outputs]).mean()

    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
        lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
        return [optimizer], [lr_scheduler]

    def train_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64))

    def val_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64))

    def test_dataloader(self):
        return torch.utils.data.DataLoader(RandomDataset(32, 64))
Repository: Lightning-AI/pytorch-lightning
File: tests/base/develop_pipelines.py (Python)

Code before:

    else:
        with torch.no_grad():
            y_hat = trained_model(x)
        y_hat = y_hat.cpu()

        # acc
        labels_hat = torch.argmax(y_hat, dim=1)

        y = y.cpu()
        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
        acc = torch.tensor(acc)
        acc = acc.item()

    assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"


@run_prediction.register(BoringModel)
def _(trained_model, dataloader, dp=False, min_acc=0.25):
    # run prediction on 1 batch
    batch = next(iter(dataloader))
    with torch.no_grad():
        output = trained_model(batch)
    acc = trained_model.loss(batch, output)

    assert acc >= min_acc, f"This model is expected to get , {min_acc} in test set (it got {acc})"


Review comment (suggestion):
```suggestion
    assert acc >= min_acc, f"This model is expected to get, {min_acc} in test set but got {acc}"
```

Code after:
        acc = torch.mean(acc).item()

    else:
        with torch.no_grad():
            y_hat = trained_model(x)
        y_hat = y_hat.cpu()

        # acc
        labels_hat = torch.argmax(y_hat, dim=1)

        y = y.cpu()
        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
        acc = torch.tensor(acc)
        acc = acc.item()

    assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"


def _boring_model_run_prediction(trained_model, dataloader, dp=False, min_acc=0.25):
    # run prediction on 1 batch
    batch = next(iter(dataloader))
    with torch.no_grad():
        output = trained_model(batch)
    acc = trained_model.loss(batch, output)

    assert acc >= min_acc, f"This model is expected to get, {min_acc} in test set but got {acc}"
Repository: Lightning-AI/pytorch-lightning
File: tests/models/test_cpu.py (Python)

Code before:


def test_default_logger_callbacks_cpu_model(tmpdir):
    """Test each of the trainer options."""
    trainer_options = dict(
        default_root_dir=tmpdir,
        max_epochs=1,
        gradient_clip_val=1.0,
        overfit_batches=0.20,
        progress_bar_refresh_rate=0,
        limit_train_batches=0.01,
        limit_val_batches=0.01,
    )

    model = BoringModel()
    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.01)

    # test freeze on cpu
    model.freeze()
    model.unfreeze()


def test_running_test_after_fitting(tmpdir):
    """Verify test() on fitted model."""
    class ModelTrainValTest(BoringModel):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)

        def validation_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            return {"x": loss}

        def validation_epoch_end(self, outputs) -> None:
            val_loss = torch.stack([x["x"] for x in outputs]).mean()
            self.log('val_loss', val_loss)

        def test_step(self, batch, batch_idx):
            output = self.layer(batch)
            loss = self.loss(batch, output)
            return {"y": loss}

        def test_epoch_end(self, outputs) -> None:
            test_loss = torch.stack([x["y"] for x in outputs]).mean()
            self.log('test_loss', test_loss)

    model = ModelTrainValTest()

    # logger file to get meta
    logger = tutils.get_default_logger(tmpdir)


Review comment (suggestion):
```suggestion
```
same as what we have in BoringModel.

Code after:


def test_default_logger_callbacks_cpu_model(tmpdir):
    """Test each of the trainer options."""
    trainer_options = dict(
        default_root_dir=tmpdir,
        max_epochs=1,
        gradient_clip_val=1.0,
        overfit_batches=0.20,
        progress_bar_refresh_rate=0,
        limit_train_batches=0.01,
        limit_val_batches=0.01,
    )

    model = BoringModel()
    tpipes.run_model_test_without_loggers(trainer_options, model, min_acc=0.01)

    # test freeze on cpu
    model.freeze()
    model.unfreeze()


def test_running_test_after_fitting(tmpdir):
    """Verify test() on fitted model."""
    class ModelTrainValTest(BoringModel):

        def validation_epoch_end(self, outputs) -> None:
            val_loss = torch.stack([x["x"] for x in outputs]).mean()
            self.log('val_loss', val_loss)

        def test_epoch_end(self, outputs) -> None:
            test_loss = torch.stack([x["y"] for x in outputs]).mean()
            self.log('test_loss', test_loss)

    model = ModelTrainValTest()

    # logger file to get meta
    logger = tutils.get_default_logger(tmpdir)

    # logger file to get weights
    checkpoint = tutils.init_checkpoint_callback(logger)

    # fit model
    trainer = Trainer(
        default_root_dir=tmpdir,
        progress_bar_refresh_rate=0,
        max_epochs=2,
        limit_train_batches=0.4,
        limit_val_batches=0.2,
        limit_test_batches=0.2,
        callbacks=[checkpoint],
Repository: Lightning-AI/pytorch-lightning
File: tests/base/boring_model.py (Python)

Code before:
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return {"id": str(index), "x": self.data[index]}

    def __len__(self):
        return self.len


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self, in_features: int = 32, out_features: int = 2,
                 optimizer_name: str = 'SGD', learning_rate: float = 0.1,
                 batch_size:int = 1):
        """
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()

        self.batch_size = batch_size
        self.learning_rate = learning_rate
        self.optimizer_name = optimizer_name
        self.layer = torch.nn.Linear(in_features, out_features)
        self.save_hyperparameters()


Review comment (suggestion):
I'd still suggest not adding any arguments or any other assignments in BoringModel. Recently I added `example_input_array` which is a lightning thing but William suggested to remove that and instead override Boring Model to keep BoringModel as it is.

Code after:
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return {"id": str(index), "x": self.data[index]}

    def __len__(self):
        return self.len


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):

    def __init__(self):
        """
        Testing PL Module

        Use as follows:
        - subclass
        - modify the behavior for what you want

        class TestModel(BaseTestModel):
            def training_step(...):
                # do your own thing

        or:

        model = BaseTestModel()
        model.training_epoch_end = None

        """
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def loss(self, batch, prediction):
        # An arbitrary loss to have a loss that updates the model weights during `Trainer.fit` calls
        return torch.nn.functional.mse_loss(prediction, torch.ones_like(prediction))
Repository: Lightning-AI/pytorch-lightning
File: tests/base/develop_pipelines.py (Python)

Code before:
    x, y = batch
    x = x.view(x.size(0), -1)

    if dp:
        with torch.no_grad():
            output = trained_model(batch, 0)
            acc = output['val_acc']
        acc = torch.mean(acc).item()

    else:
        with torch.no_grad():
            y_hat = trained_model(x)
        y_hat = y_hat.cpu()

        # acc
        labels_hat = torch.argmax(y_hat, dim=1)

        y = y.cpu()
        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
        acc = torch.tensor(acc)
        acc = acc.item()

    assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"


@run_prediction.register(BoringModel)
def _(trained_model, dataloader, dp=False, min_acc=0.25):
    # run prediction on 1 batch
    batch = next(iter(dataloader))
    with torch.no_grad():
        output = trained_model(batch)
    acc = trained_model.loss(batch, output)

    assert acc >= min_acc, f"This model is expected to get, {min_acc} in test set but got {acc}"


Review comment (suggestion):
I am still not very convinced about this wrapping, I see the beauty but it is not trivial for new users...

Code after:
    x = x.view(x.size(0), -1)

    if dp:
        with torch.no_grad():
            output = trained_model(batch, 0)
            acc = output['val_acc']
        acc = torch.mean(acc).item()

    else:
        with torch.no_grad():
            y_hat = trained_model(x)
        y_hat = y_hat.cpu()

        # acc
        labels_hat = torch.argmax(y_hat, dim=1)

        y = y.cpu()
        acc = torch.sum(y == labels_hat).item() / (len(y) * 1.0)
        acc = torch.tensor(acc)
        acc = acc.item()

    assert acc >= min_acc, f"This model is expected to get > {min_acc} in test set (it got {acc})"


def _boring_model_run_prediction(trained_model, dataloader, dp=False, min_acc=0.25):
    # run prediction on 1 batch
    batch = next(iter(dataloader))
    with torch.no_grad():
        output = trained_model(batch)
    acc = trained_model.loss(batch, output)

    assert acc >= min_acc, f"This model is expected to get, {min_acc} in test set but got {acc}"
Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/metrics/retrieval/retrieval_metric.py (Python)

Code before:
import torch
from typing import List, Optional, Callable, Any

from pytorch_lightning.metrics import Metric

from pytorch_lightning.metrics.utils import get_mini_groups


IGNORE_IDX = -100


class RetrievalMetric(Metric):
    r"""
    Works with binary data. Accepts integer or float predictions from a model output.

    Forward accepts
    - ``indexes`` (long tensor): ``(N, ...)``
    - ``preds`` (long or float tensor): ``(N, ...)`` or ``(N, 2, ...)``
    - ``target`` (long tensor): ``(N, ...)``

    Indexes and target must have the same dimension. If preds has a higher number of dimensions,
    we perform argmax on ``dim=1``. Indexes indicate to which query a prediction belongs.
    Predictions will be first grouped by indexes and then the actual metric will be computed as the mean
    of the scores over each query.

    Args:
        threshold:
            Threshold value for binary predictions. default: 0.5
        compute_on_step:
            Forward only calls ``update()`` and return None if this is set to False. default: True
        dist_sync_on_step:
            Synchronize metric state across processes at each ``forward()``
            before returning the value at the step. default: False
        process_group:
            Specify the process group on which synchronization is called. default: None (which selects

Review comment (suggestion):
```suggestion
#: todo some decribtion here
```
 pls explain what is it for...

Code after:
from abc import ABC, abstractmethod
from typing import Any, Callable, Optional

import torch

from pytorch_lightning.metrics import Metric
from pytorch_lightning.metrics.utils import get_group_indexes

#: get_group_indexes is used to group predictions belonging to the same query

IGNORE_IDX = -100


class RetrievalMetric(Metric, ABC):
    r"""
    Works with binary data. Accepts integer or float predictions from a model output.

    Forward accepts
    - ``indexes`` (long tensor): ``(N, ...)``
    - ``preds`` (float or int tensor): ``(N, ...)``
    - ``target`` (long or bool tensor): ``(N, ...)``

    `indexes`, `preds` and `target` must have the same dimension and will be flatten
    to single dimension once provided.

    `indexes` indicate to which query a prediction belongs.
    Predictions will be first grouped by indexes. Then the
    real metric, defined by overriding the `_metric` method,
    will be computed as the mean of the scores over each query.

    Args:
        query_without_relevant_docs:
            Specify what to do with queries that do not have at least a positive target. Choose from:

            - ``'skip'``: skip those queries (default); if all queries are skipped, ``0.0`` is returned