Force real BERT tensor loading in Spaces

src/models.py

@@ -1,4 +1,4 @@
-"""Models for ACSA with metadata fusion.
+"""Models for ACSA with metadata fusion.
 
 Three architectures:
   - BertMetaFusionACSAModel:  BERT + Meta Cross-Attention + per-aspect heads  (Proposed)
@@ -47,6 +47,14 @@ logger = logging.getLogger(__name__)
 META_NUM_META_TOKENS = 4
 
 
+def _load_bert_model(bert_name: str):
+    """Load BERT with real tensors in Spaces instead of lazy meta tensors."""
+    try:
+        return AutoModel.from_pretrained(bert_name, low_cpu_mem_usage=False)
+    except TypeError:
+        return AutoModel.from_pretrained(bert_name)
+
+
 # ---------------------------------------------------------------------------
 # Shared helpers
 # ---------------------------------------------------------------------------
@@ -201,6 +209,17 @@ class GatedAspectSemanticCrossAttention(nn.Module):
         self.out_proj = nn.Linear(hidden, hidden)
         self.gate = nn.Linear(hidden * 2, hidden)
         nn.init.constant_(self.gate.bias, -0.3)
+        prior_rows = []
+        prior_cfg = getattr(cfg, "CROSS_ATTN_SOURCE_PRIOR", {})
+        for aspect in cfg.ASPECTS[:num_aspects]:
+            row = list(prior_cfg.get(aspect, [1.0 / 3.0] * 3))
+            if len(row) < 3:
+                row = row + [0.0] * (3 - len(row))
+            prior_rows.append(row[:3])
+        prior = torch.tensor(prior_rows, dtype=torch.float)
+        prior = prior - prior.mean(dim=-1, keepdim=True)
+        self.register_buffer("source_prior_bias", prior, persistent=False)
+        self.source_prior_strength = float(getattr(cfg, "CROSS_ATTN_SOURCE_PRIOR_STRENGTH", 0.0))
         self.dropout = nn.Dropout(dropout)
         self.norm = nn.LayerNorm(hidden)
 
@@ -211,7 +230,13 @@ class GatedAspectSemanticCrossAttention(nn.Module):
         q = q.transpose(1, 2)
         k = self.k_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
         v = self.v_proj(meta_tokens).view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
-        attn = F.softmax(torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim), dim=-1)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if self.source_prior_strength > 0 and T == self.source_prior_bias.shape[-1]:
+            prior = self.source_prior_bias[:self.num_aspects, :T].view(
+                1, 1, self.num_aspects, T
+            )
+            scores = scores + self.source_prior_strength * prior
+        attn = F.softmax(scores, dim=-1)
         ctx = torch.matmul(self.dropout(attn), v)
         ctx = ctx.transpose(1, 2).contiguous().view(B, self.num_aspects, H)
         ctx = self.out_proj(ctx)
@@ -291,7 +316,7 @@ class BertMetaFusionACSAModel(nn.Module):
         class_weights: Optional[torch.Tensor] = None,
     ):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(bert_name)
+        self.bert = _load_bert_model(bert_name)
         hidden = self.bert.config.hidden_size
 
         self.num_aspects = num_aspects
@@ -391,7 +416,7 @@ class GatedAspectSemanticMetaFusionACSAModel(nn.Module):
         expected_dim = cfg.META_TFIDF_DIM + cfg.META_NUM_DIM
         if meta_in_dim != expected_dim:
             raise ValueError(f"Semantic meta fusion expects meta_in_dim={expected_dim}, got {meta_in_dim}")
-        self.bert = AutoModel.from_pretrained(bert_name)
+        self.bert = _load_bert_model(bert_name)
         hidden = self.bert.config.hidden_size
 
         self.num_aspects = num_aspects
@@ -546,7 +571,7 @@ class BertACSAModel(nn.Module):
         class_weights: Optional[torch.Tensor] = None,
     ):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(bert_name)
+        self.bert = _load_bert_model(bert_name)
         hidden = self.bert.config.hidden_size
         self.num_aspects = num_aspects
         self.num_classes = num_classes
@@ -581,7 +606,7 @@ class BertOverallModel(nn.Module):
                  num_classes: int = cfg.OVERALL_NUM_CLASSES,
                  dropout: float = 0.3):
         super().__init__()
-        self.bert = AutoModel.from_pretrained(bert_name)
+        self.bert = _load_bert_model(bert_name)
         hidden = self.bert.config.hidden_size
         self.classifier = nn.Sequential(
             nn.Dropout(dropout),