xpm-torch 0.1.0__tar.gz

xpm_torch-0.1.0/.git

@@ -0,0 +1 @@
+gitdir: ../.git/modules/xpm-torch

xpm_torch-0.1.0/.github/workflows/pytest.yml

@@ -0,0 +1,29 @@
+name: Python tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          enable-cache: true
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv sync --group test --group dev
+      - name: Lint with ruff
+        run: uv run ruff check .
+      - name: Run tests
+        run: uv run pytest

xpm_torch-0.1.0/.github/workflows/python-publish.yml

@@ -0,0 +1,25 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [created]
+
+permissions:
+  contents: write
+  id-token: write
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - name: Install build tools
+        run: pip install hatch
+      - name: Build package
+        run: hatch build
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

xpm_torch-0.1.0/.gitignore

@@ -0,0 +1,105 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+.ipynb
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mypy
+.mypy_cache/
+.vscode/
+src/xpmir/_version.py
+.DS_Store

xpm_torch-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,36 @@
+default_install_hook_types:
+  - pre-commit
+  - commit-msg
+
+repos:
+- hooks:
+  - id: check-yaml
+  - id: end-of-file-fixer
+  - id: trailing-whitespace
+  repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v6.0.0
+# - hooks:
+#   - id: python-check-blanket-noqa
+#   - id: python-no-eval
+#   repo: https://github.com/pre-commit/pygrep-hooks
+#   rev: v1.10.0
+- hooks:
+  - id: no-fixme-todo
+    name: No FIXME or TODO markers
+    entry: '(FIXME:|TODO:)'
+    language: pygrep
+    types: [python]
+  repo: local
+- hooks:
+  - id: ruff
+    args: [--fix]
+  - id: ruff-format
+    exclude: ^src/experimaestro/webui/data
+  repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.14.10
+- hooks:
+  - id: conventional-pre-commit
+    stages:
+    - commit-msg
+  repo: https://github.com/compilerla/conventional-pre-commit
+  rev: v4.3.0

xpm_torch-0.1.0/.readthedocs.yml

@@ -0,0 +1,18 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.12"
+  commands:
+    - asdf plugin add uv
+    - asdf install uv latest
+    - asdf global uv latest
+    # Install only docs deps (experimaestro + sphinx), not the full project.
+    # conf.py adds src/ to sys.path and mocks heavy deps (torch, lightning, etc.)
+    - uv sync --only-group docs --no-install-project
+    - uv run python -m sphinx -b html docs/source $READTHEDOCS_OUTPUT/html

xpm_torch-0.1.0/PKG-INFO

