optuna-integration 3.2.0__tar.gz

optuna-integration-3.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2018 Preferred Networks, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

optuna-integration-3.2.0/MANIFEST.in

@@ -0,0 +1 @@
+include LICENSE

optuna-integration-3.2.0/PKG-INFO

@@ -0,0 +1,112 @@
+Metadata-Version: 2.1
+Name: optuna-integration
+Version: 3.2.0
+Summary: Integration libraries of Optuna.
+License: MIT License
+        
+        Copyright (c) 2018 Preferred Networks, Inc.
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.7
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Description-Content-Type: text/markdown
+Provides-Extra: test
+Provides-Extra: checking
+Provides-Extra: document
+Provides-Extra: all
+License-File: LICENSE
+
+# Optuna-Integration
+
+[![Python](https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9%20%7C%203.10-blue)](https://www.python.org)
+[![GitHub license](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/optuna/optuna-integration)
+[![Codecov](https://codecov.io/gh/optuna/optuna-integration/branch/master/graph/badge.svg)](https://codecov.io/gh/optuna/optuna-integration/branch/master)
+<!-- [![pypi](https://img.shields.io/pypi/v/optuna.svg)](https://pypi.python.org/pypi/optuna-integration) -->
+<!-- [![conda](https://img.shields.io/conda/vn/conda-forge/optuna.svg)](https://anaconda.org/conda-forge/optuna-integration) -->
+[![Read the Docs](https://readthedocs.org/projects/optuna-integration/badge/?version=stable)](https://optuna-integration.readthedocs.io/en/stable/)
+
+[**Docs**](https://optuna-integration.readthedocs.io/en/stable/)
+
+*Optuna-Integration* is an integration module of [Optuna](https://github.com/optuna/optuna).
+This package allows us to use Optuna, an automatic Hyperparameter optimization software framework,
+integrated with many useful tools like PyTorch, sklearn, TensorFlow, etc.
+
+## Integrations
+
+Optuna-Integration API reference is [here](https://optuna-integration.readthedocs.io/en/stable/reference/index.html).
+
+* [AllenNLP](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#allennlp)  ([example](https://github.com/optuna/optuna-examples/tree/main/allennlp))
+* [Catalyst](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#catalyst)  ([example](https://github.com/optuna/optuna-examples/blob/main/pytorch/catalyst_simple.py))
+* [Chainer](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#chainer)  ([example](https://github.com/optuna/optuna-examples/tree/main/chainer/chainer_integration.py))
+* [ChainerMN](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#chainermn)  ([example](https://github.com/optuna/optuna-examples/tree/main/chainer/chainermn_simple.py))
+* [Keras](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#keras)  ([example](https://github.com/optuna/optuna-examples/tree/main/keras))
+* [skorch](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#skorch)  ([example](https://github.com/optuna/optuna-examples/tree/main/pytorch/skorch_simple.py))
+* [tf.keras](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#tensorflow)  ([example](https://github.com/optuna/optuna-examples/tree/main/tfkeras/tfkeras_integration.py))
+
+## Installation
+
+Optuna-Integration is available at [the Python Package Index](https://pypi.org/project/optuna-integration/) and on [Anaconda Cloud](https://anaconda.org/conda-forge/optuna-integration).
+
+```bash
+# PyPI
+$ pip install optuna-integration
+```
+
+```bash
+# Anaconda Cloud
+$ conda install -c conda-forge optuna-integration
+```
+
+Optuna-Integration supports from Python 3.7 to Python 3.10.
+
+Also, we also provide Optuna docker images on [DockerHub](https://hub.docker.com/r/optuna/optuna).
+
+## Communication
+
+- [GitHub Discussions] for questions.
+- [GitHub Issues] for bug reports and feature requests.
+
+[GitHub Discussions]: https://github.com/optuna/optuna-integration/discussions
+[GitHub issues]: https://github.com/optuna/optuna-integration/issues
+
+## Contribution
+
+Any contributions to Optuna-Integration are more than welcome!
+
+For general guidelines how to contribute to the project, take a look at [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+## Reference
+
+Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. 2019.
+Optuna: A Next-generation Hyperparameter Optimization Framework. In KDD ([arXiv](https://arxiv.org/abs/1907.10902)).

optuna-integration-3.2.0/README.md

@@ -0,0 +1,63 @@
+# Optuna-Integration
+
+[![Python](https://img.shields.io/badge/python-3.7%20%7C%203.8%20%7C%203.9%20%7C%203.10-blue)](https://www.python.org)
+[![GitHub license](https://img.shields.io/badge/license-MIT-blue.svg)](https://github.com/optuna/optuna-integration)
+[![Codecov](https://codecov.io/gh/optuna/optuna-integration/branch/master/graph/badge.svg)](https://codecov.io/gh/optuna/optuna-integration/branch/master)
+<!-- [![pypi](https://img.shields.io/pypi/v/optuna.svg)](https://pypi.python.org/pypi/optuna-integration) -->
+<!-- [![conda](https://img.shields.io/conda/vn/conda-forge/optuna.svg)](https://anaconda.org/conda-forge/optuna-integration) -->
+[![Read the Docs](https://readthedocs.org/projects/optuna-integration/badge/?version=stable)](https://optuna-integration.readthedocs.io/en/stable/)
+
+[**Docs**](https://optuna-integration.readthedocs.io/en/stable/)
+
+*Optuna-Integration* is an integration module of [Optuna](https://github.com/optuna/optuna).
+This package allows us to use Optuna, an automatic Hyperparameter optimization software framework,
+integrated with many useful tools like PyTorch, sklearn, TensorFlow, etc.
+
+## Integrations
+
+Optuna-Integration API reference is [here](https://optuna-integration.readthedocs.io/en/stable/reference/index.html).
+
+* [AllenNLP](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#allennlp)  ([example](https://github.com/optuna/optuna-examples/tree/main/allennlp))
+* [Catalyst](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#catalyst)  ([example](https://github.com/optuna/optuna-examples/blob/main/pytorch/catalyst_simple.py))
+* [Chainer](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#chainer)  ([example](https://github.com/optuna/optuna-examples/tree/main/chainer/chainer_integration.py))
+* [ChainerMN](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#chainermn)  ([example](https://github.com/optuna/optuna-examples/tree/main/chainer/chainermn_simple.py))
+* [Keras](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#keras)  ([example](https://github.com/optuna/optuna-examples/tree/main/keras))
+* [skorch](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#skorch)  ([example](https://github.com/optuna/optuna-examples/tree/main/pytorch/skorch_simple.py))
+* [tf.keras](https://optuna-integration.readthedocs.io/en/stable/reference/index.html#tensorflow)  ([example](https://github.com/optuna/optuna-examples/tree/main/tfkeras/tfkeras_integration.py))
+
+## Installation
+
+Optuna-Integration is available at [the Python Package Index](https://pypi.org/project/optuna-integration/) and on [Anaconda Cloud](https://anaconda.org/conda-forge/optuna-integration).
+
+```bash
+# PyPI
+$ pip install optuna-integration
+```
+
+```bash
+# Anaconda Cloud
+$ conda install -c conda-forge optuna-integration
+```
+
+Optuna-Integration supports from Python 3.7 to Python 3.10.
+
+Also, we also provide Optuna docker images on [DockerHub](https://hub.docker.com/r/optuna/optuna).
+
+## Communication
+
+- [GitHub Discussions] for questions.
+- [GitHub Issues] for bug reports and feature requests.
+
+[GitHub Discussions]: https://github.com/optuna/optuna-integration/discussions
+[GitHub issues]: https://github.com/optuna/optuna-integration/issues
+
+## Contribution
+
+Any contributions to Optuna-Integration are more than welcome!
+
+For general guidelines how to contribute to the project, take a look at [CONTRIBUTING.md](./CONTRIBUTING.md).
+
+## Reference
+
+Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. 2019.
+Optuna: A Next-generation Hyperparameter Optimization Framework. In KDD ([arXiv](https://arxiv.org/abs/1907.10902)).

optuna-integration-3.2.0/optuna_integration/_imports.py

@@ -0,0 +1,125 @@
+import importlib
+import types
+from types import TracebackType
+from typing import Any
+from typing import Optional
+from typing import Tuple
+from typing import Type
+
+
+class _DeferredImportExceptionContextManager:
+    """Context manager to defer exceptions from imports.
+
+    Catches :exc:`ImportError` and :exc:`SyntaxError`.
+    If any exception is caught, this class raises an :exc:`ImportError` when being checked.
+
+    """
+
+    def __init__(self) -> None:
+        self._deferred: Optional[Tuple[Exception, str]] = None
+
+    def __enter__(self) -> "_DeferredImportExceptionContextManager":
+        """Enter the context manager.
+
+        Returns:
+            Itself.
+
+        """
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[Exception]],
+        exc_value: Optional[Exception],
+        traceback: Optional[TracebackType],
+    ) -> Optional[bool]:
+        """Exit the context manager.
+
+        Args:
+            exc_type:
+                Raised exception type. :obj:`None` if nothing is raised.
+            exc_value:
+                Raised exception object. :obj:`None` if nothing is raised.
+            traceback:
+                Associated traceback. :obj:`None` if nothing is raised.
+
+        Returns:
+            :obj:`None` if nothing is deferred, otherwise :obj:`True`.
+            :obj:`True` will suppress any exceptions avoiding them from propagating.
+
+        """
+        if isinstance(exc_value, (ImportError, SyntaxError)):
+            if isinstance(exc_value, ImportError):
+                message = (
+                    "Tried to import '{}' but failed. Please make sure that the package is "
+                    "installed correctly to use this feature. Actual error: {}."
+                ).format(exc_value.name, exc_value)
+            elif isinstance(exc_value, SyntaxError):
+                message = (
+                    "Tried to import a package but failed due to a syntax error in {}. Please "
+                    "make sure that the Python version is correct to use this feature. Actual "
+                    "error: {}."
+                ).format(exc_value.filename, exc_value)
+            else:
+                assert False
+
+            self._deferred = (exc_value, message)
+            return True
+        return None
+
+    def is_successful(self) -> bool:
+        """Return whether the context manager has caught any exceptions.
+
+        Returns:
+            :obj:`True` if no exceptions are caught, :obj:`False` otherwise.
+
+        """
+        return self._deferred is None
+
+    def check(self) -> None:
+        """Check whether the context manager has caught any exceptions.
+
+        Raises:
+            :exc:`ImportError`:
+                If any exception was caught from the caught exception.
+
+        """
+        if self._deferred is not None:
+            exc_value, message = self._deferred
+            raise ImportError(message) from exc_value
+
+
+def try_import() -> _DeferredImportExceptionContextManager:
+    """Create a context manager that can wrap imports of optional packages to defer exceptions.
+
+    Returns:
+        Deferred import context manager.
+
+    """
+    return _DeferredImportExceptionContextManager()
+
+
+class _LazyImport(types.ModuleType):
+    """Module wrapper for lazy import.
+
+    This class wraps the specified modules and lazily imports them only when accessed.
+    Otherwise, `import optuna-integration` is slowed down by importing all submodules and
+    dependencies even if not required.
+    Within this project's usage, importlib override this module's attribute on the first
+    access and the imported submodule is directly accessed from the second access.
+
+    Args:
+        name: Name of module to apply lazy import.
+    """
+
+    def __init__(self, name: str) -> None:
+        super().__init__(name)
+        self._name = name
+
+    def _load(self) -> types.ModuleType:
+        module = importlib.import_module(self._name)
+        self.__dict__.update(module.__dict__)
+        return module
+
+    def __getattr__(self, item: str) -> Any:
+        return getattr(self._load(), item)

optuna-integration-3.2.0/optuna_integration/allennlp/__init__.py

@@ -0,0 +1,6 @@
+from optuna_integration.allennlp._dump_best_config import dump_best_config
+from optuna_integration.allennlp._executor import AllenNLPExecutor
+from optuna_integration.allennlp._pruner import AllenNLPPruningCallback
+
+
+__all__ = ["dump_best_config", "AllenNLPExecutor", "AllenNLPPruningCallback"]

optuna-integration-3.2.0/optuna_integration/allennlp/_dump_best_config.py

@@ -0,0 +1,61 @@
+import json
+
+import optuna
+
+from optuna_integration._imports import try_import
+from optuna_integration.allennlp._environment import _environment_variables
+
+
+with try_import() as _imports:
+    import _jsonnet
+
+
+def dump_best_config(input_config_file: str, output_config_file: str, study: optuna.Study) -> None:
+    """Save JSON config file with environment variables and best performing hyperparameters.
+
+    Args:
+        input_config_file:
+            Input Jsonnet config file used with
+            :class:`~optuna_integration.AllenNLPExecutor`.
+        output_config_file:
+            Output JSON config file.
+        study:
+            Instance of :class:`~optuna.study.Study`.
+            Note that :func:`~optuna.study.Study.optimize` must have been called.
+
+    """
+    _imports.check()
+
+    # Get environment variables.
+    ext_vars = _environment_variables()
+
+    # Get the best hyperparameters.
+    best_params = study.best_params
+    for key, value in best_params.items():
+        best_params[key] = str(value)
+
+    # If keys both appear in environment variables and best_params,
+    # values in environment variables are overwritten, which means best_params is prioritized.
+    ext_vars.update(best_params)
+
+    best_config = json.loads(_jsonnet.evaluate_file(input_config_file, ext_vars=ext_vars))
+
+    # `optuna_pruner` only works with Optuna.
+    # It removes when dumping configuration since
+    # the result of `dump_best_config` can be passed to
+    # `allennlp train`.
+    if "callbacks" in best_config["trainer"]:
+        new_callbacks = []
+        callbacks = best_config["trainer"]["callbacks"]
+        for callback in callbacks:
+            if callback["type"] == "optuna_pruner":
+                continue
+            new_callbacks.append(callback)
+
+        if len(new_callbacks) == 0:
+            best_config["trainer"].pop("callbacks")
+        else:
+            best_config["trainer"]["callbacks"] = new_callbacks
+
+    with open(output_config_file, "w") as f:
+        json.dump(best_config, f, indent=4)

optuna-integration-3.2.0/optuna_integration/allennlp/_environment.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+import os
+
+
+def _is_encodable(value: str) -> bool:
+    # https://github.com/allenai/allennlp/blob/master/allennlp/common/params.py#L77-L85
+    return (value == "") or (value.encode("utf-8", "ignore") != b"")
+
+
+def _environment_variables() -> dict[str, str]:
+    return {key: value for key, value in os.environ.items() if _is_encodable(value)}

optuna-integration-3.2.0/optuna_integration/allennlp/_executor.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+import warnings
+
+import optuna
+from optuna import TrialPruned
+from optuna._experimental import experimental_class
+
+from optuna_integration._imports import try_import
+from optuna_integration.allennlp._environment import _environment_variables
+from optuna_integration.allennlp._variables import _VariableManager
+from optuna_integration.allennlp._variables import OPTUNA_ALLENNLP_DISTRIBUTED_FLAG
+
+
+with try_import() as _imports:
+    import allennlp
+    import allennlp.commands
+    import allennlp.common.cached_transformers
+    import allennlp.common.util
+
+# TrainerCallback is conditionally imported because allennlp may be unavailable in
+# the environment that builds the documentation.
+if _imports.is_successful():
+    import _jsonnet
+    import psutil
+    from torch.multiprocessing.spawn import ProcessRaisedException
+
+
+def _fetch_pruner_config(trial: optuna.Trial) -> dict[str, Any]:
+    pruner = trial.study.pruner
+    kwargs: dict[str, Any] = {}
+
+    if isinstance(pruner, optuna.pruners.HyperbandPruner):
+        kwargs["min_resource"] = pruner._min_resource
+        kwargs["max_resource"] = pruner._max_resource
+        kwargs["reduction_factor"] = pruner._reduction_factor
+
+    elif isinstance(pruner, optuna.pruners.MedianPruner):
+        kwargs["n_startup_trials"] = pruner._n_startup_trials
+        kwargs["n_warmup_steps"] = pruner._n_warmup_steps
+        kwargs["interval_steps"] = pruner._interval_steps
+
+    elif isinstance(pruner, optuna.pruners.PercentilePruner):
+        kwargs["percentile"] = pruner._percentile
+        kwargs["n_startup_trials"] = pruner._n_startup_trials
+        kwargs["n_warmup_steps"] = pruner._n_warmup_steps
+        kwargs["interval_steps"] = pruner._interval_steps
+
+    elif isinstance(pruner, optuna.pruners.SuccessiveHalvingPruner):
+        kwargs["min_resource"] = pruner._min_resource
+        kwargs["reduction_factor"] = pruner._reduction_factor
+        kwargs["min_early_stopping_rate"] = pruner._min_early_stopping_rate
+
+    elif isinstance(pruner, optuna.pruners.ThresholdPruner):
+        kwargs["lower"] = pruner._lower
+        kwargs["upper"] = pruner._upper
+        kwargs["n_warmup_steps"] = pruner._n_warmup_steps
+        kwargs["interval_steps"] = pruner._interval_steps
+    elif isinstance(pruner, optuna.pruners.NopPruner):
+        pass
+    else:
+        raise ValueError("Unsupported pruner is specified: {}".format(type(pruner)))
+
+    return kwargs
+
+
+@experimental_class("1.4.0")
+class AllenNLPExecutor:
+    """AllenNLP extension to use optuna with Jsonnet config file.
+
+    See the examples of `objective function <https://github.com/optuna/optuna-examples/tree/
+    main/allennlp/allennlp_jsonnet.py>`_.
+
+    You can also see the tutorial of our AllenNLP integration on
+    `AllenNLP Guide <https://guide.allennlp.org/hyperparameter-optimization>`_.
+
+    .. note::
+        From Optuna v2.1.0, users have to cast their parameters by using methods in Jsonnet.
+        Call ``std.parseInt`` for integer, or ``std.parseJson`` for floating point.
+        Please see the `example configuration <https://github.com/optuna/optuna-examples/tree/main/
+        allennlp/classifier.jsonnet>`_.
+
+    .. note::
+        In :class:`~optuna_integration.AllenNLPExecutor`,
+        you can pass parameters to AllenNLP by either defining a search space using
+        Optuna suggest methods or setting environment variables just like AllenNLP CLI.
+        If a value is set in both a search space in Optuna and the environment variables,
+        the executor will use the value specified in the search space in Optuna.
+
+    Args:
+        trial:
+            A :class:`~optuna.trial.Trial` corresponding to the current evaluation
+            of the objective function.
+        config_file:
+            Config file for AllenNLP.
+            Hyperparameters should be masked with ``std.extVar``.
+            Please refer to `the config example <https://github.com/allenai/allentune/blob/
+            master/examples/classifier.jsonnet>`_.
+        serialization_dir:
+            A path which model weights and logs are saved.
+        metrics:
+            An evaluation metric. `GradientDescrentTrainer.train() <https://docs.allennlp.org/
+            main/api/training/gradient_descent_trainer/#train>`_ of AllenNLP
+            returns a dictionary containing metrics after training.
+            :class:`~optuna_integration.AllenNLPExecutor` accesses the dictionary
+            by the key ``metrics`` you specify and use it as a objective value.
+        force:
+            If :obj:`True`, an executor overwrites the output directory if it exists.
+        file_friendly_logging:
+            If :obj:`True`, tqdm status is printed on separate lines and slows tqdm refresh rate.
+        include_package:
+            Additional packages to include.
+            For more information, please see
+            `AllenNLP documentation <https://docs.allennlp.org/master/api/commands/train/>`_.
+
+    """
+
+    def __init__(
+        self,
+        trial: optuna.Trial,
+        config_file: str,
+        serialization_dir: str,
+        metrics: str = "best_validation_accuracy",
+        *,
+        include_package: str | list[str] | None = None,
+        force: bool = False,
+        file_friendly_logging: bool = False,
+    ):
+        _imports.check()
+
+        self._params = trial.params
+        self._config_file = config_file
+        self._serialization_dir = serialization_dir
+        self._metrics = metrics
+        self._force = force
+        self._file_friendly_logging = file_friendly_logging
+
+        if include_package is None:
+            include_package = []
+        if isinstance(include_package, str):
+            include_package = [include_package]
+
+        self._include_package = include_package + ["optuna_integration.allennlp"]
+
+        storage = trial.study._storage
+
+        if isinstance(storage, optuna.storages.RDBStorage):
+            url = storage.url
+
+        elif isinstance(storage, optuna.storages._CachedStorage):
+            assert isinstance(storage._backend, optuna.storages.RDBStorage)
+            url = storage._backend.url
+
+        else:
+            url = ""
+
+        target_pid = psutil.Process().ppid()
+        variable_manager = _VariableManager(target_pid)
+
+        pruner_kwargs = _fetch_pruner_config(trial)
+        variable_manager.set_value("study_name", trial.study.study_name)
+        variable_manager.set_value("trial_id", trial._trial_id)
+        variable_manager.set_value("storage_name", url)
+        variable_manager.set_value("monitor", metrics)
+
+        if trial.study.pruner is not None:
+            variable_manager.set_value("pruner_class", type(trial.study.pruner).__name__)
+            variable_manager.set_value("pruner_kwargs", pruner_kwargs)
+
+    def _build_params(self) -> dict[str, Any]:
+        """Create a dict of params for AllenNLP.
+
+        _build_params is based on allentune's ``train_func``.
+        For more detail, please refer to
+        https://github.com/allenai/allentune/blob/master/allentune/modules/allennlp_runner.py#L34-L65
+
+        """
+        params = _environment_variables()
+        params.update({key: str(value) for key, value in self._params.items()})
+        return json.loads(_jsonnet.evaluate_file(self._config_file, ext_vars=params))
+
+    def _set_environment_variables(self) -> None:
+        for key, value in _environment_variables().items():
+            if key is None:
+                continue
+            os.environ[key] = value
+
+    def run(self) -> float:
+        """Train a model using AllenNLP."""
+        for package_name in self._include_package:
+            allennlp.common.util.import_module_and_submodules(package_name)
+
+        # Without the following lines, the transformer model construction only takes place in the
+        # first trial (which would consume some random numbers), and the cached model will be used
+        # in trials afterwards (which would not consume random numbers), leading to inconsistent
+        # results between single trial and multiple trials. To make results reproducible in
+        # multiple trials, we clear the cache before each trial.
+        # TODO(MagiaSN) When AllenNLP has introduced a better API to do this, one should remove
+        # these lines and use the new API instead. For example, use the `_clear_caches()` method
+        # which will be in the next AllenNLP release after 2.4.0.
+        allennlp.common.cached_transformers._model_cache.clear()
+        allennlp.common.cached_transformers._tokenizer_cache.clear()
+
+        self._set_environment_variables()
+        params = allennlp.common.params.Params(self._build_params())
+
+        if "distributed" in params:
+            if OPTUNA_ALLENNLP_DISTRIBUTED_FLAG in os.environ:
+                warnings.warn(
+                    "Other process may already exists."
+                    " If you have trouble, please unset the environment"
+                    " variable `OPTUNA_ALLENNLP_USE_DISTRIBUTED`"
+                    " and try it again."
+                )
+
+            os.environ[OPTUNA_ALLENNLP_DISTRIBUTED_FLAG] = "1"
+
+        try:
+            allennlp.commands.train.train_model(
+                params=params,
+                serialization_dir=self._serialization_dir,
+                file_friendly_logging=self._file_friendly_logging,
+                force=self._force,
+                include_package=self._include_package,
+            )
+        except ProcessRaisedException as e:
+            if "raise TrialPruned()" in str(e):
+                raise TrialPruned()
+
+        metrics = json.load(open(os.path.join(self._serialization_dir, "metrics.json")))
+        return metrics[self._metrics]