Merge pull request #30 from naist-nlp/plugin

Implement plug-in module loader

Author: de9uch1

docs/index.rst

@@ -30,6 +30,7 @@ mbrs: A library for MBR decoding
       - :doc:`design`
       - :doc:`custom_metric`
       - :doc:`custom_decoder`
+      - :doc:`plugin`
       - :doc:`timer`
 
    .. grid-item-card:: :material-regular:`library_books;2em` References
@@ -66,6 +67,7 @@ mbrs: A library for MBR decoding
    design
    custom_metric
    custom_decoder
+   plugin
    timer
 
 .. toctree::

docs/plugin.rst

@@ -0,0 +1,59 @@
+Plug-in loader
+==============
+
+:code:`mbrs-decode` and :code:`mbrs-score` load plug-in modules via the :code:`--plugin_dir` option.
+
+Examples
+~~~~~~~~
+
+This tutorial explains how to load a user defined modules.
+
+.. seealso::
+
+   :doc:`How to define a new metric <./custom_metric>`
+        Detailed documentation of the metric customization.
+
+   :doc:`How to define a new decoder <./custom_decoder>`
+        Detailed documentation of the decoder customization.
+
+1. Define a new metric, decoder, or selector with :code:`@register` decorator.
+
+   .. code-block:: python
+      :emphasize-lines: 4
+
+      from mbrs.metrics import register, Metric, MetricBLEU
+
+
+      @register("my_bleu")
+      class MetricMyBLEU(MetricBLEU):
+          ...
+
+2. Prepare :code:`__init__.py` to specify classes to be loaded.
+
+   .. code-block:: python
+      :emphasize-lines: 1
+
+      from .new import MetricNew
+
+3. Then, load the modules with :code:`--plugin_dir` option with a path to the directory containing the above :code:`__init__.py`.
+
+   .. code-block:: bash
+      :emphasize-lines: 2,6
+
+      mbrs-decode \
+        --plugin_dir path/to/plugins/ \
+        hypotheses.txt \
+        --num_candidates 1024 \
+        --decoder mbr \
+        --metric my_bleu
+
+   :code:`mbrs-score` also supports plug-in loading.
+
+   .. code-block:: bash
+      :emphasize-lines: 2,5
+
+      mbrs-score \
+        --plugin_dir path/to/plugins/ \
+        hypotheses.txt \
+        -r references.txt \
+        --metric my_bleu

mbrs/args.py

@@ -1,5 +1,9 @@
+from __future__ import annotations
+
 import argparse
+import importlib
 import logging
+import os
 import sys
 from argparse import Namespace
 from pathlib import Path
@@ -12,6 +16,47 @@
 
 
 class ArgumentParser(simple_parsing.ArgumentParser):
+    IMPORTED_MODULES: set[Path] = set()
+
+    def add_plugin_argumnets(self, parser: ArgumentParser) -> None:
+        """Add arguments for plugins.
+
+        Args:
+            parser (ArgumentParser): Argument parser.
+        """
+        parser.add_argument(
+            "--plugin_dir",
+            type=Path,
+            default=None,
+            help="Path to a directory containing user defined plugins.",
+        )
+
+    def import_plugin(self, plugin_dir: Path) -> None:
+        """Import plugin modules.
+
+        Args:
+            plugin_dir (pathlib.Path): A directory containing user defined plugins.
+        """
+        plugin_dir = plugin_dir.absolute()
+        if not os.path.exists(plugin_dir) or not os.path.isdir(
+            os.path.dirname(plugin_dir)
+        ):
+            raise FileNotFoundError(plugin_dir)
+
+        module_parent, module_name = os.path.split(plugin_dir)
+        if plugin_dir not in ArgumentParser.IMPORTED_MODULES:
+            if module_name not in sys.modules:
+                sys.path.insert(0, module_parent)
+                importlib.import_module(module_name)
+            elif plugin_dir in sys.modules[module_name].__path__:
+                logger.info(f"--plugin_dir={plugin_dir} has already been imported.")
+            else:
+                raise ImportError(
+                    f"Failed to import --plugin_dir={plugin_dir} because the module name "
+                    f"({module_name}) is not globally unique."
+                )
+            self.IMPORTED_MODULES.add(plugin_dir)
+
     def preprocess_parser(self) -> None:
         """Preprocess ArgumentParser."""
         self.parse_known_args_preprocess(sys.argv[1:])
