implement lazy loading

convokit/__init__.py

@@ -1,41 +1,104 @@
 import warnings
+from typing import Any
 
-try:
-    from .model import *
-    from .util import *
-    from .coordination import *
-    from .politenessStrategies import *
-    from .transformer import *
-    from .convokitPipeline import *
-    from .hyperconvo import *
-    from .speakerConvoDiversity import *
-    from .text_processing import *
-    from .phrasing_motifs import *
-    from .prompt_types import *
-    from .classifier.classifier import *
-    from .ranker import *
-    from .forecaster import *
-    from .fighting_words import *
-    from .paired_prediction import *
-    from .bag_of_words import *
-    from .expected_context_framework import *
-    from .surprise import *
-    from .convokitConfig import *
-    from .redirection import *
-    from .pivotal_framework import *
-    from .utterance_simulator import *
-except ModuleNotFoundError as e:
-    # Don't print ModuleNotFoundError messages as they're handled by individual modules
-    if "not currently installed" not in str(e):
-        print(f"An error occurred: {e}")
-        warnings.warn(
-            "If you are using ConvoKit with Google Colab, incorrect versions of some packages (ex. scipy) may be imported while runtime start. To fix the issue, restart the session and run all codes again. Thank you!"
-        )
-except Exception as e:
-    print(f"An error occurred: {e}")
-    warnings.warn(
-        "If you are using ConvoKit with Google Colab, incorrect versions of some packages (ex. scipy) may be imported while runtime start. To fix the issue, restart the session and run all codes again. Thank you!"
-    )
-
-
-# __path__ = __import__('pkgutil').extend_path(__path__, __name__)
+# Core modules - always imported immediately
+from .model import *
+from .util import *
+from .transformer import *
+from .convokitConfig import *
+from .convokitPipeline import *
+
+# Module mapping for lazy loading
+# Each entry maps module_name -> import_path
+_LAZY_MODULES = {
+    "coordination": ".coordination",
+    "politenessStrategies": ".politenessStrategies",
+    "hyperconvo": ".hyperconvo",
+    "speakerConvoDiversity": ".speakerConvoDiversity",
+    "text_processing": ".text_processing",
+    "phrasing_motifs": ".phrasing_motifs",
+    "prompt_types": ".prompt_types",
+    "classifier": ".classifier",
+    "ranker": ".ranker",
+    "forecaster": ".forecaster",
+    "fighting_words": ".fighting_words",
+    "paired_prediction": ".paired_prediction",
+    "bag_of_words": ".bag_of_words",
+    "expected_context_framework": ".expected_context_framework",
+    "surprise": ".surprise",
+    "redirection": ".redirection",
+    "pivotal_framework": ".pivotal_framework",
+    "utterance_simulator": ".utterance_simulator",
+    "utterance_likelihood": ".utterance_likelihood",
+    "speaker_convo_helpers": ".speaker_convo_helpers",
+    "politeness_collections": ".politeness_collections",
+}
+
+# Cache for loaded modules
+_loaded_modules = {}
+
+
+def _lazy_import(module_name: str) -> Any:
+    """Import a module lazily and cache the result."""
+    if module_name in _loaded_modules:
+        return _loaded_modules[module_name]
+
+    if module_name not in _LAZY_MODULES:
+        raise AttributeError(f"module '{__name__}' has no attribute '{module_name}'")
+
+    import_path = _LAZY_MODULES[module_name]
+
+    try:
+        import importlib
+
+        module = importlib.import_module(import_path, package=__name__)
+        _loaded_modules[module_name] = module
+
+        globals_dict = globals()
+        if hasattr(module, "__all__"):
+            for name in module.__all__:
+                if hasattr(module, name):
+                    globals_dict[name] = getattr(module, name)
+        else:
+            for name in dir(module):
+                if not name.startswith("_"):
+                    globals_dict[name] = getattr(module, name)
+
+        return module
+
+    except Exception as e:
+        # Simply re-raise whatever the module throws
+        # Let each module handle its own error messaging
+        raise
+
+
+def __getattr__(name: str) -> Any:
+    """Handle attribute access for lazy-loaded modules."""
+    # Check if it's a module we can lazy load
+    if name in _LAZY_MODULES:
+        return _lazy_import(name)
+
+    # Check if it's an exported symbol from a lazy module
+    # We need to check each module to see if it exports this symbol
+    for module_name in _LAZY_MODULES:
+        if module_name not in _loaded_modules:
+            # Try to import the module to see if it has the requested attribute
+            try:
+                import importlib
+
+                import_path = _LAZY_MODULES[module_name]
+                module = importlib.import_module(import_path, package=__name__)
+
+                # Check if this module has the requested attribute
+                if hasattr(module, name):
+                    # Import the full module (which will add all symbols to globals)
+                    _lazy_import(module_name)
+                    # Return the requested attribute
+                    return getattr(module, name)
+
+            except Exception:
+                # If module fails to import, just skip it and try next module
+                # The module's own error handling will take care of proper error messages
+                continue
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

