From 969bfbec279ce159fdf84f57e6bcdfcb23a6bd44 Mon Sep 17 00:00:00 2001 From: dgisolfi3 Date: Sun, 7 Jun 2026 14:24:04 -0400 Subject: [PATCH 1/2] Started new generate pipeline for running the GAN - Setup image datasets for GAN data, both CIFAR10 and fashion-mnist to validate - DATA_DIR is now in utils as its needed by both a pipeline and data loaders - fixed gitignore so the data subpackage is not excluded now --- .gitignore | 2 +- Makefile | 2 +- belt/__main__.py | 3 ++ belt/data/image.py | 55 ++++++++++++++++++++++++++++++ belt/data/translation.py | 67 +++++++++++++++++++++++-------------- belt/registry.py | 14 +++++++- belt/supervised/generate.py | 29 ++++++++++++++++ belt/utils.py | 8 +++++ 8 files changed, 152 insertions(+), 28 deletions(-) create mode 100644 belt/data/image.py create mode 100644 belt/supervised/generate.py diff --git a/.gitignore b/.gitignore index dee40bc..e681e5b 100644 --- a/.gitignore +++ b/.gitignore @@ -204,7 +204,7 @@ cython_debug/ tempCodeRunnerFile.py # Downloaded datasets and model cache -data/ +/data/ # Ruff stuff: .ruff_cache/ diff --git a/Makefile b/Makefile index 56fea73..57575e1 100644 --- a/Makefile +++ b/Makefile @@ -4,7 +4,7 @@ install: lint: python -m black belt python -m isort belt - python -m ruff check belt tests + python -m ruff check belt --fix train: python -m belt.supervised.softmax --config configs/supervised_softmax.yaml diff --git a/belt/__main__.py b/belt/__main__.py index 381632d..3599302 100644 --- a/belt/__main__.py +++ b/belt/__main__.py @@ -10,16 +10,19 @@ import argparse from belt.supervised.classifier import deploy as softmax_train +from belt.supervised.generate import deploy as generate_train from belt.supervised.translation import deploy as translation_train _DEFAULT_CONFIGS: dict[str, str] = { "classifier": "configs/softmax_classifier.yaml", "translation": "configs/translation_small.yaml", + "generate": "configs/translation_small.yaml", } PIPELINES = { "classifier": softmax_train, "translation": translation_train, + "generate": generate_train, } diff --git a/belt/data/image.py b/belt/data/image.py new file mode 100644 index 0000000..c71e4b1 --- /dev/null +++ b/belt/data/image.py @@ -0,0 +1,55 @@ +"""Image Datasets + +- CIFAR-10 +- Fashion-MNIST + +Author(s) +--------- +Daniel Nicolas Gisolfi +""" + + +from torch.utils.data import DataLoader +from torchvision.datasets import CIFAR10, FashionMNIST + +from belt.data.iris import SupervisedData +from belt.utils import DATA_DIR + +# https://docs.pytorch.org/tutorials/beginner/basics/data_tutorial.html +fashion_mnist_class_map = { + 0: "T-Shirt", + 1: "Trouser", + 2: "Pullover", + 3: "Dress", + 4: "Coat", + 5: "Sandal", + 6: "Shirt", + 7: "Sneaker", + 8: "Bag", + 9: "Ankle Boot", +} +CLASS_NAMES = fashion_mnist_class_map.values() + + +def load_fashion_mnist(batch_size=64, **kwargs): + train = FashionMNIST(root=DATA_DIR, train=True, download=True) + test = FashionMNIST(root=DATA_DIR, train=False, download=True) + + return SupervisedData( + train_loader=DataLoader(train, batch_size=batch_size, shuffle=True), + val_loader=DataLoader(test, batch_size=batch_size, shuffle=False), + test_loader=DataLoader(test, batch_size=batch_size, shuffle=False), + class_names=CLASS_NAMES, + ) + + +def load_cifar10(batch_size=64, **kwargs): + train = CIFAR10(root=DATA_DIR, train=True, download=True) + test = CIFAR10(root=DATA_DIR, train=False, download=True) + + return SupervisedData( + train_loader=DataLoader(train, batch_size=batch_size, shuffle=True), + val_loader=DataLoader(test, batch_size=batch_size, shuffle=False), + test_loader=DataLoader(test, batch_size=batch_size, shuffle=False), + class_names=CLASS_NAMES, + ) diff --git a/belt/data/translation.py b/belt/data/translation.py index 7c0c6cc..0493794 100644 --- a/belt/data/translation.py +++ b/belt/data/translation.py @@ -1,6 +1,6 @@ """Translation Datasets -Use hugging face datasets and tokenizers to create +Use hugging face datasets and tokenizers to create language word pairings for translation task Author(s) @@ -10,15 +10,7 @@ import os from dataclasses import dataclass -from pathlib import Path -DATA_DIR = Path(__file__).parents[2] / "data" -os.environ.setdefault("HF_HOME", str(DATA_DIR)) -if (DATA_DIR / "datasets").exists(): - os.environ.setdefault("HF_DATASETS_OFFLINE", "1") - os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") - -import torch from datasets import load_dataset from torch.utils.data import DataLoader, Dataset from transformers import AutoTokenizer, PreTrainedTokenizerBase @@ -39,6 +31,7 @@ class TranslationData: source_lang: str target_lang: str + class TranslationDataset(Dataset): def __init__(self, pairs, source_tokenizer, target_tokenizer, max_length): self.pairs = pairs @@ -68,22 +61,24 @@ def __getitem__(self, index): return src_ids, tgt_ids -def _translation_loader(pairs, source_tokenizer, target_tokenizer, max_length, batch_size, shuffle): +def _translation_loader( + pairs, source_tokenizer, target_tokenizer, max_length, batch_size, shuffle +): dataset = TranslationDataset(pairs, source_tokenizer, target_tokenizer, max_length) return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle) - def lang_pairs(row, source_lang, target_lang): translation = row.get("translation") if not isinstance(translation, dict): translation = row - - source = translation.get(source_lang) + + source = translation.get(source_lang) target = translation.get(target_lang) - + return str(source), str(target) + def load_translation_data( dataset_name="multi30k", dataset_config=None, @@ -108,18 +103,29 @@ def load_translation_data( os.environ.pop("TRANSFORMERS_OFFLINE", None) model_name = f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}" - source_tokenizer = AutoTokenizer.from_pretrained(model_name, force_download=no_cache) + source_tokenizer = AutoTokenizer.from_pretrained( + model_name, force_download=no_cache + ) # opus-mt is directional, one tokenizer covers both sides target_tokenizer = source_tokenizer if dataset_name not in SUPPORTED_DATASETS: - raise ValueError(f"Unsupported dataset '{dataset_name}'. Choose from: {SUPPORTED_DATASETS}") + raise ValueError( + f"Unsupported dataset '{dataset_name}'. Choose from: {SUPPORTED_DATASETS}" + ) download_mode = "force_redownload" if no_cache else "reuse_dataset_if_exists" dataset_split = f"{split}[:{sample_size}]" if sample_size else split raw = ( - load_dataset(dataset_name, dataset_config, split=dataset_split, download_mode=download_mode) + load_dataset( + dataset_name, + dataset_config, + split=dataset_split, + download_mode=download_mode, + ) if dataset_config - else load_dataset(dataset_name, split=dataset_split, download_mode=download_mode) + else load_dataset( + dataset_name, split=dataset_split, download_mode=download_mode + ) ) if shuffle: raw = raw.shuffle(seed=seed) @@ -139,20 +145,31 @@ def load_translation_data( return TranslationData( train_loader=_translation_loader( - train_pairs, source_tokenizer, target_tokenizer, max_length, batch_size, shuffle=True + train_pairs, + source_tokenizer, + target_tokenizer, + max_length, + batch_size, + shuffle=True, ), val_loader=_translation_loader( - val_pairs, source_tokenizer, target_tokenizer, max_length, batch_size, shuffle=False + val_pairs, + source_tokenizer, + target_tokenizer, + max_length, + batch_size, + shuffle=False, ), test_loader=_translation_loader( - test_pairs, source_tokenizer, target_tokenizer, max_length, batch_size, shuffle=False + test_pairs, + source_tokenizer, + target_tokenizer, + max_length, + batch_size, + shuffle=False, ), source_tokenizer=source_tokenizer, target_tokenizer=target_tokenizer, source_lang=source_lang, target_lang=target_lang, ) - - - - diff --git a/belt/registry.py b/belt/registry.py index 3ecc618..01d09c0 100644 --- a/belt/registry.py +++ b/belt/registry.py @@ -11,6 +11,7 @@ from dataclasses import dataclass from typing import Any +from belt.data.image import load_cifar10, load_fashion_mnist from belt.data.iris import SupervisedData, iris from belt.data.translation import TranslationData, load_translation_data @@ -26,6 +27,14 @@ class DatasetLoaders: "iris": DatasetLoaders( supervised=iris, ), + # Image Data + "fashion_mnist": DatasetLoaders( + supervised=load_fashion_mnist, + ), + "cifar10": DatasetLoaders( + supervised=load_cifar10, + ), + # Translation Data "opus_books_en_fr": DatasetLoaders( translation=load_translation_data, ), @@ -37,7 +46,10 @@ class DatasetLoaders: # Basic Registry for Models and other composable components class Registry(dict[str, type]): - """Maps string names to classes and supports decorator-based registration.""" + """decorator based registration + + used by the pipelines to register available models to run + """ def register(self, name: str): def decorator(cls: type) -> type: diff --git a/belt/supervised/generate.py b/belt/supervised/generate.py new file mode 100644 index 0000000..963d99f --- /dev/null +++ b/belt/supervised/generate.py @@ -0,0 +1,29 @@ +"""Generate Pipeline + +Author(s) +--------- +Daniel Nicolas Gisolfi +""" + + +from belt.supervised.pipeline import SupervisedPipeline + + +class GeneratePipeline(SupervisedPipeline): + """Image Classification and Generation""" + + def __init__(self): + super().__init__() + + def train_batch(self, batch): + pass + + def eval_batch(self, batch): + pass + + def score(self, predictions, targets): + pass + + +def deploy(config_path, overrides=None): + return GeneratePipeline().run(config_path, overrides) diff --git a/belt/utils.py b/belt/utils.py index 3a122ce..d4c704f 100644 --- a/belt/utils.py +++ b/belt/utils.py @@ -7,6 +7,7 @@ import json import logging +import os import random from pathlib import Path from time import perf_counter @@ -23,6 +24,13 @@ logger = logging.getLogger("belt") +DATA_DIR = Path(__file__).parents[2] / "data" +os.environ.setdefault("HF_HOME", str(DATA_DIR)) +if (DATA_DIR / "datasets").exists(): + os.environ.setdefault("HF_DATASETS_OFFLINE", "1") + os.environ.setdefault("TRANSFORMERS_OFFLINE", "1") + + def load_config(path): target = Path(path) if not target.exists(): From 8e80ec54c3a7bddee8c0c180444b82f3ae4b6cc4 Mon Sep 17 00:00:00 2001 From: dgisolfi3 Date: Sun, 7 Jun 2026 14:50:55 -0400 Subject: [PATCH 2/2] Setup common TorchPipeline - The torch pipeline removes duplicate code from various pipelines which all share train and eval functions - Both translation and generate are simplified with only setup needed to be done for loading domain specific data - ran formatters --- README.md | 2 ++ belt/data/image.py | 1 - belt/supervised/generate.py | 62 ++++++++++++++++++++++++---------- belt/supervised/pipeline.py | 59 ++++++++++++++++++++++++++++++-- belt/supervised/translation.py | 48 +++++++------------------- 5 files changed, 114 insertions(+), 58 deletions(-) diff --git a/README.md b/README.md index 33b9041..bdd8d65 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,8 @@ ML Tasks are broken into pipelines allowing models to be used with multiple pipe ```bash # Softmax classifier python -m belt classifier +# Seq2Seq text translation +python -m belt translation ``` Configs live in `configs/`. Each pipeline has a YAML controlling data, model, training, and output settings. diff --git a/belt/data/image.py b/belt/data/image.py index c71e4b1..a5c2252 100644 --- a/belt/data/image.py +++ b/belt/data/image.py @@ -8,7 +8,6 @@ Daniel Nicolas Gisolfi """ - from torch.utils.data import DataLoader from torchvision.datasets import CIFAR10, FashionMNIST diff --git a/belt/supervised/generate.py b/belt/supervised/generate.py index 963d99f..a6e5014 100644 --- a/belt/supervised/generate.py +++ b/belt/supervised/generate.py @@ -5,24 +5,50 @@ Daniel Nicolas Gisolfi """ - -from belt.supervised.pipeline import SupervisedPipeline - - -class GeneratePipeline(SupervisedPipeline): - """Image Classification and Generation""" - - def __init__(self): - super().__init__() - - def train_batch(self, batch): - pass - - def eval_batch(self, batch): - pass - - def score(self, predictions, targets): - pass +import torch +from torch import nn + +from belt.data.image import load_fashion_mnist +from belt.registry import dataset_registry +from belt.supervised.models import supervised_model_registry +from belt.supervised.pipeline import TorchPipeline +from belt.utils import describe_device, get_device + + +class GeneratePipeline(TorchPipeline): + """Supervised image pipeline""" + + def setup(self, config): + self.device = get_device(config.get("device", "auto")) + self.device_label = describe_device(self.device) + + dataset_name = config.get("dataset", "fashion_mnist") + dataset_loader = dataset_registry.get(dataset_name) + if dataset_loader and dataset_loader.supervised: + data = dataset_loader.supervised(**config.get("data", {})) + else: + data = load_fashion_mnist(**config.get("data", {})) + + self.train_loader = data.train_loader + self.val_loader = data.val_loader + self.test_loader = data.test_loader + self.class_names = data.class_names + + self.model = supervised_model_registry.build( + config.get("model_name", "ImageClassifier"), + **config["model"], + num_classes=len(self.class_names), + device=self.device, + ).to(self.device) + self.criterion = nn.CrossEntropyLoss() + self.optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=config["training"]["lr"], + weight_decay=config["training"]["weight_decay"], + ) + + def extra_metrics(self, _): + return {"class_names": self.class_names} def deploy(config_path, overrides=None): diff --git a/belt/supervised/pipeline.py b/belt/supervised/pipeline.py index a29f4c4..24d9b0f 100644 --- a/belt/supervised/pipeline.py +++ b/belt/supervised/pipeline.py @@ -10,6 +10,7 @@ from typing import Any import torch +from sklearn.metrics import accuracy_score from belt.utils import ( ensure_parent, @@ -17,6 +18,7 @@ log_training_end, log_training_start, plant, + save_checkpoint, save_json, ) @@ -25,7 +27,7 @@ class SupervisedPipeline(ABC): """base class for all supervised training pipelines Subclass and override the five abstract methods to add a new supervised - task or model type. The shared run func will drive the training and evaluation loop. + task or model type. The shared run func will drive the training and evaluation loop """ train_loader: Any @@ -58,8 +60,7 @@ def checkpoint_metric(self, record: dict[str, object]) -> float: return float(record.get("val_loss", float("inf"))) def is_better(self, current: float | None, new: float) -> bool: - """Return True when new score is better than current. - Default: lower is better. Override for metrics that should be maximized.""" + """Return True when new score is better than current""" return current is None or new < current def extra_metrics(self, config: dict) -> dict[str, object]: @@ -134,3 +135,55 @@ def _eval_epoch(self, loader: Any) -> tuple[float, list, list]: all_preds.extend(preds) all_targets.extend(targets) return total_loss / max(total_n, 1), all_preds, all_targets + + +class TorchPipeline(SupervisedPipeline): + """base pytorch defaults for a supervised pipeline + + subclasses implement setup() and extra_metrics() if needed + override _compute_loss() or _extract_preds() to avoid duplicating train / eval loop + """ + + _metric_name: str = "accuracy" + + def _compute_loss(self, logits, targets): + return self.criterion(logits, targets) + + def _extract_preds(self, logits, targets): + return logits.argmax(dim=-1).cpu().tolist(), targets.cpu().tolist() + + def train_batch(self, batch): + inputs, targets = batch + inputs, targets = inputs.to(self.device), targets.to(self.device) + self.model.train() + self.optimizer.zero_grad() + logits = self.model(inputs) + loss = self._compute_loss(logits, targets) + loss.backward() + self.optimizer.step() + n = inputs.size(0) + return loss.item() * n, n + + def eval_batch(self, batch): + inputs, targets = batch + inputs, targets = inputs.to(self.device), targets.to(self.device) + self.model.eval() + logits = self.model(inputs) + loss = self._compute_loss(logits, targets) + n = inputs.size(0) + preds, tgts = self._extract_preds(logits, targets) + return loss.item() * n, n, preds, tgts + + def score(self, predictions, targets): + if not targets: + return {self._metric_name: 0.0} + return {self._metric_name: accuracy_score(targets, predictions)} + + def checkpoint_metric(self, record): + return float(record.get(self._metric_name, -float("inf"))) + + def is_better(self, current, new): + return current is None or new > current + + def save_model(self, path): + save_checkpoint(self.model, path) diff --git a/belt/supervised/translation.py b/belt/supervised/translation.py index 95c3f70..71aad35 100644 --- a/belt/supervised/translation.py +++ b/belt/supervised/translation.py @@ -6,14 +6,13 @@ """ import torch -from sklearn.metrics import accuracy_score from torch import nn from belt.data.translation import load_translation_data from belt.registry import dataset_registry from belt.supervised.models import supervised_model_registry -from belt.supervised.pipeline import SupervisedPipeline -from belt.utils import describe_device, get_device, save_checkpoint +from belt.supervised.pipeline import TorchPipeline +from belt.utils import describe_device, get_device def sequence_accuracy(logits, targets, pad_id): @@ -48,8 +47,10 @@ def sample_translations(model, data, device, limit): return samples -class TranslationPipeline(SupervisedPipeline): - """Supervised seq2seq translation pipeline.""" +class TranslationPipeline(TorchPipeline): + """Supervised seq2seq translation pipeline""" + + _metric_name = "token_accuracy" def setup(self, config): self.device = get_device(config.get("device", "auto")) @@ -91,38 +92,13 @@ def setup(self, config): weight_decay=config["training"]["weight_decay"], ) - def train_batch(self, batch): - source, target = batch - source, target = source.to(self.device), target.to(self.device) - self.model.train() - self.optimizer.zero_grad() - logits = self.model(source) - loss = self.criterion(logits.reshape(-1, logits.size(-1)), target.reshape(-1)) - loss.backward() - self.optimizer.step() - n = source.size(0) - return loss.item() * n, n - - def eval_batch(self, batch): - source, target = batch - source, target = source.to(self.device), target.to(self.device) - self.model.eval() - logits = self.model(source) - loss = self.criterion(logits.reshape(-1, logits.size(-1)), target.reshape(-1)) - n = source.size(0) + def _compute_loss(self, logits, targets): + return self.criterion(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) + + def _extract_preds(self, logits, targets): pad_id = self.data.source_tokenizer.pad_token_id - mask = target != pad_id - preds = logits.argmax(dim=-1)[mask].cpu().tolist() - targets = target[mask].cpu().tolist() - return loss.item() * n, n, preds, targets - - def score(self, predictions, targets): - if not targets: - return {"token_accuracy": 0.0} - return {"token_accuracy": accuracy_score(targets, predictions)} - - def save_model(self, path): - save_checkpoint(self.model, path) + mask = targets != pad_id + return logits.argmax(dim=-1)[mask].cpu().tolist(), targets[mask].cpu().tolist() def extra_metrics(self, config): return {