Fix PyTorch 2.6 checkpoint loading for Diac Lightning models

Register yacs CfgNode safe globals and default weights_only=False when
loading from_pretrained checkpoints (PyTorch 2.6 weights_only=True default).

Co-authored-by: AHAD <ahad-m@users.noreply.github.com>

backend/app/checkpoint_compat.py

@@ -0,0 +1,76 @@
+"""
+PyTorch 2.6+ compatibility for Lightning checkpoints from the Diac project.
+
+Lightning checkpoints embed yacs.config.CfgNode in hyperparameters. PyTorch 2.6
+defaults torch.load(..., weights_only=True), which rejects those objects unless
+allowlisted or weights_only=False is used (trusted checkpoints only).
+"""
+
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+_PATCHED = False
+
+
+def prepare_checkpoint_loading() -> None:
+    """Configure torch/Lightning so Diac HF checkpoints load on PyTorch >= 2.6."""
+    global _PATCHED
+    if _PATCHED:
+        return
+
+    _register_safe_globals()
+    _patch_lightning_load_from_checkpoint()
+    _PATCHED = True
+    logger.debug("Checkpoint loading compatibility applied (PyTorch 2.6+ / yacs)")
+
+
+def _register_safe_globals() -> None:
+    try:
+        import torch
+    except ImportError:
+        return
+
+    if not hasattr(torch.serialization, "add_safe_globals"):
+        return
+
+    entries: list[type] = []
+    try:
+        from yacs.config import CfgNode
+
+        entries.append(CfgNode)
+    except ImportError:
+        pass
+
+    try:
+        from types import SimpleNamespace
+
+        entries.append(SimpleNamespace)
+    except ImportError:
+        pass
+
+    if entries:
+        torch.serialization.add_safe_globals(entries)
+
+
+def _patch_lightning_load_from_checkpoint() -> None:
+    """
+    Diac's from_pretrained() does not pass weights_only=False.
+    Trusted NAACL/Hugging Face checkpoints require full pickle load.
+    """
+    try:
+        from lightning.pytorch.core import LightningModule
+    except ImportError:
+        return
+
+    original = LightningModule.__dict__["load_from_checkpoint"]
+
+    @classmethod
+    def load_from_checkpoint_compat(cls, *args, **kwargs):
+        if kwargs.get("weights_only") is None:
+            kwargs["weights_only"] = False
+        return original(cls, *args, **kwargs)
+
+    LightningModule.load_from_checkpoint = load_from_checkpoint_compat

backend/app/services/model_service.py

@@ -5,6 +5,8 @@ from __future__ import annotations
 import logging
 from typing import Optional
 
+from backend.app.checkpoint_compat import prepare_checkpoint_loading
+
 logger = logging.getLogger(__name__)
 
 try:
@@ -40,6 +42,7 @@ class ModelService:
             )
 
         logger.info("Loading model %s with constants at %s", model_name, constants_path)
+        prepare_checkpoint_loading()
         self._model = DiacritizationModule.from_pretrained(
             model_name,
             tokenizer_constants_path=constants_path,