convokit/forecaster/TransformerDecoderModel.py

@@ -1,20 +1,10 @@
-try:
-    import unsloth
-    from unsloth import FastLanguageModel, is_bfloat16_supported
-    from unsloth.chat_templates import get_chat_template
-    import torch
-    import torch.nn.functional as F
-    from trl import SFTTrainer, SFTConfig
-    from datasets import Dataset
-
-    UNSLOTH_AVAILABLE = True
-except (ModuleNotFoundError, ImportError) as e:
-    if "Unsloth GPU requirement not met" in str(e):
-        raise ImportError("Unsloth GPU requirement not met") from e
-    else:
-        raise ModuleNotFoundError(
-            "unsloth, torch, trl, or datasets is not currently installed. Run 'pip install convokit[llm]' if you would like to use the TransformerDecoderModel."
-        ) from e
+import unsloth
+from unsloth import FastLanguageModel, is_bfloat16_supported
+from unsloth.chat_templates import get_chat_template
+import torch
+import torch.nn.functional as F
+from trl import SFTTrainer, SFTConfig
+from datasets import Dataset
 
 import json
 import os

convokit/forecaster/TransformerEncoderModel.py

@@ -1,20 +1,13 @@
-try:
-    import torch
-    import torch.nn.functional as F
-    from datasets import Dataset, DatasetDict
-    from transformers import (
-        AutoConfig,
-        AutoModelForSequenceClassification,
-        AutoTokenizer,
-        TrainingArguments,
-        Trainer,
-    )
-
-    TRANSFORMERS_AVAILABLE = True
-except (ModuleNotFoundError, ImportError) as e:
-    raise ModuleNotFoundError(
-        "torch, transformers, or datasets is not currently installed. Run 'pip install convokit[llm]' if you would like to use the TransformerEncoderModel."
-    ) from e
+import torch
+import torch.nn.functional as F
+from datasets import Dataset, DatasetDict
+from transformers import (
+    AutoConfig,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    TrainingArguments,
+    Trainer,
+)
 
 import os
 import pandas as pd

convokit/forecaster/__init__.py

@@ -11,25 +11,36 @@
 # Import Transformer models with proper error handling
 try:
     from .TransformerDecoderModel import *
-except ImportError as e:
+except (ImportError, ModuleNotFoundError) as e:
     if "Unsloth GPU requirement not met" in str(e):
-        print(
+        raise ImportError(
             "Error from Unsloth: NotImplementedError: Unsloth currently only works on NVIDIA GPUs and Intel GPUs."
-        )
-    elif "not currently installed" in str(e):
-        print(
+        ) from e
+    elif (
+        "not currently installed" in str(e)
+        or "torch" in str(e)
+        or "unsloth" in str(e)
+        or "trl" in str(e)
+        or "datasets" in str(e)
+    ):
+        raise ImportError(
             "TransformerDecoderModel requires ML dependencies. Run 'pip install convokit[llm]' to install them."
-        )
+        ) from e
     else:
         raise
 
 try:
     from .TransformerEncoderModel import *
-except ImportError as e:
-    if "not currently installed" in str(e):
-        print(
+except (ImportError, ModuleNotFoundError) as e:
+    if (
+        "not currently installed" in str(e)
+        or "torch" in str(e)
+        or "transformers" in str(e)
+        or "datasets" in str(e)
+    ):
+        raise ImportError(
             "TransformerEncoderModel requires ML dependencies. Run 'pip install convokit[llm]' to install them."
-        )
+        ) from e
     else:
         raise
 

convokit/pivotal_framework/__init__.py

@@ -1,13 +1,18 @@
 try:
     from .pivotal import *
-except ImportError as e:
+except (ImportError, ModuleNotFoundError) as e:
     if "Unsloth GPU requirement not met" in str(e):
