tensorless 0.1.0__py3-none-any.whl

tensorless/runtime.py

@@ -0,0 +1,158 @@
+"""Runtime inference wrapper.
+
+`tl.load("model.tl")` returns a `LoadedModel`, which knows how to rebuild
+the exact architecture used at training time, load the weights, and
+expose a task-appropriate prediction API:
+
+  - text-generation      -> `.generate(prompt)` and `.chat()`
+  - text-classification  -> `.predict(text)`
+  - classification/regression (tabular) -> `.predict(record_or_records)`
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Union
+
+import torch
+
+from .models.registry import build_model
+from .tokenization.char_tokenizer import CharTokenizer
+from .data.tabular import TabularPreprocessor
+from .devices.device import get_torch_device
+from .errors import ModelError
+from .serialization.tl_format import load_tl
+
+
+class LoadedModel:
+    def __init__(self, payload: Dict[str, Any], device: str = None):
+        self.payload = payload
+        self.task: str = payload["task"]
+        self.model_type: str = payload["model_type"]
+        self.config: Dict[str, Any] = payload["config"]
+        self.meta: Dict[str, Any] = payload["meta"]
+        self.metrics: Dict[str, Any] = payload.get("metrics", {})
+        self.dataset_fingerprint: str = payload.get("dataset_fingerprint")
+
+        self.tokenizer = None
+        if payload.get("tokenizer_state") is not None:
+            self.tokenizer = CharTokenizer.from_state_dict(payload["tokenizer_state"])
+
+        self.preprocessor = None
+        if payload.get("preprocessor_state") is not None:
+            self.preprocessor = TabularPreprocessor.from_state_dict(payload["preprocessor_state"])
+
+        device_name = device or self.config.get("device", "cpu")
+        self.device = get_torch_device(device_name)
+
+        self.model = build_model(self.task, self.model_type, self.config, self.meta)
+        self.model.load_state_dict(payload["model_state_dict"])
+        self.model.to(self.device)
+        self.model.eval()
+
+    # ------------------------------------------------------------------
+    # Text generation
+    # ------------------------------------------------------------------
+    def generate(self, prompt: str = "", max_new_tokens: int = 200, temperature: float = 0.8, top_k: int = 40) -> str:
+        if self.task != "text-generation":
+            raise ModelError(f"generate() is only available for text-generation models, not '{self.task}'.")
+        ids = self.tokenizer.encode(prompt, add_special_tokens=True)[:-1]  # drop trailing eos
+        if not ids:
+            ids = [self.tokenizer.bos_id]
+        input_ids = torch.tensor([ids], dtype=torch.long, device=self.device)
+        out = self.model.generate(
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_k=top_k,
+            eos_id=self.tokenizer.eos_id,
+        )
+        return self.tokenizer.decode(out[0].tolist())
+
+    def chat(self) -> None:
+        """Interactive terminal chat loop for text-generation models."""
+        if self.task != "text-generation":
+            raise ModelError(f"chat() is only available for text-generation models, not '{self.task}'.")
+        print("Tensorless interactive chat. Type 'exit' or Ctrl+C to quit.")
+        while True:
+            try:
+                prompt = input("> ")
+            except (EOFError, KeyboardInterrupt):
+                print()
+                break
+            if prompt.strip().lower() in ("exit", "quit"):
+                break
+            reply = self.generate(prompt, max_new_tokens=200)
+            print(reply)
+
+    # ------------------------------------------------------------------
+    # Text classification
+    # ------------------------------------------------------------------
+    def _predict_text_classification(self, texts: List[str]) -> List[str]:
+        block_size = self.config["max_seq_len"]
+        batch_ids, batch_mask = [], []
+        for t in texts:
+            ids = self.tokenizer.encode(t, add_special_tokens=True)[:block_size]
+            mask = [1] * len(ids)
+            if len(ids) < block_size:
+                pad_len = block_size - len(ids)
+                ids = ids + [self.tokenizer.pad_id] * pad_len
+                mask = mask + [0] * pad_len
+            batch_ids.append(ids)
+            batch_mask.append(mask)
+        input_ids = torch.tensor(batch_ids, dtype=torch.long, device=self.device)
+        attn_mask = torch.tensor(batch_mask, dtype=torch.long, device=self.device)
+        with torch.no_grad():
+            logits = self.model(input_ids, attention_mask=attn_mask)
+        pred_idx = logits.argmax(dim=-1).tolist()
+        classes = self.meta["classes"]
+        return [classes[i] for i in pred_idx]
+
+    # ------------------------------------------------------------------
+    # Tabular classification / regression
+    # ------------------------------------------------------------------
+    def _predict_tabular(self, records: List[Dict[str, Any]]) -> List[Any]:
+        transformed = self.preprocessor.transform(records, with_target=False)
+        numeric = transformed["numeric"].to(self.device)
+        categorical = transformed["categorical"].to(self.device)
+        with torch.no_grad():
+            out = self.model(numeric, categorical)
+        if self.task == "classification":
+            pred_idx = out.argmax(dim=-1)
+            return self.preprocessor.inverse_target(pred_idx)
+        else:
+            return self.preprocessor.inverse_target(out)
+
+    # ------------------------------------------------------------------
+    # Unified predict()
+    # ------------------------------------------------------------------
+    def predict(self, x: Union[str, Dict[str, Any], List[Any]]) -> Any:
+        single = not isinstance(x, list)
+        items = [x] if single else x
+
+        if self.task == "text-classification":
+            preds = self._predict_text_classification(items)
+        elif self.task in ("classification", "regression"):
+            preds = self._predict_tabular(items)
+        elif self.task == "text-generation":
+            preds = [self.generate(prompt=str(i)) for i in items]
+        else:
+            raise ModelError(f"predict() not supported for task '{self.task}'.")
+
+        return preds[0] if single else preds
+
+    def info(self) -> Dict[str, Any]:
+        return {
+            "task": self.task,
+            "model_type": self.model_type,
+            "tensorless_version": self.payload.get("tensorless_version"),
+            "tl_format_version": self.payload.get("tl_format_version"),
+            "config": self.config,
+            "metrics": self.metrics,
+            "training_complete": self.payload.get("training_complete"),
+            "n_parameters": sum(p.numel() for p in self.model.parameters()),
+        }
+
+
+def load_model(path: str, device: str = None) -> LoadedModel:
+    payload = load_tl(path)
+    return LoadedModel(payload, device=device)

tensorless/serialization/__init__.py

@@ -0,0 +1,3 @@
+from .tl_format import save_tl, load_tl
+
+__all__ = ["save_tl", "load_tl"]

tensorless/serialization/tl_format.py

@@ -0,0 +1,89 @@
+"""The `.tl` file format.
+
+A `.tl` file is a single portable file (a torch pickle archive under the
+hood) containing everything needed for inference on a *different*
+machine, with no access to the original dataset or training code:
+
+    {
+        "tl_format_version": int,
+        "tensorless_version": str,
+        "task": str,                  # e.g. "text-generation"
+        "model_type": str,            # e.g. "transformer"
+        "config": {...},              # resolved training config
+        "meta": {...},                # vocab_size / n_classes / column info
+        "model_state_dict": {...},
+        "tokenizer_state": {...} | None,
+        "preprocessor_state": {...} | None,
+        "dataset_fingerprint": str,
+        "training_complete": bool,
+        "metrics": {...},
+    }
+
+We deliberately use a single file (rather than a directory/zip of many
+files) so users can `scp`/email/upload one `model.tl` and have it just
+work elsewhere, per the framework's "portable single file" requirement.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any, Dict
+
+import torch
+
+from .._version import __version__, TL_FORMAT_VERSION
+from ..errors import SerializationError
+
+REQUIRED_KEYS = (
+    "tl_format_version",
+    "task",
+    "model_type",
+    "config",
+    "meta",
+    "model_state_dict",
+)
+
+
+def save_tl(path: str, payload: Dict[str, Any]) -> None:
+    payload = dict(payload)
+    payload.setdefault("tl_format_version", TL_FORMAT_VERSION)
+    payload.setdefault("tensorless_version", __version__)
+
+    parent = os.path.dirname(os.path.abspath(path))
+    if parent:
+        os.makedirs(parent, exist_ok=True)
+
+    tmp_path = path + ".tmp"
+    try:
+        torch.save(payload, tmp_path)
+        os.replace(tmp_path, path)
+    except Exception as e:
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+        raise SerializationError(f"Failed to write .tl file to '{path}': {e}") from e
+
+
+def load_tl(path: str, map_location: str = "cpu") -> Dict[str, Any]:
+    if not os.path.isfile(path):
+        raise SerializationError(f"'{path}' does not exist.")
+    try:
+        payload = torch.load(path, map_location=map_location, weights_only=False)
+    except Exception as e:
+        raise SerializationError(f"Failed to read .tl file '{path}': {e}") from e
+
+    missing = [k for k in REQUIRED_KEYS if k not in payload]
+    if missing:
+        raise SerializationError(
+            f"'{path}' is missing required field(s) {missing}. It may be "
+            f"corrupt or not a valid Tensorless .tl file."
+        )
+
+    file_version = payload.get("tl_format_version")
+    if file_version > TL_FORMAT_VERSION:
+        raise SerializationError(
+            f"'{path}' was created with a newer .tl format (v{file_version}) "
+            f"than this installed version of Tensorless supports "
+            f"(v{TL_FORMAT_VERSION}). Please upgrade Tensorless."
+        )
+
+    return payload

tensorless/tokenization/__init__.py

@@ -0,0 +1,3 @@
+from .char_tokenizer import CharTokenizer
+
+__all__ = ["CharTokenizer"]

tensorless/tokenization/char_tokenizer.py

@@ -0,0 +1,84 @@
+"""Character-level tokenizer.
+
+Tensorless defaults to a character-level tokenizer because it:
+  - requires no external vocabulary files or training corpus assumptions
+  - works on any UTF-8 text out of the box (any language, code, etc.)
+  - has a small, easily-portable vocabulary that fits directly inside a
+    `.tl` file
+
+This is deliberately simple rather than a full BPE tokenizer -- Tensorless
+optimizes for "it just works with zero setup" over maximal efficiency.
+Advanced users can plug in their own tokenizer via `model_type=` /
+extension points documented in the developer docs.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Sequence
+
+PAD = "<pad>"
+UNK = "<unk>"
+BOS = "<bos>"
+EOS = "<eos>"
+SPECIAL_TOKENS = [PAD, UNK, BOS, EOS]
+
+
+class CharTokenizer:
+    def __init__(self, vocab: List[str] = None):
+        self.vocab: List[str] = vocab if vocab is not None else list(SPECIAL_TOKENS)
+        self._stoi: Dict[str, int] = {c: i for i, c in enumerate(self.vocab)}
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    @property
+    def pad_id(self) -> int:
+        return self._stoi[PAD]
+
+    @property
+    def unk_id(self) -> int:
+        return self._stoi[UNK]
+
+    @property
+    def bos_id(self) -> int:
+        return self._stoi[BOS]
+
+    @property
+    def eos_id(self) -> int:
+        return self._stoi[EOS]
+
+    @classmethod
+    def build(cls, texts: Sequence[str]) -> "CharTokenizer":
+        chars = set()
+        for t in texts:
+            chars.update(t)
+        vocab = list(SPECIAL_TOKENS) + sorted(chars)
+        return cls(vocab)
+
+    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        ids = [self._stoi.get(c, self.unk_id) for c in text]
+        if add_special_tokens:
+            ids = [self.bos_id] + ids + [self.eos_id]
+        return ids
+
+    def decode(self, ids: Sequence[int], skip_special_tokens: bool = True) -> str:
+        chars = []
+        special_ids = {self.pad_id, self.bos_id, self.eos_id} if skip_special_tokens else set()
+        for i in ids:
+            if i in special_ids:
+                continue
+            if 0 <= i < len(self.vocab):
+                tok = self.vocab[i]
+                if tok == UNK:
+                    chars.append("\ufffd")
+                elif tok not in SPECIAL_TOKENS:
+                    chars.append(tok)
+        return "".join(chars)
+
+    def state_dict(self) -> Dict:
+        return {"vocab": self.vocab}
+
+    @classmethod
+    def from_state_dict(cls, state: Dict) -> "CharTokenizer":
+        return cls(vocab=list(state["vocab"]))

tensorless/training/__init__.py

@@ -0,0 +1,4 @@
+from .trainer import run_training
+from .early_stopping import EarlyStopping
+
+__all__ = ["run_training", "EarlyStopping"]

tensorless/training/data_prep.py

@@ -0,0 +1,215 @@
+"""Turns a loaded `Dataset` + resolved config into PyTorch-ready tensors,
+train/val splits, and the tokenizer/preprocessor that produced them.
+
+Kept separate from `trainer.py` so the "how do I turn data into tensors"
+logic can be tested and extended independently of the training loop.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+from torch.utils.data import DataLoader, Dataset as TorchDataset
+
+from ..data.loader import Dataset
+from ..data.tabular import TabularPreprocessor
+from ..tokenization.char_tokenizer import CharTokenizer
+from ..auto.detector import target_column
+from ..errors import DataError
+
+
+@dataclass
+class PreparedData:
+    train_loader: DataLoader
+    val_loader: Optional[DataLoader]
+    meta: Dict[str, Any]
+    tokenizer: Optional[CharTokenizer] = None
+    preprocessor: Optional[TabularPreprocessor] = None
+
+
+class _LMChunkDataset(TorchDataset):
+    def __init__(self, token_ids: List[int], block_size: int):
+        self.ids = token_ids
+        self.block_size = block_size
+
+    def __len__(self) -> int:
+        return max(0, len(self.ids) - self.block_size)
+
+    def __getitem__(self, idx: int):
+        chunk = self.ids[idx: idx + self.block_size + 1]
+        x = torch.tensor(chunk[:-1], dtype=torch.long)
+        y = torch.tensor(chunk[1:], dtype=torch.long)
+        return x, y
+
+
+class _ClsTextDataset(TorchDataset):
+    def __init__(self, input_ids: List[List[int]], attn_masks: List[List[int]], labels: List[int]):
+        self.input_ids = input_ids
+        self.attn_masks = attn_masks
+        self.labels = labels
+
+    def __len__(self) -> int:
+        return len(self.labels)
+
+    def __getitem__(self, idx: int):
+        return (
+            torch.tensor(self.input_ids[idx], dtype=torch.long),
+            torch.tensor(self.attn_masks[idx], dtype=torch.long),
+            torch.tensor(self.labels[idx], dtype=torch.long),
+        )
+
+
+class _TabularDataset(TorchDataset):
+    def __init__(self, numeric: torch.Tensor, categorical: torch.Tensor, target: torch.Tensor):
+        self.numeric = numeric
+        self.categorical = categorical
+        self.target = target
+
+    def __len__(self) -> int:
+        return self.numeric.shape[0]
+
+    def __getitem__(self, idx: int):
+        return self.numeric[idx], self.categorical[idx], self.target[idx]
+
+
+def _split_indices(n: int, val_split: float, seed: int) -> Tuple[List[int], List[int]]:
+    idx = list(range(n))
+    random.Random(seed).shuffle(idx)
+    n_val = int(n * val_split)
+    val_idx = idx[:n_val]
+    train_idx = idx[n_val:]
+    if not train_idx:
+        train_idx = idx
+        val_idx = []
+    return train_idx, val_idx
+
+
+def prepare_text_generation(
+    ds: Dataset, cfg: Dict[str, Any], tokenizer: Optional[CharTokenizer] = None
+) -> PreparedData:
+    tokenizer = tokenizer or CharTokenizer.build(ds.texts)
+    all_ids: List[int] = []
+    for t in ds.texts:
+        all_ids.extend(tokenizer.encode(t, add_special_tokens=True))
+
+    block_size = cfg["max_seq_len"]
+    if len(all_ids) <= block_size:
+        # Pad tiny corpora so we have at least one training example.
+        all_ids = all_ids + [tokenizer.pad_id] * (block_size + 1 - len(all_ids))
+
+    n_val_tokens = int(len(all_ids) * cfg["val_split"])
+    if n_val_tokens > block_size + 1:
+        split_point = len(all_ids) - n_val_tokens
+        train_ids, val_ids = all_ids[:split_point], all_ids[split_point:]
+    else:
+        train_ids, val_ids = all_ids, None
+
+    train_ds = _LMChunkDataset(train_ids, block_size)
+    train_loader = DataLoader(
+        train_ds, batch_size=cfg["batch_size"], shuffle=True, drop_last=False
+    )
+
+    val_loader = None
+    if val_ids is not None:
+        val_ds = _LMChunkDataset(val_ids, block_size)
+        if len(val_ds) > 0:
+            val_loader = DataLoader(val_ds, batch_size=cfg["batch_size"], shuffle=False)
+
+    meta = {
+        "vocab_size": tokenizer.vocab_size,
+        "pad_id": tokenizer.pad_id,
+        "n_classes": 0,
+    }
+    return PreparedData(train_loader=train_loader, val_loader=val_loader, meta=meta, tokenizer=tokenizer)
+
+
+def prepare_text_classification(
+    ds: Dataset,
+    cfg: Dict[str, Any],
+    tokenizer: Optional[CharTokenizer] = None,
+    classes: Optional[List[str]] = None,
+) -> PreparedData:
+    tokenizer = tokenizer or CharTokenizer.build(ds.texts)
+    classes = classes or sorted(set(ds.labels))
+    label2id = {c: i for i, c in enumerate(classes)}
+
+    block_size = cfg["max_seq_len"]
+    input_ids, attn_masks, labels = [], [], []
+    for text, label in zip(ds.texts, ds.labels):
+        ids = tokenizer.encode(text, add_special_tokens=True)[:block_size]
+        mask = [1] * len(ids)
+        if len(ids) < block_size:
+            pad_len = block_size - len(ids)
+            ids = ids + [tokenizer.pad_id] * pad_len
+            mask = mask + [0] * pad_len
+        input_ids.append(ids)
+        attn_masks.append(mask)
+        labels.append(label2id.get(label, 0))
+
+    n = len(labels)
+    train_idx, val_idx = _split_indices(n, cfg["val_split"], cfg["seed"])
+
+    def subset(indices):
+        return _ClsTextDataset(
+            [input_ids[i] for i in indices],
+            [attn_masks[i] for i in indices],
+            [labels[i] for i in indices],
+        )
+
+    train_loader = DataLoader(subset(train_idx), batch_size=cfg["batch_size"], shuffle=True)
+    val_loader = (
+        DataLoader(subset(val_idx), batch_size=cfg["batch_size"], shuffle=False) if val_idx else None
+    )
+
+    meta = {
+        "vocab_size": tokenizer.vocab_size,
+        "pad_id": tokenizer.pad_id,
+        "n_classes": len(classes),
+        "classes": classes,
+    }
+    return PreparedData(train_loader=train_loader, val_loader=val_loader, meta=meta, tokenizer=tokenizer)
+
+
+def prepare_tabular(
+    ds: Dataset,
+    cfg: Dict[str, Any],
+    task: str,
+    preprocessor: Optional[TabularPreprocessor] = None,
+) -> PreparedData:
+    target_col = target_column(ds)
+    if target_col is None:
+        raise DataError("Could not determine a target column for tabular training.")
+
+    prep = preprocessor or TabularPreprocessor().fit(ds.records, ds.columns, target_col, task)
+    transformed = prep.transform(ds.records, with_target=True)
+
+    n = transformed["numeric"].shape[0]
+    train_idx, val_idx = _split_indices(n, cfg["val_split"], cfg["seed"])
+    train_idx_t = torch.tensor(train_idx, dtype=torch.long)
+
+    train_ds = _TabularDataset(
+        transformed["numeric"][train_idx_t],
+        transformed["categorical"][train_idx_t],
+        transformed["target"][train_idx_t],
+    )
+    train_loader = DataLoader(train_ds, batch_size=cfg["batch_size"], shuffle=True)
+
+    val_loader = None
+    if val_idx:
+        val_idx_t = torch.tensor(val_idx, dtype=torch.long)
+        val_ds = _TabularDataset(
+            transformed["numeric"][val_idx_t],
+            transformed["categorical"][val_idx_t],
+            transformed["target"][val_idx_t],
+        )
+        val_loader = DataLoader(val_ds, batch_size=cfg["batch_size"], shuffle=False)
+
+    meta = {
+        "n_numeric": len(prep.numeric_columns),
+        "categorical_vocab_sizes": prep.categorical_vocab_sizes(),
+        "n_classes": len(prep.classes) if task == "classification" else 0,
+    }
+    return PreparedData(train_loader=train_loader, val_loader=val_loader, meta=meta, preprocessor=prep)

tensorless/training/early_stopping.py

@@ -0,0 +1,31 @@
+"""Early stopping.
+
+Tracks a validation metric (lower is better, e.g. loss) and reports when
+training should stop because it hasn't improved by at least `min_delta`
+for `patience` consecutive checks.
+"""
+
+from __future__ import annotations
+
+
+class EarlyStopping:
+    def __init__(self, patience: int = 5, min_delta: float = 1e-4):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.best: float = float("inf")
+        self.best_state: dict = None
+        self.num_bad_checks = 0
+        self.should_stop = False
+
+    def step(self, value: float, state: dict = None) -> bool:
+        """Returns True if `value` is a new best."""
+        if value < self.best - self.min_delta:
+            self.best = value
+            self.best_state = state
+            self.num_bad_checks = 0
+            return True
+        else:
+            self.num_bad_checks += 1
+            if self.num_bad_checks >= self.patience:
+                self.should_stop = True
+            return False