@@ -0,0 +1,65 @@
+Metadata-Version: 2.4
+Name: xpm-torch
+Version: 0.1.0
+Summary: Experimaestro module for Pytorch-based experiments
+License: GPL-3
+Keywords: PyTorch,experimaestro,experiments
+Requires-Python: >=3.10
+Requires-Dist: experimaestro>=2.3
+Requires-Dist: huggingface-hub>0.17
+Requires-Dist: lightning>=2.4
+Requires-Dist: safetensors>=0.4
+Requires-Dist: setuptools<82.0
+Requires-Dist: tensorboard>=2.19.0
+Requires-Dist: torchdata>=0.10
+Description-Content-Type: text/markdown
+
+# xpm-torch
+
+PyTorch training framework built on [experimaestro](https://experimaestro-python.readthedocs.io/) and [Lightning Fabric](https://lightning.ai/docs/fabric/).
+
+## Features
+
+- **Module system**: `Module` base class combining experimaestro `Config` with `torch.nn.Module`, with safetensors serialization
+- **Training**: `Learner` task with checkpointing, validation listeners, and TensorBoard logging
+- **HuggingFace Hub**: Upload/download models via `ExperimaestroHFHub` (from experimaestro)
+- **Distributed training**: Lightning Fabric integration for multi-GPU/node training
+
+## Installation
+
+```bash
+pip install xpm-torch
+```
+
+## Quick Start
+
+```python
+from xpm_torch.module import Module
+from experimaestro import Param
+import torch
+import torch.nn as nn
+
+class MyModel(Module):
+    input_dim: Param[int]
+    hidden_dim: Param[int]
+
+    def __initialize__(self):
+        self.fc1 = nn.Linear(self.input_dim, self.hidden_dim)
+
+    def forward(self, x):
+        return self.fc1(x)
+
+# Create config, then instance
+cfg = MyModel.C(input_dim=50, hidden_dim=100)
+model = cfg.instance()
+model.initialize()
+
+# Save/load with safetensors
+from pathlib import Path
+model.save_model(Path("checkpoint/model"))
+model.load_model(Path("checkpoint/model"))
+```
+
+## License
+
+GPL-3.0

xpm_torch-0.1.0/README.md

@@ -0,0 +1,49 @@
+# xpm-torch
+
+PyTorch training framework built on [experimaestro](https://experimaestro-python.readthedocs.io/) and [Lightning Fabric](https://lightning.ai/docs/fabric/).
+
+## Features
+
+- **Module system**: `Module` base class combining experimaestro `Config` with `torch.nn.Module`, with safetensors serialization
+- **Training**: `Learner` task with checkpointing, validation listeners, and TensorBoard logging
+- **HuggingFace Hub**: Upload/download models via `ExperimaestroHFHub` (from experimaestro)
+- **Distributed training**: Lightning Fabric integration for multi-GPU/node training
+
+## Installation
+
+```bash
+pip install xpm-torch
+```
+
+## Quick Start
+
+```python
+from xpm_torch.module import Module
+from experimaestro import Param
+import torch
+import torch.nn as nn
+
+class MyModel(Module):
+    input_dim: Param[int]
+    hidden_dim: Param[int]
+
+    def __initialize__(self):
+        self.fc1 = nn.Linear(self.input_dim, self.hidden_dim)
+
+    def forward(self, x):
+        return self.fc1(x)
+
+# Create config, then instance
+cfg = MyModel.C(input_dim=50, hidden_dim=100)
+model = cfg.instance()
+model.initialize()
+
+# Save/load with safetensors
+from pathlib import Path
+model.save_model(Path("checkpoint/model"))
+model.load_model(Path("checkpoint/model"))
+```
+
+## License
+
+GPL-3.0

xpm_torch-0.1.0/docs/Makefile

@@ -0,0 +1,12 @@
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

xpm_torch-0.1.0/docs/source/conf.py

@@ -0,0 +1,50 @@
+import sys
+from pathlib import Path
+
+# Add source directory so Sphinx can import xpm_torch without installing
+# the full project (which would pull in torch, lightning, etc.)
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "src"))
+
+project = "xpm-torch"
+copyright = "2024, Benjamin Piwowarski"
+author = "Benjamin Piwowarski"
+
+extensions = [
+    # Experimaestro extension for documenting Config/Param classes
+    "experimaestro.sphinx",
+    "sphinx_rtd_theme",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    # Link to other documentations
+    "sphinx.ext.intersphinx",
+    # Auto-link identifiers in code blocks to API docs
+    "sphinx_codeautolink",
+]
+
+intersphinx_mapping = {
+    "experimaestro": (
+        "https://experimaestro-python.readthedocs.io/en/latest/",
+        None,
+    ),
+}
+
+# Mock heavy dependencies that are not needed to build the docs.
+# experimaestro is NOT mocked — it's installed in the docs group
+# and needed for experimaestro.sphinx and Config/Param introspection.
+autodoc_mock_imports = [
+    "huggingface_hub",
+    "lightning",
+    "numpy",
+    "safetensors",
+    "tensorboard",
+    "torch",
+    "torchdata",
+]
+
+templates_path = ["_templates"]
+exclude_patterns = []
+html_theme = "sphinx_rtd_theme"
+html_static_path = ["_static"]
+source_suffix = ".rst"

xpm_torch-0.1.0/docs/source/huggingface.rst

