dashai-frankenstein 0.1.0__tar.gz

dashai_frankenstein-0.1.0/.gitignore

@@ -0,0 +1,36 @@
+.venv/
+uv.lock
+wandb/
+output/
+dist/
+**/__pycache__/
+**/*.pyc
+**/temp_data/
+**/logs/
+
+**/checkpoints
+**/es_redpajama_50k.model
+**/es_redpajama_50k.vocab
+**/nohup.out**
+**/training_metrics.csv
+
+docs/**/*.aux
+docs/**/*.log
+docs/**/*.out
+docs/**/*.toc
+docs/**/*.bbl
+docs/**/*.blg
+docs/_build/
+docs/pdoc/
+
+*.csv
+
+.opencode/node_modules/
+
+# Exhaustive end-to-end harness workspace (generated tokenizers, checkpoints, deploys)
+full_tests/tmp/
+full_tests/runs/
+full_tests/**/__pycache__/
+full_tests/**/*.pyc
+full_tests/**/*.model
+full_tests/**/*.vocab

dashai_frankenstein-0.1.0/PKG-INFO

@@ -0,0 +1,61 @@
+Metadata-Version: 2.5
+Name: dashai-frankenstein
+Version: 0.1.0
+Summary: DashAI plugin that registers Frankenstein Transformer model classes (MLM encoder, causal decoder, ViT) as DashAI components.
+Author-email: Erick Merino <erickfmm@gmail.com>
+License: MIT
+Requires-Python: >=3.10
+Requires-Dist: datasets>=3.0.0
+Requires-Dist: frankenstein-transformer>=1.1.0
+Requires-Dist: numpy<3.0,>=1.26.0
+Requires-Dist: pyyaml>=6.0.1
+Requires-Dist: torch>=2.0
+Requires-Dist: transformers<5.0.0,>=4.45.0
+Provides-Extra: sbert
+Requires-Dist: sentence-transformers>=3.3.0; extra == 'sbert'
+Requires-Dist: sentencepiece>=0.2.0; extra == 'sbert'
+Description-Content-Type: text/markdown
+
+# dashai-frankenstein
+
+A [DashAI](https://docs.dash-ai.com/) plugin that registers
+[Frankenstein Transformer](https://github.com/erickfmm/frankenstein-transformer)
+model classes as DashAI components, so end users can train, evaluate, predict,
+save, and load them from the DashAI UI.
+
+## Components registered
+
+| Entry point | Class | DashAI base | Binds to task |
+|---|---|---|---|
+| `frankenstein_mlm` | `FrankensteinMLMModel` | `BaseModel` | `TextClassificationTask` |
+| `frankenstein_decoder` | `FrankensteinDecoderModel` | `BaseGenerativeModel` | `TextToTextGenerationTask` |
+| `frankenstein_vit_cls` | `FrankensteinViTClassifier` | `BaseModel` | `ImageClassificationTask` |
+| `frankenstein_vit_seg` | `FrankensteinViTSegmenter` | `BaseModel` | `SegmentationTask` |
+| `segmentation_task` | `SegmentationTask` | `BaseTask` | (new task provided by this plugin) |
+
+## Schema (v1: passthrough YAML)
+
+Each model exposes a minimal pydantic schema. The primary field,
+`frankenstein_yaml`, is a string containing a full Frankenstein training YAML,
+validated by Frankenstein's own config loader + JSON Schema (the Frankenstein
+schema remains the single source of truth). A `preset` dropdown is populated
+from Frankenstein's bundled `configs/*.yaml`; `device`, `batch_size`, and
+`num_epochs` are convenience overrides merged into the YAML before validation.
+
+## Install
+
+```bash
+pip install dashai-frankenstein            # from PyPI once published
+# or, from this repo:
+pip install -e ./dashai-frankenstein
+```
+
+DashAI discovers the plugin via the `dashai.plugins` entry-points group on
+startup — no DashAI source edits required.
+
+## Architecture
+
+See `docs/dashai-plugin-audit.md` in the Frankenstein repo for the full
+integration design (§5 component designs, §6 phased plan, §7 Frankenstein
+changes). This package is the Phase 1–3 adapter layer; it consumes the
+Frankenstein engine API (`src.engine`) added in Phase 0.

dashai_frankenstein-0.1.0/README.md

@@ -0,0 +1,43 @@
+# dashai-frankenstein
+
+A [DashAI](https://docs.dash-ai.com/) plugin that registers
+[Frankenstein Transformer](https://github.com/erickfmm/frankenstein-transformer)
+model classes as DashAI components, so end users can train, evaluate, predict,
+save, and load them from the DashAI UI.
+
+## Components registered
+
+| Entry point | Class | DashAI base | Binds to task |
+|---|---|---|---|
+| `frankenstein_mlm` | `FrankensteinMLMModel` | `BaseModel` | `TextClassificationTask` |
+| `frankenstein_decoder` | `FrankensteinDecoderModel` | `BaseGenerativeModel` | `TextToTextGenerationTask` |
+| `frankenstein_vit_cls` | `FrankensteinViTClassifier` | `BaseModel` | `ImageClassificationTask` |
+| `frankenstein_vit_seg` | `FrankensteinViTSegmenter` | `BaseModel` | `SegmentationTask` |
+| `segmentation_task` | `SegmentationTask` | `BaseTask` | (new task provided by this plugin) |
+
+## Schema (v1: passthrough YAML)
+
+Each model exposes a minimal pydantic schema. The primary field,
+`frankenstein_yaml`, is a string containing a full Frankenstein training YAML,
+validated by Frankenstein's own config loader + JSON Schema (the Frankenstein
+schema remains the single source of truth). A `preset` dropdown is populated
+from Frankenstein's bundled `configs/*.yaml`; `device`, `batch_size`, and
+`num_epochs` are convenience overrides merged into the YAML before validation.
+
+## Install
+
+```bash
+pip install dashai-frankenstein            # from PyPI once published
+# or, from this repo:
+pip install -e ./dashai-frankenstein
+```
+
+DashAI discovers the plugin via the `dashai.plugins` entry-points group on
+startup — no DashAI source edits required.
+
+## Architecture
+
+See `docs/dashai-plugin-audit.md` in the Frankenstein repo for the full
+integration design (§5 component designs, §6 phased plan, §7 Frankenstein
+changes). This package is the Phase 1–3 adapter layer; it consumes the
+Frankenstein engine API (`src.engine`) added in Phase 0.

dashai_frankenstein-0.1.0/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "dashai-frankenstein"
+version = "0.1.0"
+description = "DashAI plugin that registers Frankenstein Transformer model classes (MLM encoder, causal decoder, ViT) as DashAI components."
+readme = "README.md"
+license = { text = "MIT" }
+authors = [{ name = "Erick Merino", email = "erickfmm@gmail.com" }]
+requires-python = ">=3.10"
+dependencies = [
+    # Frankenstein ships the importable ``src`` package and the engine API.
+    # Pin to the Phase-0 release that exposes engine.build_model / train_from_config.
+    "frankenstein-transformer>=1.1.0",
+    "torch>=2.0",
+    "transformers>=4.45.0,<5.0.0",
+    "datasets>=3.0.0",
+    "numpy>=1.26.0,<3.0",
+    "PyYAML>=6.0.1",
+]
+
+[project.optional-dependencies]
+sbert = ["sentence-transformers>=3.3.0", "sentencepiece>=0.2.0"]
+
+# DashAI discovers plugins whose distribution name startswith("dashai").
+# Each entry points to a class that inherits a DashAI base (BaseModel /
+# BaseGenerativeModel / BaseTask) carrying a TYPE attribute.
+[project.entry-points."dashai.plugins"]
+frankenstein_mlm = "dashai_frankenstein:FrankensteinMLMModel"
+frankenstein_decoder = "dashai_frankenstein:FrankensteinDecoderModel"
+frankenstein_vit_cls = "dashai_frankenstein:FrankensteinViTClassifier"
+frankenstein_vit_seg = "dashai_frankenstein:FrankensteinViTSegmenter"
+segmentation_task = "dashai_frankenstein:SegmentationTask"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/dashai_frankenstein"]

dashai_frankenstein-0.1.0/src/dashai_frankenstein/__init__.py

@@ -0,0 +1,20 @@
+"""dashai-frankenstein — DashAI plugin for Frankenstein Transformer.
+
+Registers Frankenstein model classes (encoder, decoder, ViT classifier/segmenter)
+and a SegmentationTask as DashAI components, discovered via the
+``dashai.plugins`` entry-points group.
+"""
+from dashai_frankenstein.models.decoder import FrankensteinDecoderModel  # noqa: F401
+from dashai_frankenstein.models.mlm import FrankensteinMLMModel  # noqa: F401
+from dashai_frankenstein.models.vit_classifier import FrankensteinViTClassifier  # noqa: F401
+from dashai_frankenstein.models.vit_segmenter import FrankensteinViTSegmenter  # noqa: F401
+from dashai_frankenstein.tasks.segmentation import SegmentationTask  # noqa: F401
+
+__all__ = [
+    "FrankensteinMLMModel",
+    "FrankensteinDecoderModel",
+    "FrankensteinViTClassifier",
+    "FrankensteinViTSegmenter",
+    "SegmentationTask",
+]
+__version__ = "0.1.0"

dashai_frankenstein-0.1.0/src/dashai_frankenstein/adapters/__init__.py

@@ -0,0 +1,4 @@
+"""DashAI <-> Frankenstein adapters (dataset, IO, metrics)."""
+from dashai_frankenstein.adapters import dataset, io, metrics  # noqa: F401
+
+__all__ = ["dataset", "io", "metrics"]

dashai_frankenstein-0.1.0/src/dashai_frankenstein/adapters/dataset.py

@@ -0,0 +1,281 @@
+"""Dataset adapter: DashAIDataset <-> Frankenstein training tensors.
+
+A ``DashAIDataset`` is a HuggingFace ``datasets.Dataset`` wrapper. Frankenstein
+models consume tensors directly. This module extracts the relevant columns,
+tokenizes text (for NLP components), and yields PyTorch ``DataLoader`` batches
+ready for the in-process training loop.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any, List, Optional, Tuple
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+
+def find_text_column(dataset: Any, exclude: Optional[List[str]] = None) -> str:
+    """Return the single non-categorical text column of a DashAIDataset.
+
+    Parameters
+    ----------
+    dataset : DashAIDataset
+        A dataset whose ``column_names`` and ``types`` are available.
+    exclude : list of str, optional
+        Column names to ignore (e.g. the label column).
+
+    Returns
+    -------
+    str
+        The name of the text column.
+
+    Raises
+    ------
+    ValueError
+        If there is not exactly one text column.
+    """
+    from DashAI.back.types.categorical import Categorical
+
+    exclude = set(exclude or [])
+    try:
+        types = dataset.types
+    except AttributeError:
+        types = {}
+    text_cols = [
+        col
+        for col in dataset.column_names
+        if col not in exclude and not isinstance(types.get(col), Categorical)
+    ]
+    if len(text_cols) != 1:
+        raise ValueError(
+            f"Expected exactly one text column, found {text_cols} "
+            f"(columns={list(dataset.column_names)}, excluded={sorted(exclude)})."
+        )
+    return text_cols[0]
+
+
+def extract_label_column(dataset: Any) -> str:
+    """Return the (single) categorical/output column name."""
+    from DashAI.back.types.categorical import Categorical
+
+    try:
+        types = dataset.types
+    except AttributeError:
+        types = {}
+    cat_cols = [
+        col for col in dataset.column_names if isinstance(types.get(col), Categorical)
+    ]
+    if len(cat_cols) != 1:
+        raise ValueError(
+            f"Expected exactly one categorical label column, found {cat_cols}."
+        )
+    return cat_cols[0]
+
+
+def tokenized_dataloader(
+    dataset: Any,
+    tokenizer: Any,
+    text_column: str,
+    label_column: str,
+    *,
+    batch_size: int = 16,
+    max_length: int = 512,
+    device: str = "cpu",
+    shuffle: bool = True,
+) -> Any:
+    """Build a torch DataLoader of tokenized (input_ids, attention_mask, labels).
+
+    Parameters
+    ----------
+    dataset : DashAIDataset
+        Source dataset.
+    tokenizer : Any
+        HF tokenizer (or Frankenstein SPM tokenizer exposing ``__call__``
+        returning ``input_ids``).
+    text_column : str
+        Name of the text column.
+    label_column : str
+        Name of the integer-label column.
+    batch_size : int
+        Batch size.
+    max_length : int
+        Max token length.
+    device : str
+        Torch device (for pinning).
+    shuffle : bool
+        Whether to shuffle.
+
+    Returns
+    -------
+    torch.utils.data.DataLoader
+        Yields dicts with ``input_ids``, ``attention_mask``, ``labels`` tensors.
+    """
+    import torch
+    from torch.utils.data import DataLoader, TensorDataset
+
+    texts = list(dataset[text_column])
+    labels = list(dataset[label_column])
+
+    enc = tokenizer(texts, truncation=True, padding=True, max_length=max_length)
+    input_ids = torch.tensor(enc["input_ids"], dtype=torch.long)
+    attention_mask = torch.tensor(enc["attention_mask"], dtype=torch.long)
+    label_tensor = torch.tensor(np.asarray(labels).astype("int64"), dtype=torch.long)
+
+    ds = TensorDataset(input_ids, attention_mask, label_tensor)
+    return DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        pin_memory=str(device).startswith("cuda"),
+    )
+
+
+def prediction_loader(
+    dataset: Any,
+    tokenizer: Any,
+    text_column: str,
+    *,
+    batch_size: int = 32,
+    max_length: int = 512,
+    device: str = "cpu",
+) -> Any:
+    """Build a DataLoader for inference (no labels required)."""
+    import torch
+    from torch.utils.data import DataLoader, TensorDataset
+
+    texts = list(dataset[text_column])
+    enc = tokenizer(texts, truncation=True, padding=True, max_length=max_length)
+    input_ids = torch.tensor(enc["input_ids"], dtype=torch.long)
+    attention_mask = torch.tensor(enc["attention_mask"], dtype=torch.long)
+    ds = TensorDataset(input_ids, attention_mask)
+    return DataLoader(
+        ds,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=str(device).startswith("cuda"),
+    )
+
+
+# ---------------------------------------------------------------------------
+# Vision: DashAI image columns -> (pixel_values, labels) tensors
+# ---------------------------------------------------------------------------
+
+def _image_column(dataset: Any) -> str:
+    """Return the single image column of a DashAIDataset (DashAIImage type)."""
+    try:
+        from DashAI.back.types.dashai_image import DashAIImage
+    except ImportError:  # pragma: no cover
+        DashAIImage = ()  # type: ignore
+    try:
+        types = dataset.types
+    except AttributeError:
+        types = {}
+    img_cols = [
+        col
+        for col in dataset.column_names
+        if (DashAIImage and isinstance(types.get(col), DashAIImage))
+        or str(types.get(col)).lower() == "dashaiimage"
+    ]
+    if len(img_cols) != 1:
+        raise ValueError(
+            f"Expected exactly one image column, found {img_cols} "
+            f"(columns={list(dataset.column_names)})."
+        )
+    return img_cols[0]
+
+
+def image_dataloader(
+    dataset: Any,
+    y_dataset: Any = None,
+    *,
+    image_size: int = 224,
+    batch_size: int = 32,
+    device: str = "cpu",
+    shuffle: bool = True,
+    label_column: Optional[str] = None,
+) -> Any:
+    """Build a DataLoader of (pixel_values, labels) or pixel_values tensors.
+
+    Images are resized to ``image_size`` x ``image_size`` and normalized with
+    ImageNet statistics (matching the torchvision DashAI classifiers).
+
+    Parameters
+    ----------
+    dataset : DashAIDataset
+        Source dataset carrying an image column.
+    y_dataset : DashAIDataset, optional
+        Label dataset (train mode). When ``None``, the loader yields images only.
+    image_size : int
+        Target square image size.
+    batch_size, device, shuffle
+        Loader options.
+    label_column : str, optional
+        Override label column name in ``y_dataset``.
+
+    Returns
+    -------
+    torch.utils.data.DataLoader
+    """
+    import torch
+    import torch.utils.data
+    from torchvision import transforms
+
+    image_col = _image_column(dataset)
+    label_col = label_column
+    if y_dataset is not None and label_col is None:
+        label_col = y_dataset.column_names[0]
+
+    label_to_idx: dict = {}
+    if y_dataset is not None and label_col is not None:
+        cat = (getattr(y_dataset, "types", {}) or {}).get(label_col)
+        if cat is not None and getattr(cat, "categories", None):
+            unique_labels = sorted(cat.categories)
+        else:
+            unique_labels = sorted(set(y_dataset[label_col]))
+        label_to_idx = {lbl: i for i, lbl in enumerate(unique_labels)}
+
+    transform = transforms.Compose(
+        [
+            transforms.Lambda(lambda img: img.convert("RGB")),
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]
+    )
+
+    class _ImageDataset(torch.utils.data.Dataset):
+        def __init__(self):
+            self.x = dataset
+            self.y = y_dataset
+
+        def __len__(self):
+            return len(self.x)
+
+        def __getitem__(self, idx):
+            image = transform(self.x[idx][image_col].to_pil())
+            if self.y is None:
+                return image
+            label_str = self.y[idx][label_col]
+            return image, int(label_to_idx.get(label_str, -1))
+
+    def _collate_with_labels(batch):
+        images = torch.stack([b[0] for b in batch])
+        labels = torch.tensor([b[1] for b in batch], dtype=torch.long)
+        return images, labels
+
+    def _collate_images(batch):
+        return torch.stack(batch)
+
+    ds_obj = _ImageDataset()
+    return (
+        torch.utils.data.DataLoader(
+            ds_obj,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            collate_fn=_collate_with_labels if y_dataset is not None else _collate_images,
+            pin_memory=str(device).startswith("cuda"),
+        ),
+        label_to_idx,
+        len(label_to_idx),
+    )

dashai_frankenstein-0.1.0/src/dashai_frankenstein/adapters/io.py

@@ -0,0 +1,70 @@
+"""IO adapter: DashAI run directory <-> Frankenstein checkpoint bundle.
+
+Each DashAI run persists artifacts under a run-specific directory. This module
+bridges the DashAI ``save(filename)``/``load(filename)`` contract to the
+Frankenstein engine's :func:`save_checkpoint`/:func:`load_checkpoint`, which
+write a self-contained bundle (``model.pt`` + ``config.yaml`` + ``tokenizer/``
++ ``dashai_meta.json``).
+"""
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, Optional, Tuple
+
+
+def save_run(
+    filename: str,
+    model: Any,
+    loaded_config: Any,
+    tokenizer: Any,
+    *,
+    extra: Optional[Dict[str, Any]] = None,
+) -> str:
+    """Persist a Frankenstein model bundle into the DashAI run directory.
+
+    Parameters
+    ----------
+    filename : str
+        DashAI run directory (created if missing).
+    model : torch.nn.Module
+        The trained Frankenstein model.
+    loaded_config : LoadedTrainingConfig
+        The validated Frankenstein config (round-tripped as config.yaml).
+    tokenizer : Any
+        Tokenizer persisted alongside the weights.
+    extra : dict, optional
+        Extra metadata (e.g. ``{"num_labels": N, "label_encodings": {...}}``).
+
+    Returns
+    -------
+    str
+        Path to the saved ``model.pt``.
+    """
+    from dashai_frankenstein.engine import save_checkpoint
+
+    os.makedirs(filename, exist_ok=True)
+    path = save_checkpoint(filename, model, loaded_config, tokenizer, extra=extra)
+
+    # Drop a small DashAI-facing marker so the load path can sanity-check.
+    marker = {"plugin": "dashai-frankenstein", "version": "0.1.0"}
+    with open(os.path.join(filename, "plugin_marker.json"), "w", encoding="utf-8") as fh:
+        json.dump(marker, fh, indent=2)
+    return path
+
+
+def load_run(filename: str) -> Tuple[Any, Any, Any, Dict[str, Any]]:
+    """Reload a bundle saved by :func:`save_run`.
+
+    Rebuilds the ``nn.Module`` from ``config.yaml`` via the Frankenstein engine
+    and restores weights. ``extra`` (including ``num_labels``) is applied to the
+    config before building so the classification head is reconstructed.
+
+    Returns
+    -------
+    tuple
+        ``(model, loaded_config, tokenizer, extra)``.
+    """
+    from dashai_frankenstein.engine import load_checkpoint
+
+    return load_checkpoint(filename)

dashai_frankenstein-0.1.0/src/dashai_frankenstein/adapters/metrics.py

@@ -0,0 +1,74 @@
+"""Metrics adapter: stream per-epoch metrics into DashAI's metric store.
+
+DashAI models expose ``calculate_metrics(split, level, x_data, y_data, ...)``,
+which calls ``self.predict(x_data)`` and writes scores to the database. This
+module provides a tiny callback the training loop invokes at the end of each
+epoch, mirroring the HuggingFace ``MetricsCallback`` pattern.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any, Optional
+
+from DashAI.back.core.enums.metrics import LevelEnum, SplitEnum
+
+log = logging.getLogger(__name__)
+
+
+class EpochMetricsHook:
+    """Invoke ``model_instance.calculate_metrics`` once per epoch.
+
+    Parameters
+    ----------
+    model_instance : BaseModel
+        The DashAI model component (carries ``run_id``, ``train_metrics`` ...).
+    x_train, y_train : DashAIDataset
+        Training split.
+    x_val, y_val : DashAIDataset, optional
+        Validation split. When ``None``, validation metrics are skipped.
+    log_every_n_epochs : int
+        Log frequency (epochs). ``1`` logs every epoch.
+    """
+
+    def __init__(
+        self,
+        model_instance: Any,
+        x_train: Any,
+        y_train: Any,
+        x_val: Any = None,
+        y_val: Any = None,
+        *,
+        log_every_n_epochs: int = 1,
+    ) -> None:
+        self.model_instance = model_instance
+        self.x_train = x_train
+        self.y_train = y_train
+        self.x_val = x_val
+        self.y_val = y_val
+        self.log_every_n_epochs = max(1, int(log_every_n_epochs))
+        self._last_epoch = -1
+
+    def __call__(self, epoch: int, **_: Any) -> None:
+        if epoch <= self._last_epoch:
+            return
+        self._last_epoch = epoch
+        if epoch % self.log_every_n_epochs != 0:
+            return
+        try:
+            self.model_instance.calculate_metrics(
+                split=SplitEnum.TRAIN,
+                level=LevelEnum.EPOCH,
+                x_data=self.x_train,
+                y_data=self.y_train,
+                log_index=epoch,
+            )
+            if self.x_val is not None and self.y_val is not None:
+                self.model_instance.calculate_metrics(
+                    split=SplitEnum.VALIDATION,
+                    level=LevelEnum.EPOCH,
+                    x_data=self.x_val,
+                    y_data=self.y_val,
+                    log_index=epoch,
+                )
+        except Exception as exc:  # noqa: BLE001
+            log.warning("calculate_metrics failed at epoch %s: %s", epoch, exc)