Update modeling_matformer.py

modeling_matformer.py

@@ -1,156 +1,47 @@
 # modeling_matformer.py
-
-import importlib
 import os
-import site
-import subprocess
 import sys
-import sysconfig
-
-
-MATFORMER_VERSION = "alberturin-v1"
-MATFORMER_REPOSITORY = "https://github.com/mrinaldi97/matformer.git"
-MATFORMER_INSTALL_URL = (
-    f"git+{MATFORMER_REPOSITORY}@{MATFORMER_VERSION}"
-)
-
-
-def _add_matformer_root():
-    """
-    Optional development override.
-
-    If MATFORMER_ROOT is defined, prepend it to sys.path so that a local
-    checkout of Matformer can be used instead of the installed package.
-    """
-    matformer_root = os.getenv("MATFORMER_ROOT")
-
-    if not matformer_root:
-        return
-
-    matformer_root = os.path.abspath(
-        os.path.expanduser(matformer_root)
-    )
 
+matformer_root = os.getenv("MATFORMER_ROOT")
+if matformer_root:
+    matformer_root = os.path.abspath(os.path.expanduser(matformer_root))
     if matformer_root not in sys.path:
         sys.path.insert(0, matformer_root)
 
-
-def _clear_matformer_import_cache():
-    """
-    Remove partially imported Matformer modules from sys.modules.
-
-    This is necessary because the first import can fail before Matformer
-    is installed, leaving stale module entries in the running interpreter.
-    """
-    for module_name in list(sys.modules):
-        if (
-            module_name == "matformer"
-            or module_name.startswith("matformer.")
-        ):
-            del sys.modules[module_name]
-
-    importlib.invalidate_caches()
-
-
-def _refresh_site_packages():
-    """
-    Refresh site-packages after pip installs a package while the current
-    Python interpreter is already running.
-    """
-    paths = sysconfig.get_paths()
-
-    for key in ("purelib", "platlib"):
-        path = paths.get(key)
-
-        if path and os.path.isdir(path):
-            site.addsitedir(path)
-
-    importlib.invalidate_caches()
-
-
-def _install_matformer():
-    """
-    Install the exact Matformer release used by AlBERTurin.
-    """
-    print(
-        f"Matformer is not available. "
-        f"Installing Matformer {MATFORMER_VERSION}..."
-    )
-
-    subprocess.check_call(
-        [
-            sys.executable,
-            "-m",
-            "pip",
-            "install",
-            MATFORMER_INSTALL_URL,
-        ]
-    )
-
-    _refresh_site_packages()
-    _clear_matformer_import_cache()
-
-
-def _load_matformer():
-    """
-    Import and register the Hugging Face integration from Matformer.
-    """
+try:
     from matformer.huggingface_integration import (
-        MatformerConfig,
         MatformerForCausalLM,
         MatformerForMaskedLM,
         MatformerForSequenceClassification,
         MatformerModel,
-        register_matformer,
-    )
-
-    register_matformer()
-
-    return (
         MatformerConfig,
-        MatformerModel,
-        MatformerForMaskedLM,
-        MatformerForCausalLM,
-        MatformerForSequenceClassification,
+        register_matformer
     )
-
-
-# -------------------------------------------------------------------------
-# Initialization
-# -------------------------------------------------------------------------
-
-_add_matformer_root()
-
-try:
-    (
-        MatformerConfig,
-        MatformerModel,
-        MatformerForMaskedLM,
-        MatformerForCausalLM,
-        MatformerForSequenceClassification,
-    ) = _load_matformer()
-
-except ModuleNotFoundError as first_error:
-
+    register_matformer()
+except ImportError as e:
+    import subprocess
+    import tempfile
+    
+    print("Installing Matformer from GitHub...")
     try:
-        _clear_matformer_import_cache()
-        _install_matformer()
-
-        (
-            MatformerConfig,
-            MatformerModel,
-            MatformerForMaskedLM,
+        subprocess.check_call([
+            sys.executable, "-m", "pip", "install", 
+            "git+https://github.com/mrinaldi97/matformer.git@alberturin-v1"
+        ])
+        
+        from matformer.huggingface_integration import (
             MatformerForCausalLM,
+            MatformerForMaskedLM,
             MatformerForSequenceClassification,
-        ) = _load_matformer()
-
+            MatformerModel,
+            MatformerConfig,
+            register_matformer
+        )
+        register_matformer()
+        
     except Exception as install_error:
         raise ImportError(
-            "\n"
-            "Failed to automatically install or import the Matformer "
-            "version required by AlBERTurin.\n\n"
-            "You can install it manually with:\n\n"
-            f"  pip install {MATFORMER_INSTALL_URL}\n\n"
-            "Alternatively, set MATFORMER_ROOT to a local Matformer "
-            "checkout.\n"
-        ) from install_error
+            "Failed to install Matformer. Install manually:\n"
+            "  pip install git+https://github.com/mrinaldi97/matformer.git\n"
+            "Or set MATFORMER_ROOT environment variable"
+        ) from install_error