@@ -0,0 +1,253 @@
+HuggingFace Hub Integration
+===========================
+
+xpm-torch models can be pushed to and loaded from the
+`HuggingFace Hub <https://huggingface.co/>`_ via
+:class:`~xpm_torch.huggingface.TorchHFHub` (which extends experimaestro's
+``ExperimaestroHFHub``). The serialization preserves the full experimaestro
+configuration graph, so a downloaded model can be used directly in further
+experiments.
+
+Exporting models via actions
+----------------------------
+
+The recommended way to export trained models is through experimaestro's
+**action system**. When a :class:`~xpm_torch.learner.Learner` is submitted,
+it automatically registers :class:`~xpm_torch.actions.ExportAction` instances
+for the last checkpoint and for each listener's best checkpoint. After the
+experiment completes, these actions can be executed interactively via the
+experimaestro CLI (``experimaestro experiments actions``).
+
+:class:`~xpm_torch.actions.ExportAction` prompts the user to choose between
+uploading to HuggingFace Hub or saving to a local directory, then delegates
+to :class:`~xpm_torch.huggingface.TorchHFHub`.
+
+How actions are registered
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Actions are registered during task submission via the ``add_action`` callback
+provided by experimaestro:
+
+.. code-block:: python
+
+    class Learner(Task):
+        def __submit__(self, dep, add_action):
+            loader = dep(self.model.loader_config(...))
+            # Register export action for the last checkpoint
+            add_action(self.model.export_action(loader, default_name="last"))
+            ...
+
+:meth:`Module.export_action(loader, **kwargs) <xpm_torch.module.Module.export_action>`
+returns an :class:`~xpm_torch.actions.ExportAction` config by default.
+Subclasses override this method to return library-specific actions (e.g.
+``XPMIRExportAction`` adds xpmir README sections and metadata).
+
+Customizing the export action
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To customize what happens during export, subclass
+:class:`~xpm_torch.actions.ExportAction` and override
+:meth:`~xpm_torch.actions.ExportAction.get_hub` to return a
+library-specific hub wrapper:
+
+.. code-block:: python
+
+    from xpm_torch.actions import ExportAction
+
+    class MyExportAction(ExportAction):
+        def get_hub(self):
+            return MyCustomHFHub(self.loader)
+
+Then override :meth:`~xpm_torch.module.Module.export_action` on your model
+to return this action:
+
+.. code-block:: python
+
+    class MyModel(Module):
+        def export_action(self, loader, **kwargs):
+            return MyExportAction.C(loader=loader, **kwargs)
+
+Direct API usage
+~~~~~~~~~~~~~~~~
+
+You can also use :class:`~xpm_torch.huggingface.TorchHFHub` directly:
+
+.. code-block:: python
+
+    from xpm_torch.huggingface import TorchHFHub
+
+    # Push a ModuleLoader (from loader_config or validation output)
+    TorchHFHub(loader).push_to_hub("your-org/model-name")
+
+    # Or save locally first
+    TorchHFHub(loader).save_pretrained("/path/to/save")
+
+:class:`~xpm_torch.huggingface.TorchHFHub` calls
+:meth:`~xpm_torch.module.ModuleLoader.write_hub_extras` and
+:meth:`~xpm_torch.module.ModuleLoader.hub_readme_sections` on the loader,
+so format-specific files (e.g. sentence-transformers configs) and README
+sections are generated automatically.
+
+What gets uploaded
+~~~~~~~~~~~~~~~~~~
+
+The serialized directory contains:
+
+- ``experimaestro.json`` — the config definition (for reloading with xpmir)
+- Model weight directories — named after the loader's ``DataPath`` fields
+  (or customized via ``__xpm_serialize__``)
+- Format-specific configs written by ``write_hub_extras`` (e.g.
+  ``modules.json``, ``router_config.json`` for sentence-transformers)
+- ``README.md`` — assembled from base + loader sections
+
+Loading a model from the Hub
+-----------------------------
+
+Low-level (experimaestro):
+
+.. code-block:: python
+
+    from experimaestro.huggingface import ExperimaestroHFHub
+
+    # Returns a deserialized config (e.g. ModuleLoader)
+    data = ExperimaestroHFHub.from_pretrained("your-org/model-name")
+
+For xpmir models, use ``AutoModel`` which returns a
+:class:`~xpm_torch.module.ModuleLoader` directly:
+
+.. code-block:: python
+
+    from xpmir.models import AutoModel
+
+    # Returns a ModuleLoader — use as an init task in experiments
+    loader = AutoModel.load_from_hf_hub("your-org/model-name")
+    # loader.model is the model config
+
+    # Or load as a ready-to-use instance for direct inference
+    model = AutoModel.load_from_hf_hub("your-org/model-name", as_instance=True)
+
+Customizing the HF checkpoint format
+-------------------------------------
+
+There are three extension points for customizing what gets written during
+Hub export:
+
+- :meth:`Module.loader_config(path) <xpm_torch.module.Module.loader_config>` —
+  on the model, controls which ``ModuleLoader`` subclass is returned
+- :meth:`ModuleLoader.write_hub_extras(save_directory) <xpm_torch.module.ModuleLoader.write_hub_extras>` —
+  on the loader, writes additional files (e.g. ST configs)
+- :meth:`ModuleLoader.hub_readme_sections() <xpm_torch.module.ModuleLoader.hub_readme_sections>` —
+  on the loader, provides named README sections with positioning
+
+The hooks are on :class:`~xpm_torch.module.ModuleLoader` (not on
+:class:`~xpm_torch.module.Module`) because the loader is the object that
+gets serialized for Hub export and holds the ``DataPath`` references to
+model weights. ``Module`` configs are data-less.
+
+:meth:`~xpm_torch.module.Module.loader_config`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When a model is serialized for the Hub, experimaestro serializes the
+:class:`~xpm_torch.module.ModuleLoader` returned by
+:meth:`~xpm_torch.module.Module.loader_config`. Subclasses override this
+method to return a custom loader with different ``DataPath`` fields:
+
+.. code-block:: python
+
+    from xpm_torch.module import Module, SimpleModuleLoader
+
+    class MyModel(Module):
+        def loader_config(self, path):
+            # Default: single path DataPath
+            return SimpleModuleLoader.C(value=self, path=path)
+
+Loaders can also override ``__xpm_serialize__`` to control the directory
+names used during serialization (e.g. mapping field names to
+sentence-transformers conventions).
+
+:meth:`~xpm_torch.module.Module.save_model` / :meth:`~xpm_torch.module.Module.load_model`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+These methods on :class:`~xpm_torch.module.Module` control how model
+weights are written to and read from a directory. Override them to change
+the on-disk format:
+
+.. code-block:: python
+
+    class DualEncoderModel(Module):
+        encoder: Param[Module]
+        query_encoder: Param[Optional[Module]]
+
+        def save_model(self, path):
+            path.mkdir(parents=True, exist_ok=True)
+            self.encoder.save_model(path / "encoder")
+            if self.query_encoder is not None:
+                self.query_encoder.save_model(path / "query_encoder")
+
+        def load_model(self, path):
+            self.encoder.load_model(path / "encoder")
+            if (path / "query_encoder").exists():
+                self.query_encoder.load_model(path / "query_encoder")
+
+:meth:`~xpm_torch.module.ModuleLoader.write_hub_extras` and :meth:`~xpm_torch.module.ModuleLoader.hub_readme_sections`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To write additional files alongside the model weights during Hub export
+(e.g. sentence-transformers compatibility configs), create a custom
+:class:`~xpm_torch.module.ModuleLoader` subclass and override
+:meth:`~xpm_torch.module.ModuleLoader.write_hub_extras`.
+
+To add sections to the README, override
+:meth:`~xpm_torch.module.ModuleLoader.hub_readme_sections` and return a
+list of :class:`~xpm_torch.module.ReadmeSection`. Each section has a
+``key``, ``content``, and optional ``before``/``after`` constraints for
+positioning relative to the base sections (``frontmatter``,
+``description``, ``usage``, ``results``):
+
+.. code-block:: python
+
+    from xpm_torch.module import ModuleLoader, ReadmeSection
+
+    class MyCustomLoader(ModuleLoader):
+        def write_hub_extras(self, save_directory):
+            (save_directory / "my_config.json").write_text('{"format": "custom"}')
+
+        def hub_readme_sections(self):
+            return [
+                ReadmeSection(
+                    key="quick_loading",
+                    content="## Quick loading\n\n```python\nmodel = MyLib.load(...)\n```",
+                    before="usage",  # appears before the XPMIR usage section
+                ),
+            ]
+
+These hooks are only called during Hub export (by
+:class:`~xpm_torch.huggingface.TorchHFHub`), not during checkpoint saving.
+Then override :meth:`~xpm_torch.module.Module.loader_config` on your model
+to return this loader:
+
+.. code-block:: python
+
+    class MyModel(Module):
+        def loader_config(self, path):
+            return MyCustomLoader.C(value=self, path=path)
+
+Class hierarchy
+~~~~~~~~~~~~~~~
+
+- :class:`~experimaestro.huggingface.ExperimaestroHFHub` — base serialization
+  (experimaestro)
+- :class:`~xpm_torch.huggingface.TorchHFHub` — calls ``write_hub_extras``
+  and ``hub_readme_sections`` on the loader (xpm-torch)
+- ``XPMIRHFHub`` — adds xpmir README sections (frontmatter, usage, results)
+  and TensorBoard logs (xpmir)
+
+Utility functions
+-----------------
+
+Helper functions for working with HuggingFace Hub (cache checks,
+downloads, config retrieval):
+
+.. automodule:: xpm_torch.huggingface
+   :members:
+   :exclude-members: TorchHFHub

xpm_torch-0.1.0/docs/source/index.rst

@@ -0,0 +1,19 @@
+xpm-torch documentation
+========================
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+   overview
+   module
+   training
+   optimization
+   huggingface
+
+Indices and tables
+==================
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`