Move FewShotEvaluator to contrib

kauldron/evals/__init__.py

@@ -18,7 +18,6 @@
 
 from kauldron.evals.evaluators import Evaluator
 from kauldron.evals.evaluators import EvaluatorBase
-from kauldron.evals.fewshot_evaluator import FewShotEvaluator
 
 # RunStrategy are available in both XM and Kauldron side
 from kauldron.xm._src.run_strategies import RunEvery

kauldron/kd.py

@@ -29,6 +29,8 @@
   pytest.register_assert_rewrite('kauldron.utils.assert_utils')
 del pytest, sys
 
+from etils import epy as _epy
+
 # Namespaces
 from kauldron import checkpoints as ckpts
 from kauldron import data
@@ -51,6 +53,12 @@
 # Register the default `ConfigDict` overwrites
 from kauldron.xm._src import default_values as _default_values
 
+# Import contrib at the end as they can use all
+
+with _epy.lazy_imports(
+):
+  from kauldron import contrib  # pylint: disable=g-bad-import-order  # pytype: disable=import-error
+
 # TODO(epot): This could be optional for the top-level module
 # Automated documentation info
 # See: https://github.com/conchylicultor/sphinx-apitree

kauldron/modules/misc.py

@@ -62,6 +62,7 @@ class Rearrange(nn.Module):
   """Wrapper around `einops.rearrange` for usage e.g. in `nn.Sequential`.
 
   Example:
+
   ```
   cfg.model = kd.nn.Sequential(
       inputs="batch.image",
@@ -93,6 +94,7 @@ class Reduce(nn.Module):
   """Wrapper around `einops.reduce` for usage e.g. in `nn.Sequential`.
 
   Example:
+
   ```
   cfg.model = kd.nn.Sequential(
       inputs="batch.image",