Defaultflash (#2537)

* set flash as default * remove model_opt from test
OpenNMT · Dec 18, 2023 · 23f7916 · 23f7916
1 parent f01bea1
commit 23f7916
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/onmt/opts.py b/onmt/opts.py
@@ -1826,15 +1826,21 @@ def translate_opts(parser, dynamic=False):
         "Necessary for models whose output layers can assign "
         "zero probability.",
     )
-
+    group.add(
+        "--self_attn_type",
+        "-self_attn_type",
+        type=str,
+        default="scaled-dot-flash",
+        help="Self attention type in Transformer decoder "
+        'layer -- currently "scaled-dot", "scaled-dot-flash" or "average" ',
+    )
     group = parser.add_argument_group("Data")
     group.add(
         "--data_type",
         "-data_type",
         default="text",
         help="Type of the source input. Options: [text].",
     )
-
     group.add(
         "--src",
         "-src",

diff --git a/onmt/tests/test_inference_engines.py b/onmt/tests/test_inference_engines.py
@@ -9,7 +9,6 @@ def _get_parser():
     parser = ArgumentParser(description="simple_inference_engine_py.py")
     opts.config_opts(parser)
     opts.translate_opts(parser, dynamic=True)
-    opts.model_opts(parser)
     return parser