@@ -67,6 +112,14 @@ def parse_known_args_preprocess(
                     help="Path to a config file containing default values to use.",
                 )
 
+            # Plugin loader
+            self.add_plugin_argumnets(temp_parser)
+            args_with_plugin_dir, args = temp_parser.parse_known_args(args)
+            plugin_dir: Path | None = args_with_plugin_dir.plugin_dir
+            if plugin_dir is not None:
+                self.import_plugin(plugin_dir)
+            self.add_plugin_argumnets(self)
+
         assert isinstance(args, list)
         self._preprocessing(args=args, namespace=namespace)
 

mbrs/args_test.py

@@ -1,11 +1,44 @@
+import os
 import pathlib
 
+import pytest
 import yaml
 
 from mbrs.cli.decode import get_argparser
+from mbrs.metrics import get_metric
 
 
 class TestArgumentParser:
+    def test_plugin_load(self, tmp_path: pathlib.Path):
+        config_path = tmp_path / "config.yaml"
+        hyps_path = tmp_path / "hyps.txt"
+        cfg_dict = {
+            "common": {
+                "metric": "my_bleu",
+                "hypotheses": str(hyps_path),
+                "num_candidates": 2,
+            },
+        }
+
+        with open(config_path, mode="w") as f:
+            yaml.dump(cfg_dict, f)
+
+        with open(hyps_path, mode="w") as f:
+            f.writelines(["tests", "a test"])
+
+        cmd_args = ["--config_path", str(config_path)]
+        with pytest.raises(NotImplementedError):
+            parser = get_argparser(cmd_args)
+
+        plugin_dir = os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "tests", "plugins"
+        )
+        cmd_args += ["--plugin_dir", plugin_dir]
+        parser = get_argparser(cmd_args)
+        args = parser.parse_args(args=cmd_args)
+        assert args.common.metric == "my_bleu"
+        assert get_metric(args.common.metric).__name__ == "MetricMyBLEU"
+
     def test_config_load(self, tmp_path: pathlib.Path):
         config_path = tmp_path / "config.yaml"
         hyps_path = tmp_path / "hyps.txt"

mbrs/cli/decode.py

@@ -14,7 +14,7 @@
     format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
     datefmt="%Y-%m-%d %H:%M:%S",
     level=os.environ.get("LOGLEVEL", "INFO").upper(),
-    stream=sys.stdout,
+    stream=sys.stderr,
 )
 logger = logging.getLogger(__name__)
 
@@ -36,6 +36,7 @@
 from mbrs.metrics import Metric, MetricEnum, get_metric
 from mbrs.selectors import Selector, get_selector
 
+
 simple_parsing.parsing.logger.setLevel(logging.ERROR)
 dataclass_wrapper.logger.setLevel(logging.ERROR)
 
@@ -66,11 +67,17 @@ class CommonArguments:
     # Number of references for each sentence.
     num_references: int | None = field(default=None)
     # Type of the decoder.
-    decoder: str = choice(*registry.get_registry("decoder").keys(), default="mbr")
+    decoder: str = field(
+        default="mbr", metadata={"choices": registry.get_registry("decoder")}
+    )
     # Type of the metric.
-    metric: str = choice(*registry.get_registry("metric").keys(), default="bleu")
+    metric: str = field(
+        default="bleu", metadata={"choices": registry.get_registry("metric")}
+    )
     # Type of the selector.
-    selector: str = choice(*registry.get_registry("selector").keys(), default="nbest")
+    selector: str = field(
+        default="nbest", metadata={"choices": registry.get_registry("selector")}
+    )
     # Return the n-best hypotheses.
     nbest: int = field(default=1)
     # No verbose information and report.
