init commit

training/__init__.py

@@ -0,0 +1,40 @@
+"""
+Training modules for BGE models
+"""
+
+def get_trainer(trainer_type: str):
+    """Lazy loading of trainers to avoid importing all at once"""
+    if trainer_type == 'base':
+        from .trainer import BaseTrainer
+        return BaseTrainer
+    elif trainer_type == 'm3':
+        from .m3_trainer import BGEM3Trainer
+        return BGEM3Trainer
+    elif trainer_type == 'reranker':
+        from .reranker_trainer import BGERerankerTrainer
+        return BGERerankerTrainer
+    elif trainer_type == 'joint':
+        from .joint_trainer import JointTrainer
+        return JointTrainer
+    else:
+        raise ValueError(f"Unknown trainer type: {trainer_type}")
+
+# For backward compatibility, import the commonly used ones
+from .trainer import BaseTrainer
+from .m3_trainer import BGEM3Trainer
+
+# Only import reranker and joint trainers when explicitly requested
+def __getattr__(name):
+    if name == 'BGERerankerTrainer':
+        from .reranker_trainer import BGERerankerTrainer
+        return BGERerankerTrainer
+    elif name == 'JointTrainer':
+        from .joint_trainer import JointTrainer
+        return JointTrainer
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+__all__ = [
+    'BaseTrainer',
+    'BGEM3Trainer',
+    'get_trainer'
+]