Move mkdocs-autoref-plugin to new external tool (#96)

* Move mkdocs-autoref-plugin to new external tool Signed-off-by: Fabrice Normandin <[email protected]> * Minor changes to docstrings for docs Signed-off-by: Fabrice Normandin <[email protected]> * Use autoref-plugin in try-except block for tests Signed-off-by: Fabrice Normandin <[email protected]> --------- Signed-off-by: Fabrice Normandin <[email protected]>
mila-iqia · Dec 4, 2024 · b0b624f · b0b624f
1 parent 7008e89
commit b0b624f
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 306 deletions.
diff --git a/docs/macros.py b/docs/macros.py
@@ -8,9 +8,27 @@
 from pathlib import Path
 from typing import Any
 
+import torch
+
 if typing.TYPE_CHECKING:
     from mkdocs_macros.plugin import MacrosPlugin
 
+import lightning
+
+try:
+    from mkdocs_autoref_plugin.autoref_plugin import default_reference_sources
+
+    default_reference_sources.extend(
+        [
+            lightning.Trainer,
+            lightning.LightningModule,
+            lightning.LightningDataModule,
+            torch.nn.Module,
+        ]
+    )
+except ImportError:
+    pass
+
 logger = logging.getLogger(__name__)
 
 

diff --git a/project/datamodules/__init__.py b/project/datamodules/__init__.py
@@ -1,6 +1,6 @@
 """Datamodules (datasets + preprocessing + dataloading)
 
-See the :ref:`lightning.LightningDataModule` class for more information.
+See the `lightning.LightningDataModule` class for more information.
 """
 
 from .image_classification import ImageClassificationDataModule

diff --git a/project/datamodules/image_classification/imagenet.py b/project/datamodules/image_classification/imagenet.py
@@ -46,10 +46,10 @@ class ImageNetDataModule(ImageClassificationDataModule):
     - Made this a subclass of VisionDataModule
 
     Notes:
+
     - train_dataloader uses the train split of imagenet2012 and puts away a portion of it for the validation split.
     - val_dataloader uses the part of the train split of imagenet2012  that was not used for training via
         `num_imgs_per_val_class`
-        - TODO: needs to pass split='val' to UnlabeledImagenet.
     - test_dataloader uses the validation split of imagenet2012 for testing.
         - TODO: need to pass num_imgs_per_class=-1 for test dataset and split="test".
     """

diff --git a/project/utils/autoref_plugin.py b/project/utils/autoref_plugin.py
diff --git a/project/utils/autoref_plugin_test.py b/project/utils/autoref_plugin_test.py
diff --git a/pyproject.toml b/pyproject.toml
@@ -45,9 +45,6 @@ dependencies = [
 readme = "README.md"
 requires-python = ">= 3.10"
 
-[project.entry-points."mkdocs.plugins"]
-custom_autoref_plugin = "project.utils.autoref_plugin:CustomAutoRefPlugin"
-
 [dependency-groups]
 dev = [
     "mktestdocs>=0.2.2",
@@ -75,6 +72,7 @@ docs = [
     "mkdocs-video>=1.5.0",
     "mkdocs-section-index>=0.3.9",
     "mkdocs-macros-plugin>=1.0.5",
+    "mkdocs-autoref-plugin",
 ]
 gpu = ["jax[cuda12]>=0.4.31"]
 
@@ -118,6 +116,7 @@ managed = true
 
 [tool.uv.sources]
 remote-slurm-executor = { git = "https://github.com/lebrice/remote-slurm-executor", branch = "master" }
+mkdocs-autoref-plugin = { git = "https://github.com/lebrice/mkdocs-autoref-plugin" }
 
 # [tool.uv.tasks]
 # serve-docs = "mkdocs serve"