Fix modeling_ablang2paired.py to work with integrated files instead of ablang2 package

modeling_ablang2paired.py

@@ -3,34 +3,39 @@ import os
 from torch import nn
 from transformers import PreTrainedModel
 
-# Import logic to handle the nested directory structure
-import importlib.util
-import os
-
-def load_ablang_module():
-    """Load the AbLang module from the nested directory structure."""
-    # Try to find the ablang.py file in the nested directory
-    current_dir = os.path.dirname(os.path.abspath(__file__))
-    ablang_path = os.path.join(current_dir, "ablang2", "models", "ablang2", "ablang.py")
-    
-    if os.path.exists(ablang_path):
-        spec = importlib.util.spec_from_file_location("ablang", ablang_path)
-        ablang_module = importlib.util.module_from_spec(spec)
-        spec.loader.exec_module(ablang_module)
-        return ablang_module.AbLang
-    else:
-        # Fallback to absolute import
-        from ablang2.models.ablang2.ablang import AbLang
-        return AbLang
-
-# Load the AbLang class
-AbLang2 = load_ablang_module()
-
 # Import configuration
 try:
     from .configuration_ablang2paired import AbLang2PairedConfig
 except ImportError:
-    from ablang2.configuration_ablang2paired import AbLang2PairedConfig
+    from configuration_ablang2paired import AbLang2PairedConfig
+
+# Import the AbLang model from the integrated files
+try:
+    from ablang2.models.ablang2.ablang import AbLang
+except ImportError:
+    # Fallback: try to find the ablang.py file in the current directory structure
+    import importlib.util
+    import os
+    
+    def load_ablang_module():
+        """Load the AbLang module from the nested directory structure."""
+        # Try to find the ablang.py file in the nested directory
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        ablang_path = os.path.join(current_dir, "ablang2", "models", "ablang2", "ablang.py")
+        
+        if os.path.exists(ablang_path):
+            spec = importlib.util.spec_from_file_location("ablang", ablang_path)
+            ablang_module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(ablang_module)
+            return ablang_module.AbLang
+        else:
+            # If not found, raise an error with helpful message
+            raise ImportError(
+                "Could not find AbLang module. Please ensure the ablang2 directory structure is present "
+                "or install the ablang2 package."
+            )
+    
+    AbLang = load_ablang_module()
 
 class AbLang2PairedHFModel(PreTrainedModel):
     config_class = AbLang2PairedConfig
@@ -38,7 +43,7 @@ class AbLang2PairedHFModel(PreTrainedModel):
 
     def __init__(self, config: AbLang2PairedConfig):
         super().__init__(config)
-        self.model = AbLang2(
+        self.model = AbLang(
             vocab_size=config.vocab_size,
             hidden_embed_size=config.hidden_embed_size,
             n_attn_heads=config.n_attn_heads,