# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/nateraw/hf-hub-lightning/blob/main/examples/hf_hub_lightning_demo.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="CiutL6CAWcT7"
# %%capture
# ! pip install hf-hub-lightning
# ! apt install git-lfs
# ! git config --global credential.helper store

# + id="LMAf7Y4WWylI"
# ! huggingface-cli login

# + id="U6bMY62HWfry"
import torch
from huggingface_hub import hf_hub_download
from pytorch_lightning import LightningModule, Trainer
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import MNIST

from hf_hub_lightning import HuggingFaceHubCallback


# + id="MZdwu2thWlPZ"
class MNISTModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)

    def forward(self, x):
        return torch.relu(self.l1(x.view(x.size(0), -1)))

    def training_step(self, batch, batch_idx):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        self.log('loss', loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


# + id="FEHT5l-Rtd0E"
model_name = 'nateraw/lit-mnist-demo'

# + id="G3WPOLBHWoEe"
# Init our model
mnist_model = MNISTModel()

# Init DataLoader from MNIST Dataset
train_ds = MNIST('.', train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(train_ds, batch_size=64)

# Initialize a trainer
trainer = Trainer(
    gpus=torch.cuda.device_count(),
    max_epochs=3,
    callbacks=[HuggingFaceHubCallback(model_name)],
)

# Train the model ⚡
trainer.fit(mnist_model, train_loader)

# + id="7UVUngAItaAt"
reloaded_model = MNISTModel.load_from_checkpoint(hf_hub_download(model_name, 'lit_model.ckpt'))

StarcoderdataJupyterScriptsDedupFiltered