neuralk-oss 0.1.0__tar.gz

neuralk_oss-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,75 @@
+# Builds the sdist and wheel and publishes them: to TestPyPI when run by hand
+# from the Actions tab, to PyPI when a `v*` tag is pushed (the tag must match
+# the version in pyproject.toml). Trusted Publishing: each index trusts this
+# workflow in its environment (`testpypi`, `pypi`), no token is stored. A
+# version is uploaded once per index; bump it in pyproject.toml before another
+# run.
+name: Release
+
+on:
+  workflow_dispatch:
+  push:
+    tags: ["v*"]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          persist-credentials: false
+      - uses: astral-sh/setup-uv@68559e99cc807c1eb7e9238fd43a15bba460eec2 # v10.3.3
+      - name: Check the tag matches the version
+        if: github.ref_type == 'tag'
+        env:
+          TAG: ${{ github.ref_name }}
+        run: |
+          version=$(uv version --short)
+          if [ "$TAG" != "v$version" ]; then
+            echo "::error::tag $TAG does not match version $version in pyproject.toml"
+            exit 1
+          fi
+      - run: uv build
+      - run: uvx twine check --strict dist/*
+      - uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
+        with:
+          name: dist
+          path: dist/
+          if-no-files-found: error
+
+  testpypi:
+    if: github.event_name == 'workflow_dispatch'
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/neuralk-oss
+    permissions:
+      id-token: write # Trusted Publishing
+    steps:
+      - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3
+        with:
+          repository-url: https://test.pypi.org/legacy/
+
+  pypi:
+    if: github.event_name == 'push' && github.ref_type == 'tag'
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/neuralk-oss
+    permissions:
+      id-token: write # Trusted Publishing
+    steps:
+      - uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3

neuralk_oss-0.1.0/.gitignore

@@ -0,0 +1,18 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+build/
+dist/
+.venv/
+.pytest_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+.DS_Store
+
+# Model weights
+*.pt
+*.pth
+*.ckpt
+*.safetensors
+!tests/golden/**/*.safetensors

neuralk_oss-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

neuralk_oss-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Neuralk-AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

neuralk_oss-0.1.0/PKG-INFO

