split lightning_utils into schedulers and callbacks

webis-de · Dec 16, 2024 · 044ab99 · 044ab99
1 parent 17d4eb7
commit 044ab99
Show file tree

Hide file tree

Showing 7 changed files with 26 additions and 18 deletions.
diff --git a/lightning_ir/__init__.py b/lightning_ir/__init__.py
@@ -21,6 +21,7 @@
     BiEncoderTokenizer,
     ScoringFunction,
 )
+from .callbacks import IndexCallback, RankCallback, RegisterLocalDatasetCallback, ReRankCallback, SearchCallback
 from .cross_encoder import (
     CrossEncoderConfig,
     CrossEncoderModel,
@@ -42,18 +43,6 @@
     TrainBatch,
     TupleDataset,
 )
-from .lightning_utils import (
-    ConstantLRSchedulerWithLinearWarmup,
-    GenericConstantSchedulerWithLinearWarmup,
-    GenericConstantSchedulerWithQuadraticWarmup,
-    GenericLinearSchedulerWithLinearWarmup,
-    IndexCallback,
-    LinearLRSchedulerWithLinearWarmup,
-    RankCallback,
-    ReRankCallback,
-    SearchCallback,
-    WarmupLRScheduler,
-)
 from .loss import (
     ApproxMRR,
     ApproxNDCG,
@@ -102,6 +91,14 @@
     SparseSearchConfig,
     SparseSearcher,
 )
+from .schedulers import (
+    ConstantLRSchedulerWithLinearWarmup,
+    GenericConstantSchedulerWithLinearWarmup,
+    GenericConstantSchedulerWithQuadraticWarmup,
+    GenericLinearSchedulerWithLinearWarmup,
+    LinearLRSchedulerWithLinearWarmup,
+    WarmupLRScheduler,
+)
 
 AutoConfig.register(BiEncoderConfig.model_type, BiEncoderConfig)
 AutoModel.register(BiEncoderConfig, BiEncoderModel)
@@ -188,6 +185,7 @@
     "RankBatch",
     "RankNet",
     "RankSample",
+    "RegisterLocalDatasetCallback",
     "ReRankCallback",
     "RunDataset",
     "ScoreBasedInBatchCrossEntropy",

diff --git a/lightning_ir/callbacks/__init__.py b/lightning_ir/callbacks/__init__.py
@@ -0,0 +1,9 @@
+from .callbacks import IndexCallback, RankCallback, RegisterLocalDatasetCallback, ReRankCallback, SearchCallback
+
+__all__ = [
+    "IndexCallback",
+    "RankCallback",
+    "RegisterLocalDatasetCallback",
+    "ReRankCallback",
+    "SearchCallback",
+]
diff --git a/lightning_ir/lightning_utils/callbacks.py → lightning_ir/callbacks/callbacks.py b/lightning_ir/lightning_utils/callbacks.py → lightning_ir/callbacks/callbacks.py
diff --git a/lightning_ir/main.py b/lightning_ir/main.py
@@ -11,7 +11,7 @@
 from typing_extensions import override
 
 import lightning_ir  # noqa: F401
-from lightning_ir.lightning_utils.lr_schedulers import LR_SCHEDULERS, WarmupLRScheduler
+from lightning_ir.schedulers.lr_schedulers import LR_SCHEDULERS, WarmupLRScheduler
 
 if torch.cuda.is_available():
     torch.set_float32_matmul_precision("medium")

diff --git a/lightning_ir/lightning_utils/__init__.py → lightning_ir/schedulers/__init__.py b/lightning_ir/lightning_utils/__init__.py → lightning_ir/schedulers/__init__.py
@@ -1,4 +1,9 @@
-from .callbacks import IndexCallback, RankCallback, ReRankCallback, SearchCallback
+"""
+Module containing utility classes and functions for PyTorch Lightning.
+
+This module provides callbacks .
+"""
+
 from .lr_schedulers import ConstantLRSchedulerWithLinearWarmup, LinearLRSchedulerWithLinearWarmup, WarmupLRScheduler
 from .schedulers import (
     GenericConstantSchedulerWithLinearWarmup,
@@ -11,10 +16,6 @@
     "GenericConstantSchedulerWithLinearWarmup",
     "GenericConstantSchedulerWithQuadraticWarmup",
     "GenericLinearSchedulerWithLinearWarmup",
-    "IndexCallback",
     "LinearLRSchedulerWithLinearWarmup",
-    "RankCallback",
-    "ReRankCallback",
-    "SearchCallback",
     "WarmupLRScheduler",
 ]
diff --git a/...tning_ir/lightning_utils/lr_schedulers.py → lightning_ir/schedulers/lr_schedulers.py b/...tning_ir/lightning_utils/lr_schedulers.py → lightning_ir/schedulers/lr_schedulers.py
diff --git a/lightning_ir/lightning_utils/schedulers.py → lightning_ir/schedulers/schedulers.py b/lightning_ir/lightning_utils/schedulers.py → lightning_ir/schedulers/schedulers.py