-        print(
+        raise ImportError(
             "Error from Unsloth: NotImplementedError: Unsloth currently only works on NVIDIA GPUs and Intel GPUs."
-        )
-    elif "not currently installed" in str(e):
-        print(
+        ) from e
+    elif (
+        "not currently installed" in str(e)
+        or "torch" in str(e)
+        or "unsloth" in str(e)
+        or "transformers" in str(e)
+    ):
+        raise ImportError(
             "Pivotal framework requires ML dependencies. Run 'pip install convokit[llm]' to install them."
-        )
+        ) from e
     else:
         raise

convokit/pivotal_framework/pivotal.py

@@ -8,10 +8,7 @@
 from convokit.forecaster.forecasterModel import ForecasterModel
 from convokit.forecaster.forecaster import Forecaster
 
-try:
-    from convokit.utterance_simulator.utteranceSimulatorModel import UtteranceSimulatorModel
-except NotImplementedError as e:
-    raise ImportError("Unsloth GPU requirement not met") from e
+from convokit.utterance_simulator.utteranceSimulatorModel import UtteranceSimulatorModel
 from convokit.utterance_simulator.utteranceSimulator import UtteranceSimulator
 from .util import ContextTuple, DEFAULT_LABELER
 

convokit/redirection/__init__.py

@@ -1,11 +1,19 @@
-from .redirection import *
+try:
+    from .redirection import *
+except (ImportError, ModuleNotFoundError) as e:
+    if "torch" in str(e) or "not currently installed" in str(e):
+        raise ImportError(
+            "Redirection module requires ML dependencies. Run 'pip install convokit[llm]' to install them."
+        ) from e
+    else:
+        raise
 
 try:
     from .likelihoodModel import *
-except ImportError as e:
+except (ImportError, ModuleNotFoundError) as e:
     if "not currently installed" in str(e):
-        print(
+        raise ImportError(
             "LikelihoodModel requires ML dependencies. Run 'pip install convokit[llm]' to install them."
-        )
+        ) from e
     else:
         raise

convokit/redirection/config.py

@@ -1,13 +1,6 @@
-try:
-    from peft import LoraConfig
-    from transformers import BitsAndBytesConfig
-    import torch
-
-    REDIRECTION_ML_AVAILABLE = True
-except (ModuleNotFoundError, ImportError) as e:
-    raise ModuleNotFoundError(
-        "peft, transformers, or torch is not currently installed. Run 'pip install convokit[llm]' if you would like to use the redirection module."
-    ) from e
+from peft import LoraConfig
+from transformers import BitsAndBytesConfig
+import torch
 
 DEFAULT_BNB_CONFIG = BitsAndBytesConfig(
     load_in_4bit=True,

convokit/redirection/gemmaLikelihoodModel.py

@@ -1,20 +1,13 @@
-try:
-    import torch
-    from peft import LoraConfig, get_peft_model, AutoPeftModelForCausalLM, PeftModel
-    from transformers import (
-        AutoTokenizer,
-        AutoModelForCausalLM,
-        BitsAndBytesConfig,
-        DataCollatorForLanguageModeling,
-        TrainingArguments,
-    )
-    from trl import SFTTrainer
-
-    GEMMA_ML_AVAILABLE = True
-except (ModuleNotFoundError, ImportError) as e:
-    raise ModuleNotFoundError(
-        "torch, peft, transformers, or trl is not currently installed. Run 'pip install convokit[llm]' if you would like to use the GemmaLikelihoodModel."
-    ) from e
+import torch
+from peft import LoraConfig, get_peft_model, AutoPeftModelForCausalLM, PeftModel
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
+    DataCollatorForLanguageModeling,
+    TrainingArguments,
+)
+from trl import SFTTrainer
 
 from .likelihoodModel import LikelihoodModel
 from .config import DEFAULT_TRAIN_CONFIG, DEFAULT_BNB_CONFIG, DEFAULT_LORA_CONFIG

convokit/redirection/preprocessing.py

@@ -1,11 +1,4 @@
-try:
-    from datasets import Dataset
-
-    DATASETS_AVAILABLE = True
-except (ModuleNotFoundError, ImportError) as e:
-    raise ModuleNotFoundError(
-        "datasets is not currently installed. Run 'pip install convokit[llm]' if you would like to use the redirection preprocessing functionality."
-    ) from e
+from datasets import Dataset
 
 
 def default_speaker_prefixes(roles):