Bug fixes + prepare

.github/workflows/build.yaml

@@ -1,26 +1,93 @@
-name: Build wheels
+name: Publish Python 🐍 distribution 📦 to PyPI and TestPyPI
 
-on:
-  release:
-    types: [ published ]
+on: push
 
 jobs:
-  pypi-publish:
-    name: Build and upload release to PyPI
+  build:
+    name: Build distribution 📦
     runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: "3.x"
+    - name: Install pypa/build
+      run: >-
+        python3 -m
+        pip install
+        build
+        --user
+    - name: Build a binary wheel and a source tarball
+      run: python3 -m build
+    - name: Store the distribution packages
+      uses: actions/upload-artifact@v3
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: startsWith(github.ref, 'refs/tags/')  # only publish to PyPI on tag pushes
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/femr  # Replace <package-name> with your PyPI project name
+    permissions:
+      id-token: write  # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+    - name: Download all the dists
+      uses: actions/download-artifact@v3
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish distribution 📦 to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: >-
+      Sign the Python 🐍 distribution 📦 with Sigstore
+      and upload them to GitHub Release
+    needs:
+    - publish-to-pypi
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write  # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write  # IMPORTANT: mandatory for sigstore
+
     steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.x'
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install build
-      - name: Build package
-        run: python -m build
-      - name: Publish package
-        uses: pypa/gh-action-pypi-publish@release/v1
-        with:
-          password: ${{ secrets.pypi_token }}
+    - name: Download all the dists
+      uses: actions/download-artifact@v3
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Sign the dists with Sigstore
+      uses: sigstore/[email protected]
+      with:
+        inputs: >-
+          ./dist/*.tar.gz
+          ./dist/*.whl
+    - name: Create GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release create
+        '${{ github.ref_name }}'
+        --repo '${{ github.repository }}'
+        --notes ""
+    - name: Upload artifact signatures to GitHub Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      # Upload to GitHub Release using the `gh` CLI.
+      # `dist/` contains the built packages, and the
+      # sigstore-produced signatures and certificates.
+      run: >-
+        gh release upload
+        '${{ github.ref_name }}' dist/**
+        --repo '${{ github.repository }}'

.gitignore

@@ -7,3 +7,4 @@ ignore/*
 *.ipynb_checkpoints*
 tutorials/trash
 tutorials/tmp_trainer
+_version.py

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools", "setuptools-scm"]
+requires = ["setuptools >= 69.0", "setuptools-scm>=8.0"]
 build-backend = "setuptools.build_meta"
 
 [project]
@@ -15,15 +15,18 @@ dependencies = [
     "icecream == 2.1.3",
     "nptyping == 2.4.1",
     "msgpack >= 1.0.5",
-    "meds == 0.1.1",
+    "meds == 0.1.3",
     "torch >= 2.1.2",
     "transformers >= 4.25",
     "datasets >= 2.15",
     "polars >= 0.20",
     "dill >= 0.3.7",
 ]
 requires-python=">3.9"
-version = "0.2.2"
+dynamic = ["version"]
+
+[tool.setuptools_scm]
+version_file = "src/femr/_version.py"
 
 [project.scripts]
 

src/femr/__init__.py

@@ -0,0 +1 @@
+from femr._version import __version__  # noqa

src/femr/models/config.py

@@ -20,6 +20,20 @@ def __init__(
         hidden_act: str = "gelu",
         **kwargs,
     ) -> None:
+        """Defined a configuration for a FEMR Transformer.
+        
+        Arguments:
+            vocab_size: The number of tokens in the vocabulary
+            is_hierarchical: Whether to use a hierarchical vocabulary. See FEMRTokenizer for more information.
+            hidden_size: The internal representation size
+            intermediate_size: The size of the FFN in the transformer layers
+            n_heads: The number of attention heads
+            n_layers: The number of transformer encoder layers
+            attention_width: FEMR by default uses a local attention transformer. This defines the width of those attention windows.
+            use_normed_ages: Whether or not to provide normalized ages as a feature to the model
+            use_bias: Whether or not to use bias terms in the transformer layers. Current research indicates that this should be set to False.
+            hidden_act: The type of activation function to use in the transformer.
+            """
         super().__init__(**kwargs)
 
         self.vocab_size = vocab_size
@@ -39,18 +53,31 @@ def __init__(
 
 class FEMRTaskConfig(transformers.PretrainedConfig):
     def __init__(self, task_type: str = "", task_kwargs: Mapping[str, Any] = {}, **kwargs):
+        """A generic FEMR task definition. This holds state used for initalizing a tasks.py class.
+
+        Task.get_task_config returns the task type and kwargs used to initialize this.
+        
+        Arguments:
+            task_type: The name of the task. 
+            task_kwargs: Arbitrary arguments used to store state for that task.
+        """
         super().__init__(**kwargs)
         self.task_type = task_type
         self.task_kwargs = task_kwargs
 
 
 class FEMRModelConfig(transformers.PretrainedConfig):
+    """A model config is defined as the combination of a transformer config and a task config."""
     def __init__(
         self,
         transformer_config: Optional[Dict[str, Any]] = None,
         task_config: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
+        """A combination of a transformer config and a task config.
+        
+        It is possible to initialize this with only a transformer config, in which case the model will be configured for inference only.
+        """
         super().__init__(**kwargs)
         if transformer_config is None:
             transformer_config = {}
@@ -67,6 +94,9 @@ def __init__(
     def from_transformer_task_configs(
         cls, transformer_config: FEMRTransformerConfig, task_config: FEMRTaskConfig
     ) -> FEMRModelConfig:
+        """
+        Combine a transformer configuration and task configuration into a model configuration.
+        """
         if task_config is not None:
             task_config_dict = task_config.to_dict()
         else: