clean up

mosaicml · Apr 5, 2024 · 01e1d7e · 01e1d7e
1 parent f327dc3
commit 01e1d7e
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/llmfoundry/layers_registry.py b/llmfoundry/layers_registry.py
@@ -15,14 +15,18 @@
                         entry_points=True,
                         description=_norm_description)
 
-_ffns_description = """The ffns registry is used to register functions that build ffn layers."""
+_ffns_description = (
+    'The ffns registry is used to register functions that build ffn layers.' +
+    'See ffn.py for examples.')
 ffns = create_registry('llmfoundry',
                        'ffns',
                        generic_type=Callable,
                        entry_points=True,
                        description=_ffns_description)
 
-_ffns_with_norm_description = """The ffns_with_norm registry is used to register functions that build ffn layers that apply a normalization layer."""
+_ffns_with_norm_description = (
+    'The ffns_with_norm registry is used to register functions that build ffn layers that apply a normalization layer.'
+    + 'See ffn.py for examples.')
 ffns_with_norm = create_registry('llmfoundry',
                                  'ffns_with_norm',
                                  generic_type=Callable,

diff --git a/llmfoundry/registry.py b/llmfoundry/registry.py
@@ -12,7 +12,7 @@
 from transformers import PreTrainedTokenizerBase
 
 from llmfoundry.interfaces import CallbackWithConfig
-from llmfoundry.layers_registry import norms
+from llmfoundry.layers_registry import ffns, ffns_with_norm, norms
 from llmfoundry.utils.registry_utils import create_registry
 
 _loggers_description = (
@@ -121,4 +121,6 @@
     'metrics',
     'dataloaders',
     'norms',
+    'ffns',
+    'ffns_with_norm',
 ]
diff --git a/tests/test_registry.py b/tests/test_registry.py
@@ -31,6 +31,8 @@ def test_expected_registries_exist():
         'metrics',
         'models',
         'norms',
+        'ffns',
+        'ffns_with_norm',
     }
 
     assert existing_registries == expected_registry_names