Boilerplate for TorchText Driven Deep Learning Research with PyTorch Lightning and Lightning Fabric
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

from setuptools import setup

rootdir = Path(__file__).parent
long_description = (rootdir / "README.md").read_text()

setup(
    name="textlab",
    version="0.0.1",
    package_dir={"": "src"},
    packages=["textlab"],
    description="",
    long_description=long_description,
    long_description_content_type="text/markdown",
    author="",
    license="Apache 2.0",
    author_email="",
    classifiers=[
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Scientific/Engineering :: Information Analysis",
        "Operating System :: OS Independent",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
    ],
)

setup.py
JustinGoheen/text-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from textlab.config import Config

src/textlab/__init__.py
JustinGoheen/text-lab
import os
from pathlib import Path

import typer

FILEPATH = Path(__file__)
PROJECTPATH = FILEPATH.parents[2]
PKGPATH = FILEPATH.parents[1]

app = typer.Typer()
docs_app = typer.Typer()
app.add_typer(docs_app, name="docs")


@app.callback()
def callback() -> None:
    pass


@docs_app.command("build")
def build_docs() -> None:
    import shutil

    os.system("mkdocs build")
    shutil.copyfile(src="README.md", dst="docs/index.md")


@docs_app.command("serve")
def serve_docs() -> None:
    os.system("mkdocs serve")

src/textlab/cli/lab.py
JustinGoheen/text-lab
import os
from pathlib import Path

# SET PATHS
filepath = Path(__file__)


class Config:
    PROJECTPATH = filepath.parents[1]
    LOGSPATH = os.path.join(PROJECTPATH, "logs")
    TORCHPROFILERPATH = os.path.join(LOGSPATH, "torch_profiler")
    SIMPLEPROFILERPATH = os.path.join(LOGSPATH, "simple_profiler")
    CHKPTSPATH = os.path.join(PROJECTPATH, "models", "trials")
    MODELPATH = os.path.join(PROJECTPATH, "models", "production", "model.onnx")
    DATAPATH = os.path.join(PROJECTPATH, "data")
    PREDSPATH = os.path.join(PROJECTPATH, "data", "predictions", "predictions.pt")
    SPLITSPATH = os.path.join(PROJECTPATH, "data", "training_split")
    WANDBPATH = os.path.join(PROJECTPATH, "logs", "wandb_logs")
    OPTUNAPATH = os.path.join(PROJECTPATH, "logs", "optuna")

src/textlab/config.py
JustinGoheen/text-lab
from textlab.core.trainer import LabTrainer
from textlab.core.module import LabModule

src/textlab/core/__init__.py
JustinGoheen/text-lab
"""Demo of a simple transformer language model.

Code is adapted from the PyTorch examples at
https://github.com/pytorch/examples/blob/main/word_language_model

"""
import math
from typing import Any, Optional


import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.nn.modules import MultiheadAttention

import pytorch_lightning as pl

if hasattr(MultiheadAttention, "_reset_parameters") and not hasattr(MultiheadAttention, "reset_parameters"):
    # See https://github.com/pytorch/pytorch/issues/107909
    MultiheadAttention.reset_parameters = MultiheadAttention._reset_parameters


class Transformer(nn.Module):
    def __init__(
        self, vocab_size: int, ninp: int = 200, nhead: int = 2, nhid: int = 200, nlayers: int = 2, dropout: float = 0.2
    ) -> None:
        super().__init__()
        self.pos_encoder = PositionalEncoding(ninp, dropout)
        self.embedding = nn.Embedding(vocab_size, ninp)
        self.transformer = nn.Transformer(
            d_model=ninp,
            nhead=nhead,
            num_encoder_layers=nlayers,
            num_decoder_layers=nlayers,
            dim_feedforward=nhid,
            dropout=dropout,
            batch_first=True,
        )
        self.decoder = nn.Linear(ninp, vocab_size)

        self.ninp = ninp
        self.vocab_size = vocab_size
        self.src_mask = None

    def forward(self, inputs: Tensor, target: Tensor, mask: Optional[Tensor] = None) -> Tensor:
        b, t = inputs.shape

        # we assume target is already shifted w.r.t. inputs
        if mask is None:
            mask = torch.tril(torch.ones(t, t, device=inputs.device)) == 1
            mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, float(0.0))

        src = self.pos_encoder(self.embedding(inputs) * math.sqrt(self.ninp))
        target = self.pos_encoder(self.embedding(target) * math.sqrt(self.ninp))
        output = self.transformer(src, target, tgt_mask=mask)
        output = self.decoder(output)
        output = F.log_softmax(output, dim=-1)
        output = output.view(-1, self.vocab_size)
        return output


