Upload folder using huggingface_hub

modeling.py

@@ -3,21 +3,11 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import math
-
-import sys
-if "torch_xla" not in sys.modules:
-    sys.modules["torch_xla"] = type(sys)("torch_xla")
-    # Recreate the submodule structure
-    sys.modules["torch_xla"].core = type(sys)("core")
-    sys.modules["torch_xla"].distributed = type(sys)("distributed")
-    sys.modules["torch_xla"].distributed.spmd = type(sys)("spmd")
-    # Set required attributes to None (matching your try-except)
-    sys.modules["torch_xla"].core.xla_model = None
-    sys.modules["torch_xla"].distributed.spmd.xla_sharding = None
+import importlib
 
 try:
-    import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.spmd.xla_sharding as xs
+    xm = importlib.import_module('torch_xla.core.xla_model')
+    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
 except ImportError:
     xm = None
     xs = None

psi.py

@@ -4,6 +4,7 @@ PSI Model Definition
 
 
 import math
+import importlib
 from typing import Tuple, Union, List, Optional, Callable, Dict
 from transformers import PreTrainedModel
 import torch
@@ -18,8 +19,8 @@ from .modeling import (
 )
 
 try:
-    import torch_xla.core.xla_model as xm
-    import torch_xla.distributed.spmd.xla_sharding as xs
+    xm = importlib.import_module('torch_xla.core.xla_model')
+    xs = importlib.import_module('torch_xla.distributed.spmd.xla_sharding')
 except ImportError:
     xm = None
     xs = None