Fix meta tensor checkpoint materialization

src/evaluator.py

@@ -50,6 +50,34 @@ def _load_ckpt(path: Path, device):
     return torch.load(path, map_location="cpu", weights_only=False)
 
 
+def _construct_model_on_cpu(factory):
+    """Build modules with real CPU tensors even if the runtime default device is meta."""
+    old_device = None
+    can_restore = False
+    if hasattr(torch, "get_default_device") and hasattr(torch, "set_default_device"):
+        try:
+            old_device = torch.get_default_device()
+            torch.set_default_device("cpu")
+            can_restore = True
+        except Exception:
+            old_device = None
+            can_restore = False
+    try:
+        try:
+            with torch.device("cpu"):
+                return factory()
+        except Exception:
+            return factory()
+    finally:
+        # Do not restore a meta default device inside Spaces; that would make
+        # later tensors empty again. Restoring normal devices is safe.
+        if can_restore and old_device is not None and str(old_device) != "meta":
+            try:
+                torch.set_default_device(old_device)
+            except Exception:
+                pass
+
+
 def _move_model_to_device(model, device):
     """Move safely when Transformers creates meta tensors in Spaces."""
     try:
@@ -61,12 +89,43 @@ def _move_model_to_device(model, device):
         return model.to_empty(device=device)
 
 
+def _has_meta_tensors(model) -> bool:
+    return any(p.device.type == "meta" for p in model.parameters()) or any(
+        b is not None and b.device.type == "meta" for b in model.buffers()
+    )
+
+
+def _reset_known_bert_buffers(model, device="cpu"):
+    """Reset non-persistent BERT embedding buffers after to_empty()."""
+    bert = getattr(model, "bert", None)
+    emb = getattr(bert, "embeddings", None)
+    if emb is None:
+        return model
+    device = torch.device(device)
+    for name, buffer in list(emb._buffers.items()):
+        if buffer is None or name not in {"position_ids", "token_type_ids"}:
+            continue
+        shape = tuple(buffer.shape)
+        if not shape:
+            continue
+        if name == "position_ids":
+            values = torch.arange(shape[-1], dtype=torch.long, device=device)
+            emb._buffers[name] = values.view(1, -1).expand(shape).clone()
+        elif name == "token_type_ids":
+            emb._buffers[name] = torch.zeros(shape, dtype=torch.long, device=device)
+    return model
+
+
 def _load_model_state(model, state_dict, strict=True):
     """Load checkpoint weights and report mismatches instead of hiding them."""
+    if _has_meta_tensors(model):
+        logger.warning("Model was built with meta tensors; materializing empty CPU tensors before loading checkpoint.")
+        model.to_empty(device="cpu")
     try:
         result = model.load_state_dict(state_dict, strict=strict, assign=True)
     except TypeError:
         result = model.load_state_dict(state_dict, strict=strict)
+    _reset_known_bert_buffers(model, device="cpu")
     missing = list(getattr(result, "missing_keys", []))
     unexpected = list(getattr(result, "unexpected_keys", []))
     if missing:
@@ -130,12 +189,12 @@ def load_meta_acsa(checkpoint_dir: Path = None, meta_encoder: Optional[MetaEncod
             f"Re-fit the encoder on the same train split used for training."
         )
     if architecture == "gated_aspect_semantic_meta_acsa":
-        model = GatedAspectSemanticMetaFusionACSAModel(
+        model = _construct_model_on_cpu(lambda: GatedAspectSemanticMetaFusionACSAModel(
             bert_name=bert_name,
             meta_in_dim=meta_in_dim,
-        )
+        ))
     else:
-        model = BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim)
+        model = _construct_model_on_cpu(lambda: BertMetaFusionACSAModel(bert_name=bert_name, meta_in_dim=meta_in_dim))
     _load_model_state(model, ckpt["model_state_dict"], strict=False)
     model = _fix_meta_buffers_only(model, device)
     model = model.to(device)
@@ -152,7 +211,7 @@ def load_acsa(checkpoint_dir: Path = None, device=None):
         device = get_device()
     ckpt = _load_ckpt(checkpoint_dir / "best.pt", device)
     bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME)
-    model = BertACSAModel(bert_name=bert_name)
+    model = _construct_model_on_cpu(lambda: BertACSAModel(bert_name=bert_name))
     _load_model_state(model, ckpt["model_state_dict"])
     model = _fix_meta_buffers_only(model, device)
     model = model.to(device)
@@ -169,7 +228,7 @@ def load_bert_overall(checkpoint_dir: Path = None, device=None):
         device = get_device()
     ckpt = _load_ckpt(checkpoint_dir / "best.pt", device)
     bert_name = ckpt.get("config", {}).get("bert_name", cfg.BERT_MODEL_NAME)
-    model = BertOverallModel(bert_name=bert_name)
+    model = _construct_model_on_cpu(lambda: BertOverallModel(bert_name=bert_name))
     _load_model_state(model, ckpt["model_state_dict"])
     model = _fix_meta_buffers_only(model, device)
     model = model.to(device)