class PositionalEncoding(nn.Module):
    def __init__(self, dim: int, dropout: float = 0.1, max_len: int = 5000) -> None:
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.dim = dim
        self.max_len = max_len

        pe = self._init_pos_encoding()
        # workaround, can't use buffer, see https://github.com/pytorch/pytorch/issues/68407
        self.register_parameter("pe", nn.Parameter(pe, requires_grad=False))

    def reset_parameters(self) -> None:
        self.pe.copy_(self._init_pos_encoding())  # type: ignore[operator]

    def forward(self, x: Tensor) -> Tensor:
        x + self.pe[: x.size(0), :]  # type: ignore[index]
        return self.dropout(x)

    def _init_pos_encoding(self) -> Tensor:
        pe = torch.zeros(self.max_len, self.dim)
        position = torch.arange(0, self.max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, self.dim, 2).float() * (-math.log(10000.0) / self.dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        return pe


class LabModule(pl.LightningModule):
    def __init__(self, vocab_size: int = 33278):
        super().__init__()
        self.model = Transformer(vocab_size=vocab_size)

    def forward(self, inputs, target):
        return self.model(inputs, target)

    def training_step(self, batch, batch_idx):
        inputs, target = batch
        output = self(inputs, target)
        loss = torch.nn.functional.nll_loss(output, target.view(-1))
        return loss

    def validation_step(self, batch, batch_idx):
        inputs, target = batch
        output = self(inputs, target)
        loss = torch.nn.functional.nll_loss(output, target.view(-1))
        self.log("val-loss", loss)

    def test_step(self, batch, batch_idx):
        inputs, target = batch
        output = self(inputs, target)
        loss = torch.nn.functional.nll_loss(output, target.view(-1))
        self.log("test-loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.model.parameters(), lr=0.1)

src/textlab/core/module.py
JustinGoheen/text-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import torch
import pytorch_lightning as pl
from pytorch_lightning import seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning.loggers import Logger, TensorBoardLogger
from pytorch_lightning.profilers import Profiler, PyTorchProfiler

from textlab import config


class LabTrainer(pl.Trainer):
    def __init__(
        self,
        logger: Optional[Logger] = None,
        profiler: Optional[Profiler] = None,
        callbacks: Optional[List] = [],
        plugins: Optional[List] = [],
        set_seed: bool = True,
        **trainer_init_kwargs: Dict[str, Any]
    ) -> None:
        # SET SEED
        if set_seed:
            seed_everything(config.GLOBALSEED, workers=True)
        super().__init__(
            logger=logger or TensorBoardLogger(config.LOGSPATH, name="tensorboard"),
            profiler=profiler or PyTorchProfiler(dirpath=config.TORCHPROFILERPATH, filename="profiler"),
            callbacks=callbacks + [ModelCheckpoint(dirpath=config.CHKPTSPATH, filename="model")],
            plugins=plugins,
            **trainer_init_kwargs
        )

    def persist_predictions(self, predictions_dir: Optional[Union[str, Path]] = config.PREDSPATH) -> None:
        self.test(ckpt_path="best", datamodule=self.datamodule)
        predictions = self.predict(self.model, self.datamodule.val_dataloader())
        torch.save(predictions, predictions_dir)

src/textlab/core/trainer.py
JustinGoheen/text-lab
from pytorch_lightning.demos.transformer import WikiText2 as LabDataset  # noqa
from textlab.pipeline.datamodule import LabDataModule

src/textlab/pipeline/__init__.py
JustinGoheen/text-lab
import multiprocessing
from pathlib import Path

from torch.utils.data import DataLoader, random_split

import pytorch_lightning as pl
from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS

from textlab import Config
from textlab.pipeline import LabDataset

NUMWORKERS = int(multiprocessing.cpu_count() // 2)


class LabDataModule(pl.LightningDataModule):
    def __init__(
        self,
        num_workers: int = NUMWORKERS,
        data_dir: Path = Config.DATAPATH,
        block_size: int = 35,
        download: bool = True,
        train_size: float = 0.8,
    ) -> None:
        super().__init__()
        self.data_dir = data_dir
        self.block_size = block_size
        self.download = download
        self.num_workers = num_workers
        self.train_size = train_size
        self.dataset = None

    def prepare_data(self) -> None:
        self.dataset = LabDataset(data_dir=self.data_dir, block_size=self.block_size, download=self.download)

    def setup(self, stage: str) -> None:
        if stage == "fit" or stage is None:
            train_size = int(len(self.dataset) * self.train_size)
            test_size = len(self.dataset) - train_size
            self.train_data, self.val_data = random_split(self.dataset, lengths=[train_size, test_size])
        if stage == "test" or stage is None:
            self.test_data = self.val_data

    def train_dataloader(self) -> TRAIN_DATALOADERS:
        return DataLoader(self.train_data, num_workers=self.num_workers)

    def val_dataloader(self) -> EVAL_DATALOADERS:
        return DataLoader(self.val_data, num_workers=self.num_workers)

    def test_dataloader(self) -> EVAL_DATALOADERS:
        return DataLoader(self.test_data, num_workers=self.num_workers)

src/textlab/pipeline/datamodule.py
JustinGoheen/text-lab
"""Demo of a simple transformer language model.

Code is adapted from the PyTorch examples at
https://github.com/pytorch/examples/blob/main/word_language_model

"""
import os
from pathlib import Path
from typing import Dict, List, Tuple

import requests
import torch
from torch import Tensor

from torch.utils.data import Dataset


class WikiText2(Dataset):
    """Mini version of WikiText2."""

    def __init__(self, data_dir: Path = Path("./data"), block_size: int = 35, download: bool = True) -> None:
        super().__init__()
        self.path = data_dir / "wikitext-2.txt"
        if download:
            self.download(self.path)
        self.data, self.dictionary = tokenize(self.path)
        self.block_size = block_size

    @property
    def vocab_size(self) -> int:
        return len(self.dictionary)

    def __len__(self) -> int:
        return len(self.data) // self.block_size - 1

    def __getitem__(self, index: int) -> Tuple[Tensor, Tensor]:
        start = index * self.block_size
        end = start + self.block_size
        inputs = self.data[start:end]
        target = self.data[(start + 1) : (end + 1)]
        return inputs, target

    @staticmethod
    def download(destination: Path) -> None:
        os.makedirs(destination.parent, exist_ok=True)
        url = "https://raw.githubusercontent.com/pytorch/examples/main/word_language_model/data/wikitext-2/train.txt"
        if os.path.exists(destination):
            return
        with open(destination, "w") as f:
            f.write(requests.get(url).text)


class Dictionary:
    def __init__(self) -> None:
        self.word2idx: Dict[str, int] = {}
        self.idx2word: List[str] = []

    def add_word(self, word: str) -> int:
        if word not in self.word2idx:
            self.idx2word.append(word)
            self.word2idx[word] = len(self.idx2word) - 1
        return self.word2idx[word]

    def __len__(self) -> int:
        return len(self.idx2word)


def tokenize(path: Path) -> Tuple[Tensor, Dictionary]:
    dictionary = Dictionary()

    assert os.path.exists(path)
    # Add words to the dictionary
    with open(path, encoding="utf8") as f:
        for line in f:
            words = line.split() + ["<eos>"]
            for word in words:
                dictionary.add_word(word)

    # Tokenize file content
    with open(path, encoding="utf8") as f:
        idss: List[Tensor] = []
        for line in f:
            words = line.split() + ["<eos>"]
            ids: List[int] = []
            for word in words:
                ids.append(dictionary.word2idx[word])
            idss.append(torch.tensor(ids).type(torch.int64))

    return torch.cat(idss), dictionary

src/textlab/pipeline/dataset.py
JustinGoheen/text-lab
from pathlib import Path

import streamlit as st

# create the title
file = Path(__file__)
rootdir = file.parents[1]
long_description = (rootdir / "README.md").read_text()
st.markdown(body=long_description)

streamlit/app.py
JustinGoheen/text-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from pathlib import Path

import torch

from textlab.pipeline.datamodule import LabDataModule


def test_module_not_abstract():
    _ = LabDataModule()


def test_prepare_data():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    networkpath = Path(__file__).parent
    projectpath = networkpath.parents[0]
    datapath = os.path.join(projectpath, "data", "cache")
    assert "LabDataset" in os.listdir(datapath)


def test_setup():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    data_keys = ["train_data", "test_data", "val_data"]
    assert all(key in dir(datamodule) for key in data_keys)


def test_trainloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.train_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)


def test_testloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.test_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)


def test_valloader():
    datamodule = LabDataModule()
    datamodule.prepare_data()
    datamodule.setup()
    loader = datamodule.val_dataloader()
    sample = loader.dataset[0][0]
    assert isinstance(sample, torch.Tensor)

tests/test_datamodule.py
JustinGoheen/text-lab
# Copyright Justin R. Goheen.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch

from textlab.core.module import LabModule


def test_module_not_abstract():
    _ = LabModule()


def test_module_forward():
    input_sample = torch.randn((1, 784))
    model = LabModule()
    preds, label = model.forward(input_sample)
    assert preds.shape == input_sample.shape


def test_module_training_step():
    input_sample = torch.randn((1, 784)), 1
    model = LabModule()
    loss = model.training_step(input_sample)
    assert isinstance(loss, torch.Tensor)


def test_optimizer():
    model = LabModule()
    optimizer = model.configure_optimizers()
    optimizer_base_class = optimizer.__class__.__base__.__name__
    assert optimizer_base_class == "Optimizer"

tests/test_network_module.py
JustinGoheen/text-lab