add trainer

webis-de · Aug 29, 2024 · b076c21 · b076c21
1 parent ba3337c
commit b076c21
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/lightning_ir/__init__.py b/lightning_ir/__init__.py
@@ -61,6 +61,7 @@
     RankNet,
     SupervisedMarginMSE,
 )
+from .main import LightningIRTrainer
 from .models import ColConfig, ColModel, SpladeConfig, SpladeModel, XTRConfig, XTRModel
 from .retrieve import (
     FaissFlatIndexConfig,
@@ -146,6 +147,7 @@
     "LightningIRModule",
     "LightningIROutput",
     "LightningIRTokenizer",
+    "LightningIRTrainer",
     "LinearLRSchedulerWithLinearWarmup",
     "LocalizedContrastiveEstimation",
     "LR_SCHEDULERS",