Fix faulty rstrip in module loading (huggingface#31108)

vasqu · May 29, 2024 · 0bef4a2 · 0bef4a2
1 parent 97a58a5
commit 0bef4a2
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/transformers/dynamic_module_utils.py b/src/transformers/dynamic_module_utils.py
@@ -198,7 +198,10 @@ def get_class_in_module(class_name: str, module_path: Union[str, os.PathLike]) -
  Returns:
  `typing.Type`: The class looked for.
  """
- name = os.path.normpath(module_path).rstrip(".py").replace(os.path.sep, ".")
+ name = os.path.normpath(module_path)
+ if name.endswith(".py"):
+ name = name[:-3]
+ name = name.replace(os.path.sep, ".")
  module_spec = importlib.util.spec_from_file_location(name, location=Path(HF_MODULES_CACHE) / module_path)
  module = sys.modules.get(name)
  if module is None: