tensorless 0.1.0__py3-none-any.whl

tensorless/data/tabular.py

@@ -0,0 +1,179 @@
+"""Tabular preprocessing.
+
+Fits simple, fully-reversible preprocessing on tabular records:
+  - numeric columns -> standardized (mean/std), missing values imputed
+    with the training-set mean
+  - categorical columns -> integer-indexed vocabulary (+ <unk>/<missing>),
+    fed into per-column embeddings by the MLP model
+
+The fitted state is small and JSON-serializable, so it can be embedded
+directly inside a `.tl` file and reproduced exactly at inference time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+_MISSING = "<missing>"
+_UNK = "<unk>"
+
+
+def _try_float(v: Any) -> Optional[float]:
+    if v is None:
+        return None
+    if isinstance(v, str) and v.strip() == "":
+        return None
+    try:
+        return float(v)
+    except (TypeError, ValueError):
+        return None
+
+
+@dataclass
+class ColumnStats:
+    kind: str  # "numeric" | "categorical"
+    mean: float = 0.0
+    std: float = 1.0
+    vocab: List[str] = field(default_factory=list)
+
+
+@dataclass
+class TabularPreprocessor:
+    feature_columns: List[str] = field(default_factory=list)
+    target_column: Optional[str] = None
+    task: str = "classification"
+    column_stats: Dict[str, ColumnStats] = field(default_factory=dict)
+    classes: List[str] = field(default_factory=list)  # for classification targets
+    target_mean: float = 0.0
+    target_std: float = 1.0
+
+    @property
+    def numeric_columns(self) -> List[str]:
+        return [c for c in self.feature_columns if self.column_stats[c].kind == "numeric"]
+
+    @property
+    def categorical_columns(self) -> List[str]:
+        return [c for c in self.feature_columns if self.column_stats[c].kind == "categorical"]
+
+    def fit(
+        self,
+        records: List[Dict[str, Any]],
+        columns: List[str],
+        target_column: str,
+        task: str,
+    ) -> "TabularPreprocessor":
+        self.target_column = target_column
+        self.task = task
+        self.feature_columns = [c for c in columns if c != target_column]
+
+        for col in self.feature_columns:
+            values = [r.get(col) for r in records]
+            numeric_vals = [_try_float(v) for v in values]
+            n_present = sum(1 for v in values if v not in (None, ""))
+            n_numeric = sum(1 for v in numeric_vals if v is not None)
+            if n_present > 0 and n_numeric / n_present > 0.95:
+                nums = [v for v in numeric_vals if v is not None]
+                mean = sum(nums) / len(nums) if nums else 0.0
+                var = sum((x - mean) ** 2 for x in nums) / len(nums) if nums else 1.0
+                std = max(var ** 0.5, 1e-6)
+                self.column_stats[col] = ColumnStats(kind="numeric", mean=mean, std=std)
+            else:
+                cats = sorted({str(v) for v in values if v not in (None, "")})
+                vocab = [_MISSING, _UNK] + cats
+                self.column_stats[col] = ColumnStats(kind="categorical", vocab=vocab)
+
+        target_vals = [r.get(target_column) for r in records]
+        if task == "regression":
+            nums = [v for v in (_try_float(v) for v in target_vals) if v is not None]
+            self.target_mean = sum(nums) / len(nums) if nums else 0.0
+            var = sum((x - self.target_mean) ** 2 for x in nums) / len(nums) if nums else 1.0
+            self.target_std = max(var ** 0.5, 1e-6)
+        else:
+            self.classes = sorted({str(v) for v in target_vals if v not in (None, "")})
+
+        return self
+
+    def transform(
+        self, records: List[Dict[str, Any]], with_target: bool = True
+    ) -> Dict[str, torch.Tensor]:
+        n = len(records)
+        num_cols = self.numeric_columns
+        cat_cols = self.categorical_columns
+
+        numeric = torch.zeros(n, max(len(num_cols), 1), dtype=torch.float32)
+        categorical = torch.zeros(n, max(len(cat_cols), 1), dtype=torch.long)
+
+        for i, r in enumerate(records):
+            for j, col in enumerate(num_cols):
+                stats = self.column_stats[col]
+                v = _try_float(r.get(col))
+                if v is None:
+                    v = stats.mean
+                numeric[i, j] = (v - stats.mean) / stats.std
+            for j, col in enumerate(cat_cols):
+                stats = self.column_stats[col]
+                raw = r.get(col)
+                key = _MISSING if raw in (None, "") else str(raw)
+                idx = stats.vocab.index(key) if key in stats.vocab else stats.vocab.index(_UNK)
+                categorical[i, j] = idx
+
+        out = {"numeric": numeric, "categorical": categorical}
+
+        if with_target and self.target_column is not None:
+            if self.task == "regression":
+                target = torch.zeros(n, dtype=torch.float32)
+                for i, r in enumerate(records):
+                    v = _try_float(r.get(self.target_column))
+                    v = self.target_mean if v is None else v
+                    target[i] = (v - self.target_mean) / self.target_std
+                out["target"] = target
+            else:
+                target = torch.zeros(n, dtype=torch.long)
+                for i, r in enumerate(records):
+                    raw = str(r.get(self.target_column))
+                    idx = self.classes.index(raw) if raw in self.classes else 0
+                    target[i] = idx
+                out["target"] = target
+
+        return out
+
+    def inverse_target(self, values: torch.Tensor) -> List[Any]:
+        if self.task == "regression":
+            return [(v.item() * self.target_std + self.target_mean) for v in values]
+        return [self.classes[int(v.item())] for v in values]
+
+    def categorical_vocab_sizes(self) -> List[int]:
+        return [len(self.column_stats[c].vocab) for c in self.categorical_columns]
+
+    def state_dict(self) -> Dict:
+        return {
+            "feature_columns": self.feature_columns,
+            "target_column": self.target_column,
+            "task": self.task,
+            "column_stats": {
+                c: {"kind": s.kind, "mean": s.mean, "std": s.std, "vocab": s.vocab}
+                for c, s in self.column_stats.items()
+            },
+            "classes": self.classes,
+            "target_mean": self.target_mean,
+            "target_std": self.target_std,
+        }
+
+    @classmethod
+    def from_state_dict(cls, state: Dict) -> "TabularPreprocessor":
+        prep = cls(
+            feature_columns=list(state["feature_columns"]),
+            target_column=state["target_column"],
+            task=state["task"],
+            classes=list(state.get("classes", [])),
+            target_mean=state.get("target_mean", 0.0),
+            target_std=state.get("target_std", 1.0),
+        )
+        prep.column_stats = {
+            c: ColumnStats(kind=v["kind"], mean=v.get("mean", 0.0), std=v.get("std", 1.0), vocab=list(v.get("vocab", [])))
+            for c, v in state["column_stats"].items()
+        }
+        return prep

tensorless/devices/__init__.py

@@ -0,0 +1,3 @@
+from .device import auto_select_device, get_torch_device
+
+__all__ = ["auto_select_device", "get_torch_device"]

tensorless/devices/device.py

@@ -0,0 +1,107 @@
+"""Hardware auto-detection.
+
+Preference order is TPU -> GPU -> CPU, but selection is "intelligent"
+rather than blind: we verify each backend is actually usable (not just
+importable) before choosing it, and we fall back gracefully -- including
+at *runtime*, if a chosen device turns out to error out mid-training, the
+trainer (see `training/trainer.py`) will catch that and fall back too.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+
+
+def _tpu_available() -> bool:
+    try:
+        import torch_xla.core.xla_model as xm  # noqa: F401
+
+        return True
+    except Exception:
+        return False
+
+
+def _cuda_available() -> bool:
+    try:
+        return torch.cuda.is_available() and torch.cuda.device_count() > 0
+    except Exception:
+        return False
+
+
+def _mps_available() -> bool:
+    try:
+        return torch.backends.mps.is_available()
+    except Exception:
+        return False
+
+
+def _cuda_supports_bf16() -> bool:
+    try:
+        return torch.cuda.is_bf16_supported()
+    except Exception:
+        return False
+
+
+def auto_select_device(user_device: Optional[str], user_precision: Optional[str]) -> Tuple[str, str]:
+    """Resolve the device and precision to use.
+
+    `user_device` / `user_precision` are honored if given (with a
+    graceful downgrade if the requested device isn't actually available).
+    Otherwise we pick automatically: tpu > cuda > mps > cpu.
+    """
+    if user_device is not None:
+        device = user_device
+        if device == "tpu" and not _tpu_available():
+            device = "cuda" if _cuda_available() else "cpu"
+        elif device == "cuda" and not _cuda_available():
+            device = "cpu"
+        elif device == "mps" and not _mps_available():
+            device = "cpu"
+    else:
+        if _tpu_available():
+            device = "tpu"
+        elif _cuda_available():
+            device = "cuda"
+        elif _mps_available():
+            device = "mps"
+        else:
+            device = "cpu"
+
+    if user_precision is not None:
+        precision = user_precision
+    else:
+        if device == "cuda" and _cuda_supports_bf16():
+            precision = "bf16"
+        elif device == "cuda":
+            precision = "fp16"
+        elif device == "tpu":
+            precision = "bf16"
+        else:
+            # CPU and MPS: stick to fp32 for correctness/stability by default.
+            precision = "fp32"
+
+    return device, precision
+
+
+def get_torch_device(device: str) -> torch.device:
+    """Convert our string device name into a torch.device, with a
+    runtime fallback to CPU if the requested backend is unavailable.
+    """
+    try:
+        if device == "tpu":
+            import torch_xla.core.xla_model as xm
+
+            return xm.xla_device()
+        if device == "cuda":
+            if not _cuda_available():
+                return torch.device("cpu")
+            return torch.device("cuda")
+        if device == "mps":
+            if not _mps_available():
+                return torch.device("cpu")
+            return torch.device("mps")
+        return torch.device("cpu")
+    except Exception:
+        return torch.device("cpu")

tensorless/errors.py

@@ -0,0 +1,36 @@
+"""Tensorless error hierarchy.
+
+All errors raised by Tensorless inherit from :class:`TensorlessError`, so
+callers can do::
+
+    try:
+        tl.train("./data")
+    except tl.TensorlessError as e:
+        ...
+
+instead of catching bare Exception.
+"""
+
+
+class TensorlessError(Exception):
+    """Base class for all Tensorless errors."""
+
+
+class DataError(TensorlessError):
+    """Raised when a dataset cannot be read, is empty, or is malformed."""
+
+
+class ConfigError(TensorlessError):
+    """Raised when user-supplied configuration is invalid or contradictory."""
+
+
+class ModelError(TensorlessError):
+    """Raised for unsupported model/task combinations or model build failures."""
+
+
+class CheckpointError(TensorlessError):
+    """Raised when a checkpoint is missing, corrupt, or incompatible."""
+
+
+class SerializationError(TensorlessError):
+    """Raised when a `.tl` file cannot be written or read."""

tensorless/models/__init__.py

@@ -0,0 +1,5 @@
+from .registry import build_model
+from .transformer import TinyTransformer
+from .mlp import TabularMLP
+
+__all__ = ["build_model", "TinyTransformer", "TabularMLP"]

tensorless/models/mlp.py

@@ -0,0 +1,64 @@
+"""MLP model for tabular data (classification / regression).
+
+Numeric columns feed directly into the network; each categorical column
+gets its own small embedding table, and all features are concatenated
+before the hidden layers.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _embedding_dim(vocab_size: int) -> int:
+    # Common rule of thumb, capped to keep tiny-data models tiny.
+    return max(2, min(32, round(1.6 * (vocab_size ** 0.56))))
+
+
+class TabularMLP(nn.Module):
+    def __init__(
+        self,
+        n_numeric: int,
+        categorical_vocab_sizes: List[int],
+        d_model: int,
+        layers: int,
+        dropout: float,
+        task: str,
+        n_classes: int = 0,
+    ):
+        super().__init__()
+        self.task = task
+        self.n_numeric = n_numeric
+        self.categorical_vocab_sizes = categorical_vocab_sizes
+
+        self.embeddings = nn.ModuleList(
+            [nn.Embedding(v, _embedding_dim(v)) for v in categorical_vocab_sizes]
+        )
+        cat_dim = sum(_embedding_dim(v) for v in categorical_vocab_sizes)
+        in_dim = n_numeric + cat_dim
+
+        dims = [in_dim] + [d_model] * layers
+        modules = []
+        for i in range(len(dims) - 1):
+            modules += [nn.Linear(dims[i], dims[i + 1]), nn.ReLU(), nn.Dropout(dropout)]
+        self.backbone = nn.Sequential(*modules)
+
+        out_dim = n_classes if task == "classification" else 1
+        self.head = nn.Linear(dims[-1], out_dim)
+
+    def forward(self, numeric: torch.Tensor, categorical: torch.Tensor) -> torch.Tensor:
+        parts = []
+        if self.n_numeric > 0:
+            parts.append(numeric)
+        for i, emb in enumerate(self.embeddings):
+            parts.append(emb(categorical[:, i]))
+        x = torch.cat(parts, dim=1) if parts else numeric
+        x = self.backbone(x)
+        out = self.head(x)
+        if self.task == "regression":
+            return out.squeeze(-1)
+        return out

tensorless/models/registry.py

@@ -0,0 +1,53 @@
+"""Model registry.
+
+Central place that knows how to build a model given a task + resolved
+config + the dataset-derived metadata (vocab size, n_classes, etc). Kept
+separate from `auto/config.py` so new model types/backends can be added
+here without touching auto-configuration logic, per the "extensible
+architecture" requirement.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import torch.nn as nn
+
+from .transformer import TinyTransformer
+from .mlp import TabularMLP
+from ..errors import ModelError
+
+
+def build_model(task: str, model_type: str, cfg: Dict[str, Any], meta: Dict[str, Any]) -> nn.Module:
+    """Build a fresh, randomly-initialized model.
+
+    `cfg` is the resolved training config (dict). `meta` carries
+    task-specific sizing info produced during data prep, e.g.:
+      - text tasks: {"vocab_size": int, "pad_id": int, "n_classes": int}
+      - tabular tasks: {"n_numeric": int, "categorical_vocab_sizes": [...], "n_classes": int}
+    """
+    if model_type == "transformer":
+        return TinyTransformer(
+            vocab_size=meta["vocab_size"],
+            d_model=cfg["d_model"],
+            layers=cfg["layers"],
+            heads=cfg["heads"],
+            ff_mult=cfg["ff_mult"],
+            dropout=cfg["dropout"],
+            max_seq_len=cfg["max_seq_len"],
+            task=task,
+            n_classes=meta.get("n_classes", 0),
+            pad_id=meta.get("pad_id", 0),
+        )
+    elif model_type == "mlp":
+        return TabularMLP(
+            n_numeric=meta["n_numeric"],
+            categorical_vocab_sizes=meta["categorical_vocab_sizes"],
+            d_model=cfg["d_model"],
+            layers=cfg["layers"],
+            dropout=cfg["dropout"],
+            task=task,
+            n_classes=meta.get("n_classes", 0),
+        )
+    else:
+        raise ModelError(f"Unknown model_type '{model_type}'.")

tensorless/models/transformer.py

@@ -0,0 +1,175 @@
+"""A small, dependency-free (beyond PyTorch) GPT-style decoder transformer.
+
+Used for:
+  - "text-generation": next-token prediction over the char vocabulary
+  - "text-classification": same backbone, with a classification head on
+    the final token's hidden state instead of a language-modeling head
+
+Kept intentionally compact -- this is not meant to compete with
+production LLM training frameworks, it's meant to give Tensorless a real,
+working, from-scratch model that trains fast enough on CPU for the
+"zero setup" experience to actually be pleasant.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(self, d_model: int, heads: int, dropout: float):
+        super().__init__()
+        assert d_model % heads == 0, "d_model must be divisible by heads"
+        self.heads = heads
+        self.head_dim = d_model // heads
+        self.qkv = nn.Linear(d_model, 3 * d_model)
+        self.proj = nn.Linear(d_model, d_model)
+        self.dropout = dropout
+        self.resid_drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, T, C = x.shape
+        qkv = self.qkv(x)
+        q, k, v = qkv.split(C, dim=2)
+        q = q.view(B, T, self.heads, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.heads, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.heads, self.head_dim).transpose(1, 2)
+
+        out = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0.0, is_causal=attn_mask is None
+        )
+        out = out.transpose(1, 2).contiguous().view(B, T, C)
+        return self.resid_drop(self.proj(out))
+
+
+class MLP(nn.Module):
+    def __init__(self, d_model: int, ff_mult: int, dropout: float):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, d_model * ff_mult)
+        self.fc2 = nn.Linear(d_model * ff_mult, d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(self.fc2(F.gelu(self.fc1(x))))
+
+
+class Block(nn.Module):
+    def __init__(self, d_model: int, heads: int, ff_mult: int, dropout: float):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(d_model)
+        self.attn = CausalSelfAttention(d_model, heads, dropout)
+        self.ln2 = nn.LayerNorm(d_model)
+        self.mlp = MLP(d_model, ff_mult, dropout)
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        x = x + self.attn(self.ln1(x), attn_mask=attn_mask)
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+
+class TinyTransformer(nn.Module):
+    """Decoder-only transformer usable for LM or sequence classification."""
+
+    def __init__(
+        self,
+        vocab_size: int,
+        d_model: int,
+        layers: int,
+        heads: int,
+        ff_mult: int,
+        dropout: float,
+        max_seq_len: int,
+        task: str = "text-generation",
+        n_classes: int = 0,
+        pad_id: int = 0,
+    ):
+        super().__init__()
+        self.task = task
+        self.max_seq_len = max_seq_len
+        self.pad_id = pad_id
+
+        self.tok_emb = nn.Embedding(vocab_size, d_model)
+        self.pos_emb = nn.Embedding(max_seq_len, d_model)
+        self.drop = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList(
+            [Block(d_model, heads, ff_mult, dropout) for _ in range(layers)]
+        )
+        self.ln_f = nn.LayerNorm(d_model)
+
+        if task == "text-generation":
+            self.head = nn.Linear(d_model, vocab_size, bias=False)
+            self.head.weight = self.tok_emb.weight  # weight tying
+        elif task == "text-classification":
+            assert n_classes > 0, "n_classes must be set for text-classification"
+            self.head = nn.Linear(d_model, n_classes)
+        else:
+            raise ValueError(f"Unsupported task for TinyTransformer: {task}")
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, T = input_ids.shape
+        assert T <= self.max_seq_len, (
+            f"Sequence length {T} exceeds max_seq_len {self.max_seq_len}"
+        )
+        pos = torch.arange(T, device=input_ids.device).unsqueeze(0)
+        x = self.tok_emb(input_ids) + self.pos_emb(pos)
+        x = self.drop(x)
+
+        attn_mask = None
+        if attention_mask is not None:
+            # Combine causal mask with padding mask.
+            causal = torch.tril(torch.ones(T, T, device=input_ids.device, dtype=torch.bool))
+            pad = attention_mask.bool().unsqueeze(1).unsqueeze(1)  # B,1,1,T
+            attn_mask = (causal.unsqueeze(0).unsqueeze(0) & pad)
+
+        for block in self.blocks:
+            x = block(x, attn_mask=attn_mask)
+        x = self.ln_f(x)
+
+        if self.task == "text-generation":
+            return self.head(x)  # B, T, vocab_size
+        else:
+            if attention_mask is not None:
+                lengths = attention_mask.sum(dim=1).clamp(min=1) - 1
+            else:
+                lengths = torch.full((B,), T - 1, device=input_ids.device)
+            pooled = x[torch.arange(B, device=input_ids.device), lengths]
+            return self.head(pooled)  # B, n_classes
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int,
+        temperature: float = 0.8,
+        top_k: Optional[int] = 40,
+        eos_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        self.eval()
+        for _ in range(max_new_tokens):
+            cond = input_ids[:, -self.max_seq_len:]
+            logits = self(cond)
+            logits = logits[:, -1, :] / max(temperature, 1e-5)
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            input_ids = torch.cat([input_ids, next_id], dim=1)
+            if eos_id is not None and (next_id == eos_id).all():
+                break
+        return input_ids