@@ -174,16 +181,11 @@ def main(args: Namespace) -> None:
             reference_lprobs = f.readlines()
         assert len(references) == len(reference_lprobs)
 
-    metric_type = get_metric(args.common.metric)
-    metric: Metric = metric_type(args.metric)
-
-    selector_type = get_selector(args.common.selector)
-    selector: Selector = selector_type(args.selector)
-
-    decoder_type = get_decoder(args.common.decoder)
-    decoder: DecoderReferenceBased | DecoderReferenceless = decoder_type(
-        args.decoder, metric, selector
-    )
+    metric: Metric = get_metric(args.common.metric)(args.metric)
+    selector: Selector = get_selector(args.common.selector)(args.selector)
+    decoder: DecoderReferenceBased | DecoderReferenceless = get_decoder(
+        args.common.decoder
+    )(args.decoder, metric, selector)
 
     num_cands = args.common.num_candidates
     num_refs = args.common.num_references or num_cands

mbrs/cli/score.py

@@ -9,26 +9,25 @@
 from dataclasses import asdict, dataclass
 from typing import Sequence
 
+logging.basicConfig(
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO").upper(),
+    stream=sys.stderr,
+)
+logger = logging.getLogger(__name__)
+
 import simple_parsing
 from simple_parsing import choice, field, flag
 from simple_parsing.wrappers import dataclass_wrapper
 
 from mbrs import registry
 from mbrs.args import ArgumentParser
-from mbrs.metrics import Metric, get_metric
-from mbrs.metrics.base import MetricReferenceless
+from mbrs.metrics import Metric, MetricReferenceless, get_metric
 
 simple_parsing.parsing.logger.setLevel(logging.ERROR)
 dataclass_wrapper.logger.setLevel(logging.ERROR)
 
-logging.basicConfig(
-    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S",
-    level=os.environ.get("LOGLEVEL", "INFO").upper(),
-    stream=sys.stdout,
-)
-logger = logging.getLogger(__name__)
-
 
 class Format(enum.Enum):
     plain = "plain"
@@ -48,7 +47,9 @@ class CommonArguments:
     # Output format.
     format: Format = choice(Format, default=Format.json)
     # Type of the metric.
-    metric: str = choice(*registry.get_registry("metric").keys(), default="bleu")
+    metric: str = field(
+        default="bleu", metadata={"choices": registry.get_registry("metric")}
+    )
     # No verbose information and report.
     quiet: bool = flag(default=False)
     # Number of digits for values of float point.
@@ -61,11 +62,12 @@ def get_argparser(args: Sequence[str] | None = None) -> ArgumentParser:
     for _field in meta_parser._wrappers[0].fields:
         _field.required = False
     known_args, _ = meta_parser.parse_known_args(args=args)
-    metric_type = get_metric(known_args.common.metric)
 
     parser = ArgumentParser(add_help=True, add_config_path_arg=True)
     parser.add_arguments(CommonArguments, "common")
-    parser.add_arguments(metric_type.Config, "metric", prefix="metric.")
+    parser.add_arguments(
+        get_metric(known_args.common.metric).Config, "metric", prefix="metric."
+    )
     return parser
 
 
@@ -92,8 +94,7 @@ def main(args: Namespace) -> None:
             references = f.readlines()
         assert num_sents == len(references)
 
-    metric_type = get_metric(args.common.metric)
-    metric: Metric = metric_type(args.metric)
+    metric: Metric | MetricReferenceless = get_metric(args.common.metric)(args.metric)
 
     if isinstance(metric, MetricReferenceless):
         assert sources is not None

mbrs/tests/plugins/__init__.py

@@ -0,0 +1 @@
+from .my_bleu import MetricMyBLEU

mbrs/tests/plugins/my_bleu.py

@@ -0,0 +1,6 @@
+from mbrs.metrics import MetricBLEU, register
+
+
+@register("my_bleu")
+class MetricMyBLEU(MetricBLEU):
+    """My customized metric class."""