featfix(transformers): lotts

Author: dv-ar

emm/indexing/__init__.py

@@ -45,25 +45,18 @@
     "PandasSortedNeighbourhoodIndexer",
 ]
 
-# Optional Spark support
+# Feature detection for sentence transformers
+HAS_SENTENCE_TRANSFORMER = False
 try:
-    import pyspark
-    from emm.indexing.spark_cos_sim_matcher import SparkCosSimIndexer
-    from emm.indexing.spark_candidate_selection import SparkCandidateSelectionEstimator
-    from emm.indexing.spark_sni import SparkSortedNeighbourhoodIndexer
-    __all__.extend([
-        "SparkCosSimIndexer",
-        "SparkCandidateSelectionEstimator", 
-        "SparkSortedNeighbourhoodIndexer"
-    ])
+    import sentence_transformers
+    HAS_SENTENCE_TRANSFORMER = True
 except ImportError:
     pass
 
-# Optional Sentence Transformer support 
-try:
-    from emm.indexing.pandas_sentence_transformer import PandasSentenceTransformerIndexer
-    __all__.extend([
-        "PandasSentenceTransformerIndexer"
-    ])
-except ImportError:
-    pass  # Transformer features unavailable
+# Only import if dependencies are available
+if HAS_SENTENCE_TRANSFORMER:
+    try:
+        from emm.indexing.pandas_sentence_transformer import PandasSentenceTransformerIndexer
+        __all__.append("PandasSentenceTransformerIndexer")
+    except ImportError:
+        HAS_SENTENCE_TRANSFORMER = False

emm/models/__init__.py

@@ -1,25 +1,40 @@
+"""Optional model extensions for entity matching.
+
+This module provides optional model-based features that can be enabled by installing
+additional dependencies:
+
+1. Sentence Transformers: Install with `pip install emm[transformers]`
+2. Model Tuning: Install with `pip install emm[tuning]`
+"""
+
 from __future__ import annotations
 
 __all__ = []
 
 # Core functionality is empty - models are all optional features
 # This makes it explicit that this module only provides optional extensions
 
-# Primary optional feature: Sentence Transformers
-try:
-    import sentence_transformers
-    from emm.models.sentence_transformer.base import BaseSentenceTransformerComponent
-    __all__.append("BaseSentenceTransformerComponent")
+def _import_sentence_transformers():
+    """Helper to import sentence transformer components"""
+    try:
+        import sentence_transformers
+        from emm.models.sentence_transformer.base import BaseSentenceTransformerComponent
+        __all__.append("BaseSentenceTransformerComponent")
+        return True
+    except ImportError:
+        return False
 
-    # Secondary optional feature: Model Tuning
-    # Only available if sentence transformers is installed
+def _import_tuning():
+    """Helper to import tuning components"""
     try:
         import lightning
         import wandb
         from emm.models.sentence_transformer.tuning import SentenceTransformerTuner, TuningConfig
         __all__.extend(["SentenceTransformerTuner", "TuningConfig"])
+        return True
     except ImportError:
-        pass  # Tuning features unavailable
+        return False
 
-except ImportError:
-    pass  # Transformer features unavailable 
+# Try to import optional features
+HAS_TRANSFORMERS = _import_sentence_transformers()
+HAS_TUNING = HAS_TRANSFORMERS and _import_tuning()  # Tuning requires transformers 

emm/models/sentence_transformer/__init__.py

@@ -1,7 +1,7 @@
 """Sentence transformer functionality for entity matching.
 
-Note: This module assumes sentence-transformers is installed.
-Import checks are handled at the package level in emm.models.__init__.py
+This module requires the sentence-transformers package.
+Install with: pip install emm[transformers]
 """
 
 from __future__ import annotations
@@ -10,7 +10,7 @@
 
 __all__ = ["BaseSentenceTransformerComponent"]
 
-# Move optional imports inside try block
+# Import tuning if available
 try:
     from emm.models.sentence_transformer.tuning import TuningConfig, SentenceTransformerTuner
     __all__ += ["TuningConfig", "SentenceTransformerTuner"]

emm/models/sentence_transformer/base.py

@@ -1,10 +1,28 @@
+"""Base sentence transformer component for entity matching.
+
+This module requires sentence-transformers and torch.
+Install with: pip install emm[transformers]
+"""
+
 from __future__ import annotations
 
 from typing import List, Dict, Optional, Any, Tuple
-import torch
-import numpy as np
-from sentence_transformers import SentenceTransformer, util, SimilarityFunction
+import logging
+
+# Check required dependencies at import time
+try:
+    import torch
+    import sentence_transformers
+    from sentence_transformers import SentenceTransformer, util, SimilarityFunction
+except ImportError as e:
+    raise ImportError(
+        "sentence-transformers and torch are required for this module. "
+        "Install with: pip install emm[transformers]"
+    ) from e
+
+# Rest of imports
 from functools import lru_cache
+import numpy as np
 
 # Single numpy typing import with fallback
 try:
@@ -13,6 +31,8 @@
     from typing import Any
     NDArray = Any
 
