import unittest
import torch
import torch.nn as nn
import torch.optim as optim
import importlib
import io
import sys
import pandas as pd


def _dynamic(test, cond, success, fail):
    if cond:
        test._testMethodName = success
        test.assertTrue(True)
    else:
        test._testMethodName = fail
        test.fail(fail)


class TestUserCode(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.df = pd.read_csv(
            "https://content-media-cdn.codefinity.com/courses/1dd2b0f6-6ec0-40e6-a570-ed0ac2209666/section_2/hours_scores.csv"
        )
        cls.X = torch.tensor(cls.df["Hours Studied"].values, dtype=torch.float32).reshape(-1, 1)
        cls.Y = torch.tensor(cls.df["Test Score"].values, dtype=torch.float32).reshape(-1, 1)

    # ---------- INPUT TENSORS ----------
    def test_X_tensor_correct(self):
        import user_code
        cond = hasattr(user_code, "X_tensor")
        if cond:
            cond = torch.allclose(user_code.X_tensor, self.X) and user_code.X_tensor.ndim == 2
        _dynamic(
            self,
            cond,
            "`X_tensor` correctly contains reshaped hour values as a 2D tensor.",
            "`X_tensor` is incorrect. Ensure you convert X to a float tensor and reshape to [-1, 1]."
        )

    def test_Y_tensor_correct(self):
        import user_code
        cond = hasattr(user_code, "Y_tensor")
        if cond:
            cond = torch.allclose(user_code.Y_tensor, self.Y) and user_code.Y_tensor.ndim == 2
        _dynamic(
            self,
            cond,
            "`Y_tensor` correctly contains reshaped score values as a 2D tensor.",
            "`Y_tensor` is incorrect. Ensure you convert Y to a float tensor and reshape to [-1, 1]."
        )

    # ---------- MODEL ----------
    def test_model_is_linear(self):
        import user_code

        cond = isinstance(user_code.model, nn.Linear) and \
               user_code.model.in_features == 1 and user_code.model.out_features == 1

        _dynamic(
            self,
            cond,
            "The model is correctly defined as `nn.Linear(1, 1)`.",
            "The model must be `nn.Linear(1, 1)`."
        )

    # ---------- LOSS FUNCTION ----------
    def test_loss_function(self):
        import user_code

        cond = isinstance(user_code.criterion, nn.MSELoss)

        _dynamic(
            self,
            cond,
            "Loss function correctly set to MSE.",
            "Loss function must be `nn.MSELoss()`."
        )

    # ---------- OPTIMIZER ----------
    def test_optimizer(self):
        import user_code
        cond = isinstance(user_code.optimizer, optim.SGD)
        if cond:
            lr_correct = user_code.optimizer.param_groups[0]["lr"] == 0.01
            cond = cond and lr_correct

        _dynamic(
            self,
            cond,
            "Optimizer correctly set to SGD with learning rate 0.01.",
            "Optimizer must be `optim.SGD(model.parameters(), lr=0.01)`."
        )

    # ---------- TRAINING STEPS ----------
    def test_predictions_exist(self):
        import user_code
        cond = hasattr(user_code, "predictions")
        _dynamic(
            self,
            cond,
            "Predictions computed successfully during training.",
            "Predictions missing. Compute them using `model(X_tensor)`."
        )

    def test_loss_exists(self):
        import user_code
        cond = hasattr(user_code, "loss")
        _dynamic(
            self,
            cond,
            "Loss computed successfully during training.",
            "Loss missing. Compute it using `criterion(predictions, Y_tensor)`."
        )

    def test_weights_and_bias_extracted(self):
        import user_code

        cond = hasattr(user_code, "weights") and hasattr(user_code, "bias")
        if cond:
            cond = isinstance(user_code.weights, torch.Tensor) and isinstance(user_code.bias, torch.Tensor)

        _dynamic(
            self,
            cond,
            "Model parameters (weights and bias) successfully extracted.",
            "You must extract parameters using `model.weight.data` and `model.bias.data`."
        )

    # ---------- PRINT OUTPUT ----------
    def test_print_output(self):
        import user_code
        buffer = io.StringIO()
        sys.stdout = buffer

        importlib.reload(user_code)
        sys.stdout = sys.__stdout__

        cond = len(buffer.getvalue().strip()) > 0

        _dynamic(
            self,
            cond,
            "Training logs and model parameters printed successfully.",
            "No output detected. Make sure to print loss updates and final parameters."
        )


if __name__ == "__main__":
    unittest.main()


test_code.py

Learn the fundamental and advanced concepts needed to work with PyTorch efficiently. Gain a solid understanding tensors, including creation, operations, and reshaping. Explore the essentials of gradients, backpropagation, and linear regression before moving on to handling datasets. Master the skill needed build, train, and evaluate neural networks.

Challenge: Implementing Linear Regression

Solution