Lint

Author: MartinGleize

src/fairseq2/models/opt/__init__.py

@@ -6,16 +6,10 @@
 
 from __future__ import annotations
 
-from fairseq2.models.opt.interop import (
-    convert_opt_state_dict as convert_opt_state_dict,
-)
 from fairseq2.models.opt.config import OPT_MODEL_FAMILY as OPT_MODEL_FAMILY
 from fairseq2.models.opt.config import OPTConfig as OPTConfig
-from fairseq2.models.opt.config import (
-    register_opt_configs as register_opt_configs,
-)
+from fairseq2.models.opt.config import register_opt_configs as register_opt_configs
 from fairseq2.models.opt.factory import OPTFactory as OPTFactory
-from fairseq2.models.opt.factory import (
-    create_opt_model as create_opt_model,
-)
+from fairseq2.models.opt.factory import create_opt_model as create_opt_model
 from fairseq2.models.opt.hub import get_opt_model_hub as get_opt_model_hub
+from fairseq2.models.opt.interop import convert_opt_state_dict as convert_opt_state_dict

src/fairseq2/models/opt/factory.py

@@ -33,7 +33,6 @@
     Linear,
     PositionEncoder,
     Projection,
-    RMSNorm,
     StandardEmbedding,
     StandardLayerNorm,
 )
@@ -156,7 +155,9 @@ def create_self_attention(self) -> MultiheadAttention:
     def create_ffn(self) -> FeedForwardNetwork:
         config = self._config
 
-        return StandardFeedForwardNetwork(config.model_dim, config.ffn_inner_dim, bias=True)
+        return StandardFeedForwardNetwork(
+            config.model_dim, config.ffn_inner_dim, bias=True
+        )
 
     def create_layer_norm(self) -> LayerNorm:
         config = self._config

src/fairseq2/models/opt/hub.py

@@ -11,6 +11,6 @@
 
 # isort: split
 
-from fairseq2.models.opt.config import OPTConfig, OPT_MODEL_FAMILY
+from fairseq2.models.opt.config import OPT_MODEL_FAMILY, OPTConfig
 
 get_opt_model_hub = ModelHubAccessor(OPT_MODEL_FAMILY, TransformerLM, OPTConfig)

src/fairseq2/models/opt/interop.py

@@ -12,7 +12,6 @@
 
 from fairseq2.models.opt.config import OPTConfig
 
-
 _OPT_HG_KEY_MAP = {
     # fmt: off
     r"^model\.decoder\.embed_tokens\.": r"decoder_frontend.embed.",

tests/unit/models/opt/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import torch
+
+from fairseq2.models.opt.config import OPTConfig
+from fairseq2.models.opt.factory import OPTFactory
+from fairseq2.nn import BatchLayout
+
+
+class TestOptFactory:
+    def test_opt_factory(self) -> None:
+        """Sanity check for factory + forward on a small model."""
+        config = OPTConfig()  # by default opt-125m
+        config.num_layers = 2
+        config.vocab_size = 258
+        config.model_dim = 24
+        config.ffn_inner_dim = 48
+
+        factory = OPTFactory(config)
+
+        model = factory.create_model()
+
+        model.eval()
+
+        _ = model.forward(
+            seqs=torch.randint(0, config.vocab_size, (2, 10)),
+            seqs_layout=BatchLayout((2, 10), [10, 10]),
+        )