+logger = logging.getLogger(__name__)
+
 class BaseSentenceTransformerComponent:
     """Base component for sentence transformer functionality in EMM.
     

emm/models/sentence_transformer/examples/company_name_tuning.py

@@ -1,7 +1,22 @@
+"""Example of company name tuning with sentence transformers.
+
+This example requires the full tuning dependencies.
+Install with: pip install emm[tuning]
+"""
+
 from pathlib import Path
 from typing import List
-import torch
-from torch.utils.data import Dataset, DataLoader
+
+# Check required dependencies at import time
+try:
+    import torch
+    from torch.utils.data import Dataset, DataLoader
+except ImportError as e:
+    raise ImportError(
+        "This example requires sentence-transformers and tuning dependencies. "
+        "Install with: pip install emm[tuning]"
+    ) from e
+
 from emm.models.sentence_transformer.tuning import SentenceTransformerTuner, TuningConfig
 
 class CompanyNameDataset(Dataset):

emm/models/sentence_transformer/examples/sentence_transformer_indexing.py

@@ -1,7 +1,18 @@
-"""Example script demonstrating entity matching using sentence transformers for indexing."""
+"""Example of entity matching using sentence transformers for indexing.
 
-import pandas as pd
-from emm import PandasEntityMatching
+This example requires sentence-transformers.
+Install with: pip install emm[transformers]
+"""
+
+# Check required dependencies at import time
+try:
+    import pandas as pd
+    from emm import PandasEntityMatching
+except ImportError as e:
+    raise ImportError(
+        "This example requires sentence-transformers. "
+        "Install with: pip install emm[transformers]"
+    ) from e
 
 def main():
     # Create sample ground truth data

emm/models/sentence_transformer/tuning/__init__.py

@@ -1,7 +1,11 @@
 """Model tuning functionality for sentence transformers.
 
-Note: This module assumes lightning and wandb are installed.
-Import checks are handled at the package level in emm.models.__init__.py
+This module requires additional dependencies:
+- sentence-transformers
+- lightning 
+- wandb
+
+Install with: pip install emm[tuning]
 """
 
 from __future__ import annotations

emm/models/sentence_transformer/tuning/config.py

@@ -1,7 +1,25 @@
+"""Configuration for sentence transformer tuning.
+
+This module requires lightning and wandb.
+Install with: pip install emm[tuning]
+"""
+
+from __future__ import annotations
+
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
+# Check required dependencies at import time
+try:
+    import lightning
+    import wandb
+except ImportError as e:
+    raise ImportError(
+        "lightning and wandb are required for tuning functionality. "
+        "Install with: pip install emm[tuning]"
+    ) from e
+
 @dataclass
 class TuningConfig:
     """Configuration for fine-tuning sentence transformers

emm/models/sentence_transformer/tuning/tuner.py

@@ -1,20 +1,45 @@
+"""Tuning functionality for sentence transformers.
+
+This module requires additional dependencies:
+- sentence-transformers
+- lightning
+- wandb
+
+Install with: pip install emm[tuning]
+"""
+
 from __future__ import annotations
 
 from typing import Optional, Dict, Any, List
 import logging
-import lightning as L
-from sentence_transformers import SentenceTransformer, losses
-import torch
-import wandb
 from pathlib import Path
-from torch.utils.data import DataLoader
-from torch.cuda.amp import GradScaler
 import numpy as np
 
+# Defer imports until actually needed
+HAS_TUNING_DEPS = False
+try:
+    import torch
+    import lightning as L
+    import wandb
+    from sentence_transformers import SentenceTransformer, losses
+    from torch.utils.data import DataLoader
+    from torch.cuda.amp import GradScaler
+    HAS_TUNING_DEPS = True
+except ImportError:
+    pass
+
 from emm.models.sentence_transformer.tuning.config import TuningConfig
 
 logger = logging.getLogger(__name__)
 
+def check_tuning_dependencies():
+    """Check if tuning dependencies are available"""
+    if not HAS_TUNING_DEPS:
+        raise ImportError(
+            "sentence-transformers, torch, lightning, and wandb are required for tuning. "
+            "Install with: pip install emm[tuning]"
+        )
+
 class SentenceTransformerTuner:
     """Fine-tuning for sentence transformers specialized for company name matching"""
 
@@ -24,6 +49,7 @@ def __init__(self, config: TuningConfig):
         Args:
             config: Tuning configuration object
         """
+        check_tuning_dependencies()
         self.config = config
 
         # Setup Lightning Fabric for distributed training

emm/pipeline/__init__.py

@@ -17,12 +17,24 @@
 # 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.
 
-from emm.helper import spark_installed
+"""Pipeline implementations for entity matching.
+
+Core pipeline:
+- PandasEntityMatching: Pipeline using pandas DataFrames
+
+Optional pipelines:
+- SparkEntityMatching (requires pyspark)
+"""
+
+from __future__ import annotations
+
 from emm.pipeline.pandas_entity_matching import PandasEntityMatching
 
 __all__ = ["PandasEntityMatching"]
 
-if spark_installed:
+# Optional Spark support
+try:
     from emm.pipeline.spark_entity_matching import SparkEntityMatching
-
-    __all__ += ["SparkEntityMatching"]
+    __all__.append("SparkEntityMatching")
+except ImportError:
+    pass