@@ -0,0 +1,121 @@
+Metadata-Version: 2.5
+Name: neuralk-oss
+Version: 0.1.0
+Summary: Neuralk-AI open-source models: loading and inference
+Project-URL: Repository, https://github.com/Neuralk-AI/neuralk_oss
+Author: Neuralk-AI
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: huggingface-hub
+Requires-Dist: numpy
+Requires-Dist: pandas<3,>=2.2
+Requires-Dist: pydantic>=2
+Requires-Dist: safetensors
+Requires-Dist: structured-data-models==0.2.1
+Requires-Dist: torch>=2.7
+Description-Content-Type: text/markdown
+
+# Neuralk OSS
+
+Neuralk-AI open-source models.
+
+## Seldon
+
+Seldon is a tabular in-context-learning foundation model: `fit` keeps the
+training table as context, `predict_proba` classifies new rows against it in
+one forward pass, with no gradient step.
+
+```python
+from neuralk_oss import SeldonClassifier
+
+clf = SeldonClassifier.from_pretrained("seldon-small-classifier")
+clf.fit(X_train, y_train)  # pandas DataFrame (or array) and labels
+proba = clf.predict_proba(X_test)  # (n_rows, n_classes), columns in clf.classes_ order
+```
+
+### Checkpoints
+
+| name | task | parameters |
+| --- | --- | --- |
+| `seldon-small-classifier` | classification, up to 10 classes per model call | 27.4M |
+
+A checkpoint is a directory holding `config.json` and `model.safetensors`.
+`from_pretrained` takes:
+
+- a checkpoint name, such as `"seldon-small-classifier"`: downloaded from the
+  Hugging Face repository [`Neuralk-AI/seldon`](https://huggingface.co/Neuralk-AI/seldon),
+  where each checkpoint is a directory, on the first call, then read from the
+  Hugging Face cache (`~/.cache/huggingface`, or `HF_HOME`);
+- a local directory, read in place: `from_pretrained("path/to/seldon-small-classifier")`;
+- any Hugging Face repository, with the checkpoint's directory and a revision:
+  `from_pretrained("org/repo", subfolder="seldon-small-classifier", revision="v1")`.
+
+The repository is private for now: log in once with an account that can read
+it (`hf auth login`), or pass `token=...` to `from_pretrained`.
+
+Column types are inferred (categorical columns are embedded as such, NaNs are
+kept for the model's own imputation), each of the `n_estimators` ensemble
+members (8 by default) sees its own preprocessing, and their logits are
+averaged. Inference runs in float16 autocast on CUDA by default
+(`amp=False` for float32). Tables with more classes than the model's head
+(10) go through error-correcting output codes.
+
+### Preprocessing
+
+The preprocessing and ensembling come from structured-data-models (`sdm`,
+Apache-2.0), pinned in `pyproject.toml`. `SeldonRecipe` builds its
+recipe one overridable method per step:
+
+```python
+import sdm.processing as sp
+from neuralk_oss import SeldonClassifier, SeldonRecipe
+
+
+class MyRecipe(SeldonRecipe):
+    max_features = 200
+
+    def numerical_transforms(self):
+        return [sp.Identity(), sp.RankGaussian(max_knots=self.rank_gaussian_knots)]
+
+
+clf = SeldonClassifier.from_pretrained("seldon-small-classifier", recipe=MyRecipe())
+```
+
+### The model alone
+
+```python
+from neuralk_oss import Seldon
+
+model = Seldon.from_pretrained("seldon-small-classifier")
+context = model.encode_context(X_context, y_context, categorical_mask)  # (B, T, H), (B, T), (B, H)
+logits = model.predict_query(X_query, context, categorical_mask)  # (B, T_query, 10)
+```
+
+## Layout
+
+```
+src/neuralk_oss/
+├── model/       # Seldon's architecture, config and checkpoint format
+└── inference/   # sdm estimator, preprocessing recipe, scikit-learn style classifier
+tests/
+```
+
+## Development
+
+With [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync
+uv run pytest
+uv run ruff check . && uv run ruff format --check .
+```
+
+`tests/test_golden.py` compares the outputs of tiny random-weight checkpoints
+(`tests/golden/`) to saved references. After an intended change of outputs,
+regenerate them with `UPDATE_GOLDEN=1 uv run pytest tests/test_golden.py` and
+commit the new `tests/golden/expected/`. With CUDA, the same cases also run on
+the GPU in float16 autocast; the classifier's CUDA references
+(`tests/golden/expected/cuda/`) are only regenerated on a CUDA machine.

neuralk_oss-0.1.0/README.md

@@ -0,0 +1,101 @@
+# Neuralk OSS
+
+Neuralk-AI open-source models.
+
+## Seldon
+
+Seldon is a tabular in-context-learning foundation model: `fit` keeps the
+training table as context, `predict_proba` classifies new rows against it in
+one forward pass, with no gradient step.
+
+```python
+from neuralk_oss import SeldonClassifier
+
+clf = SeldonClassifier.from_pretrained("seldon-small-classifier")
+clf.fit(X_train, y_train)  # pandas DataFrame (or array) and labels
+proba = clf.predict_proba(X_test)  # (n_rows, n_classes), columns in clf.classes_ order
+```
+
+### Checkpoints
+
+| name | task | parameters |
+| --- | --- | --- |
+| `seldon-small-classifier` | classification, up to 10 classes per model call | 27.4M |
+
+A checkpoint is a directory holding `config.json` and `model.safetensors`.
+`from_pretrained` takes:
+
+- a checkpoint name, such as `"seldon-small-classifier"`: downloaded from the
+  Hugging Face repository [`Neuralk-AI/seldon`](https://huggingface.co/Neuralk-AI/seldon),
+  where each checkpoint is a directory, on the first call, then read from the
+  Hugging Face cache (`~/.cache/huggingface`, or `HF_HOME`);
+- a local directory, read in place: `from_pretrained("path/to/seldon-small-classifier")`;
+- any Hugging Face repository, with the checkpoint's directory and a revision:
+  `from_pretrained("org/repo", subfolder="seldon-small-classifier", revision="v1")`.
+
+The repository is private for now: log in once with an account that can read
+it (`hf auth login`), or pass `token=...` to `from_pretrained`.
+
+Column types are inferred (categorical columns are embedded as such, NaNs are
+kept for the model's own imputation), each of the `n_estimators` ensemble
+members (8 by default) sees its own preprocessing, and their logits are
+averaged. Inference runs in float16 autocast on CUDA by default
+(`amp=False` for float32). Tables with more classes than the model's head
+(10) go through error-correcting output codes.
+
+### Preprocessing
+
+The preprocessing and ensembling come from structured-data-models (`sdm`,
+Apache-2.0), pinned in `pyproject.toml`. `SeldonRecipe` builds its
+recipe one overridable method per step:
+
+```python
+import sdm.processing as sp
+from neuralk_oss import SeldonClassifier, SeldonRecipe
+
+
+class MyRecipe(SeldonRecipe):
+    max_features = 200
+
+    def numerical_transforms(self):
+        return [sp.Identity(), sp.RankGaussian(max_knots=self.rank_gaussian_knots)]
+
+
+clf = SeldonClassifier.from_pretrained("seldon-small-classifier", recipe=MyRecipe())
+```
+
+### The model alone
+
+```python
+from neuralk_oss import Seldon
+
+model = Seldon.from_pretrained("seldon-small-classifier")
+context = model.encode_context(X_context, y_context, categorical_mask)  # (B, T, H), (B, T), (B, H)
+logits = model.predict_query(X_query, context, categorical_mask)  # (B, T_query, 10)
+```
+
+## Layout
+
+```
+src/neuralk_oss/
+├── model/       # Seldon's architecture, config and checkpoint format
+└── inference/   # sdm estimator, preprocessing recipe, scikit-learn style classifier
+tests/
+```
+
+## Development
+
+With [uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync
+uv run pytest
+uv run ruff check . && uv run ruff format --check .
+```
+
+`tests/test_golden.py` compares the outputs of tiny random-weight checkpoints
+(`tests/golden/`) to saved references. After an intended change of outputs,
+regenerate them with `UPDATE_GOLDEN=1 uv run pytest tests/test_golden.py` and
+commit the new `tests/golden/expected/`. With CUDA, the same cases also run on
+the GPU in float16 autocast; the classifier's CUDA references
+(`tests/golden/expected/cuda/`) are only regenerated on a CUDA machine.

neuralk_oss-0.1.0/pyproject.toml

@@ -0,0 +1,51 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "neuralk-oss"
+version = "0.1.0"
+description = "Neuralk-AI open-source models: loading and inference"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Neuralk-AI" }]
+requires-python = ">=3.11"
+dependencies = [
+    "huggingface_hub",
+    "numpy",
+    "pandas>=2.2,<3",
+    "pydantic>=2",
+    "safetensors",
+    # Pinned: the preprocessing and ensembling Seldon was evaluated with.
+    "structured-data-models==0.2.1",
+    "torch>=2.7",
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.urls]
+Repository = "https://github.com/Neuralk-AI/neuralk_oss"
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+    "ruff>=0.6",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/neuralk_oss"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]
+# The formatter wraps code; docstring summaries stay on one line.
+ignore = ["E501"]

neuralk_oss-0.1.0/src/neuralk_oss/__init__.py

@@ -0,0 +1,7 @@
+"""Neuralk-AI open-source models."""
+
+from neuralk_oss.inference import SeldonClassifier, SeldonICL, SeldonRecipe
+from neuralk_oss.model import Seldon, SeldonConfig
+
+__all__ = ["Seldon", "SeldonClassifier", "SeldonConfig", "SeldonICL", "SeldonRecipe"]
+__version__ = "0.1.0"

neuralk_oss-0.1.0/src/neuralk_oss/inference/__init__.py

@@ -0,0 +1,7 @@
+"""Inference: Seldon behind sdm's preprocessing and ensembling, as a classifier."""
+
+from neuralk_oss.inference.classifier import SeldonClassifier
+from neuralk_oss.inference.estimator import SeldonICL
+from neuralk_oss.inference.recipe import SeldonRecipe
+
+__all__ = ["SeldonClassifier", "SeldonICL", "SeldonRecipe"]

neuralk_oss-0.1.0/src/neuralk_oss/inference/classifier.py

@@ -0,0 +1,165 @@
+"""Seldon as a scikit-learn style classifier over pandas tables."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator
+from contextlib import contextmanager, nullcontext
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import sdm
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from neuralk_oss.inference.estimator import SeldonICL
+from neuralk_oss.inference.recipe import SeldonRecipe
+from neuralk_oss.model import Seldon
+
+_TARGET = "target"
+# Every attention kernel PyTorch has but cuDNN's (see ``cudnn_attention``).
+_NON_CUDNN_ATTENTION = [
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+
+
+class SeldonClassifier:
+    """A Seldon classifier: ``fit`` keeps the training table as context, ``predict_proba`` reads the test rows against it.
+
+    sdm infers each column's type (categorical columns are embedded as such,
+    NaNs kept for the model's imputation and missingness term), preprocesses
+    each ensemble member with ``recipe``, and averages the members' logits
+    before the softmax.
+
+    Parameters
+    ----------
+    model : Seldon
+        A classifier, e.g. ``Seldon.from_pretrained(path)``.
+    n_estimators : int, default=8
+        Ensemble members.
+    device : str or torch.device, optional
+        Device of the model and the tables; CUDA when available by default.
+    max_context_size : int, default=200_000
+        Training rows one member reads in context; a larger table gives each
+        member its own random subset.
+    amp : bool, default=True
+        float16 autocast on CUDA; False runs in float32.
+    cudnn_attention : bool, default=False
+        Let PyTorch's attention use cuDNN's kernel. It is faster once a table
+        shape has run, but builds a plan per new shape, seconds each, which a
+        table fitted once pays in full; by default the FlashAttention and
+        memory-efficient kernels run instead.
+    random_state : int, default=0
+        Seed of the context subsets and of sdm's per-member preprocessing.
+    recipe : SeldonRecipe, optional
+        The preprocessing; ``SeldonRecipe()`` by default.
+    query_chunk_size : int, default=8192
+        Test rows run through the model together, per member.
+    """
+
+    def __init__(
+        self,
+        model: Seldon,
+        *,
+        n_estimators: int = 8,
+        device: str | torch.device | None = None,
+        max_context_size: int = 200_000,
+        amp: bool = True,
+        random_state: int = 0,
+        cudnn_attention: bool = False,
+        recipe: SeldonRecipe | None = None,
+        query_chunk_size: int = 8192,
+    ) -> None:
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = torch.device(device)
+        self.n_estimators = n_estimators
+        self.max_context_size = max_context_size
+        self.amp = amp
+        self.cudnn_attention = cudnn_attention
+        self.random_state = random_state
+        self.estimator = SeldonICL(model, recipe, query_chunk_size).to(self.device)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        name_or_path: str | os.PathLike,
+        *,
+        subfolder: str | None = None,
+        revision: str | None = None,
+        token: str | None = None,
+        **kwargs: Any,
+    ) -> SeldonClassifier:
+        """The classifier of a saved model: a local directory, a checkpoint name such as ``"seldon-small-classifier"``, or a Hub repository (see ``Seldon.from_pretrained``); ``kwargs`` are the constructor's."""
+        model = Seldon.from_pretrained(
+            name_or_path, subfolder=subfolder, revision=revision, token=token
+        )
+        return cls(model, **kwargs)
+
+    def fit(self, X: pd.DataFrame | np.ndarray, y: pd.Series | np.ndarray) -> SeldonClassifier:
+        """Keep ``X`` and ``y`` as the context of every member."""
+        X = _frame(X)
+        self.classes_, codes = np.unique(np.asarray(y), return_inverse=True)
+        self._stypes = sdm.infer_stypes(X, _low_cardinality="infer")
+        x_context = sdm.TableTensor.from_pandas(df=X, stypes=self._stypes)
+        y_context = sdm.TableTensor.from_pandas(
+            df=pd.DataFrame({_TARGET: codes}), stypes={_TARGET: "categorical"}
+        )
+        num_estimators: int | None = self.n_estimators
+        self._subsampled = len(X) > self.max_context_size
+        if self._subsampled:
+            generator = torch.Generator().manual_seed(self.random_state)
+            repeats = -(-self.n_estimators * self.max_context_size // len(X))
+            permutations = [torch.randperm(len(X), generator=generator) for _ in range(repeats)]
+            shape = (self.n_estimators, self.max_context_size)
+            index = torch.cat(permutations)[: shape[0] * shape[1]]
+            x_context = x_context[index].unflatten(0, shape)
+            y_context = y_context[index].unflatten(0, shape)
+            num_estimators = None
+        with self._inference():
+            self.estimator.fit(
+                x=x_context.to(self.device),
+                y=y_context.to(self.device),
+                num_estimators=num_estimators,
+                generator=torch.Generator(self.device).manual_seed(self.random_state),
+            )
+        return self
+
+    def predict_proba(self, X: pd.DataFrame | np.ndarray) -> np.ndarray:
+        """Class probabilities of the rows of ``X``, ``(n_rows, n_classes)`` in the order of ``classes_``."""
+        X = _frame(X)
+        query = sdm.TableTensor.from_pandas(df=X, stypes=self._stypes, device=self.device)
+        if self._subsampled:
+            query = query.expand(self.n_estimators, *query.size())
+        with self._inference():
+            out = self.estimator.predict(query)
+        codes = [int(code) for code in out.columns[sdm.Stype.numerical]]
+        proba = np.zeros((len(X), len(self.classes_)), dtype=np.float32)
+        proba[:, codes] = out.numerical.float().cpu().numpy()
+        return proba
+
+    def predict(self, X: pd.DataFrame | np.ndarray) -> np.ndarray:
+        """The most probable class of each row of ``X``."""
+        return self.classes_[self.predict_proba(X).argmax(axis=1)]
+
+    @contextmanager
+    def _inference(self) -> Iterator[None]:
+        amp = self.amp and self.device.type == "cuda"
+        attention = nullcontext() if self.cudnn_attention else sdpa_kernel(_NON_CUDNN_ATTENTION)
+        with (
+            torch.inference_mode(),
+            torch.autocast(self.device.type, dtype=torch.float16, enabled=amp),
+            attention,
+        ):
+            yield
+
+
+def _frame(X: pd.DataFrame | np.ndarray) -> pd.DataFrame:
+    """``X`` as a DataFrame with string column names, which sdm requires."""
+    if isinstance(X, pd.DataFrame):
+        return X.rename(columns=str)
+    X = np.asarray(X)
+    return pd.DataFrame(X, columns=[f"x{i}" for i in range(X.shape[1])])