Adding simplified pages for __init__

Author: KevinMusgrave

docs/gen_ref_pages.py

@@ -4,6 +4,7 @@
 
 import mkdocs_gen_files
 from griffe.collections import ModulesCollection
+from griffe.dataclasses import Alias, Module
 from griffe.loader import GriffeLoader
 
 TOP_LEVEL_NAME = "pytorch_adapt"
@@ -14,34 +15,63 @@ def remove_pytorch_adapt(x):
     return [z for z in list(x.parts) if z != TOP_LEVEL_NAME]
 
 
-def main():
-    collection = ModulesCollection()
-    loader = GriffeLoader(modules_collection=collection)
-    loader.load_module(Path("src", TOP_LEVEL_NAME))
-    nav = mkdocs_gen_files.Nav()
+def get_init_entries(module_instance, init_entries, prefix=""):
+    for k, v in module_instance.members.items():
+        if isinstance(v, Module) and v.is_init_module:
+            init_key = f"{prefix}.{k}" if prefix else k
+            init_entries[init_key] = sorted(
+                [
+                    (name, member.target_path)
+                    for name, member in v.members.items()
+                    if isinstance(member, Alias)
+                ]
+            )
+            get_init_entries(v, init_entries, prefix=init_key)
+    return init_entries
+
+
+def set_init_pages(module_instance, nav):
+    init_entries = get_init_entries(module_instance, {}, prefix="")
+    for k, v in init_entries.items():
+        k_split = k.split(".")
+        doc_path = Path(*k_split, "index").with_suffix(".md")
+        full_doc_path = Path(FOLDER, doc_path)
+        nav[k_split] = doc_path
+        to_write = "\n".join([f"- [{x[0]}][{x[1]}]" for x in v])
+        with mkdocs_gen_files.open(full_doc_path, "w") as fd:
+            fd.write(to_write)
 
+
+def set_non_init_pages(collection, nav):
     for path in sorted(Path("src").rglob("*.py")):
         module_path = path.relative_to("src").with_suffix("")
-        parts = list(module_path.parts)
+        parts = tuple(module_path.parts)
         if (parts[-1] in ["__init__", "__main__"]) or (
             not collection[module_path.parts].has_docstrings
         ):
             continue
 
         doc_path = path.relative_to("src").with_suffix(".md")
-        full_doc_path = Path(FOLDER, doc_path)
         doc_path = Path(*remove_pytorch_adapt(doc_path))
-        full_doc_path = Path(*remove_pytorch_adapt(full_doc_path))
+        full_doc_path = Path(FOLDER, doc_path)
 
         for_nav = remove_pytorch_adapt(module_path)
         nav[for_nav] = doc_path
 
         with mkdocs_gen_files.open(full_doc_path, "w") as fd:
             ident = ".".join(parts)
-            print("::: " + ident, file=fd)
+            fd.write(f"::: {ident}")
 
         mkdocs_gen_files.set_edit_path(full_doc_path, path)
 
+
+def main():
+    collection = ModulesCollection()
+    loader = GriffeLoader(modules_collection=collection)
+    module_instance = loader.load_module(Path("src", TOP_LEVEL_NAME))
+    nav = mkdocs_gen_files.Nav()
+    set_init_pages(module_instance, nav)
+    set_non_init_pages(collection, nav)
     with mkdocs_gen_files.open(f"{FOLDER}/SUMMARY.md", "w") as nav_file:
         nav_file.writelines(nav.build_literate_nav())
 

mkdocs.yml

@@ -32,7 +32,7 @@ plugins:
       python:
         setup_commands:
           - import sys
-          - sys.path.append("src")
+          - sys.path.insert(0, "src")
         rendering:
           show_root_toc_entry: false
 - section-index

src/pytorch_adapt/hooks/gvb.py

@@ -98,10 +98,13 @@ def call(self, inputs, losses):
         outputs = self.hook(inputs, losses)[0]
         strs = c_f.filter(self.hook.out_keys, f"_[a-z]bridge$", ["^src", "^target"])
         [src_bridge, target_bridge] = c_f.extract([outputs, inputs], strs)
-        return outputs, {
-            f"src_bridge_loss": self.loss_fn(src_bridge),
-            f"target_bridge_loss": self.loss_fn(target_bridge),
-        }
+        return (
+            outputs,
+            {
+                f"src_bridge_loss": self.loss_fn(src_bridge),
+                f"target_bridge_loss": self.loss_fn(target_bridge),
+            },
+        )
 
     def _loss_keys(self):
         return [f"src_bridge_loss", f"target_bridge_loss"]

src/pytorch_adapt/hooks/utils.py

@@ -40,9 +40,10 @@ def __init__(self, loss_names: List[str], out_names: List[str], **kwargs):
     def call(self, inputs, losses):
         """"""
         out_keys = set(self.out_names) - inputs.keys()
-        return {k: None for k in out_keys}, {
-            k: c_f.zero_loss() for k in self.loss_names
-        }
+        return (
+            {k: None for k in out_keys},
+            {k: c_f.zero_loss() for k in self.loss_names},
+        )
 
     def _loss_keys(self):
         """"""

src/pytorch_adapt/hooks/vada.py

@@ -28,10 +28,13 @@ def call(self, inputs, losses):
         )
         src_vat_loss = self.loss_fn(src_imgs, src_logits, combined_model)
         target_vat_loss = self.loss_fn(target_imgs, target_logits, combined_model)
-        return outputs, {
-            "src_vat_loss": src_vat_loss,
-            "target_vat_loss": target_vat_loss,
-        }
+        return (
+            outputs,
+            {
+                "src_vat_loss": src_vat_loss,
+                "target_vat_loss": target_vat_loss,
+            },
+        )
 
     def _loss_keys(self):
         return ["src_vat_loss", "target_vat_loss"]

src/pytorch_adapt/validators/class_cluster_validator.py

@@ -103,14 +103,6 @@ def get_clustering_performance(
     centroid_init=None,
     feat_normalizer=None,
 ):
-    """
-    :param feats: N x out numpy vector
-    :param labels:  N numpy vector
-    :param num_classes: int
-    :param src_feats
-    :param pca_size
-    :return: silhouette and calinski harabasz scores
-    """
     num_target_feats = feats.shape[0]
 
     if src_feats is not None: