Guard BERT embedding inputs during forward

src/inference.py

@@ -201,7 +201,7 @@ class AspectPredictor:
             review_text,
             max_length=self._max_inference_length(),
             truncation=True,
-            padding="max_length",
+            padding=True,
             return_tensors="pt",
         )
         input_ids = self._prepare_input_ids(enc["input_ids"]).to(self.device)
@@ -245,7 +245,7 @@ class AspectPredictor:
                 [r["review_text"] for r in chunk],
                 max_length=self._max_inference_length(),
                 truncation=True,
-                padding="max_length",
+                padding=True,
                 return_tensors="pt",
             )
             input_ids = self._prepare_input_ids(enc["input_ids"]).to(self.device)

src/models.py

@@ -60,6 +60,40 @@ def _load_bert_model(bert_name: str):
     return model
 
 
+def _safe_bert_forward(bert, input_ids, attention_mask, output_attentions=False):
+    embeddings = getattr(bert, "embeddings", None)
+    word_embeddings = getattr(embeddings, "word_embeddings", None)
+    position_embeddings = getattr(embeddings, "position_embeddings", None)
+    token_type_embeddings = getattr(embeddings, "token_type_embeddings", None)
+
+    if word_embeddings is not None:
+        vocab_size = int(word_embeddings.num_embeddings)
+        if vocab_size > 0:
+            input_ids = input_ids.clamp(min=0, max=vocab_size - 1)
+
+    if position_embeddings is not None:
+        max_pos = int(position_embeddings.num_embeddings)
+        if max_pos > 0 and input_ids.size(1) > max_pos:
+            input_ids = input_ids[:, :max_pos]
+            attention_mask = attention_mask[:, :max_pos]
+
+    token_type_ids = torch.zeros_like(input_ids)
+    if token_type_embeddings is not None:
+        type_size = int(token_type_embeddings.num_embeddings)
+        if type_size <= 0:
+            token_type_ids = None
+        elif type_size == 1:
+            token_type_ids = token_type_ids.clamp(max=0)
+
+    return bert(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        token_type_ids=token_type_ids,
+        output_attentions=output_attentions,
+        return_dict=True,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Shared helpers
 # ---------------------------------------------------------------------------
@@ -363,12 +397,7 @@ class BertMetaFusionACSAModel(nn.Module):
         overall_labels: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
     ):
-        bert_out = self.bert(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            return_dict=True,
-        )
+        bert_out = _safe_bert_forward(self.bert, input_ids, attention_mask, output_attentions)
         text_vec = bert_out.last_hidden_state[:, 0, :]   # [CLS]
 
         meta_vec = self.meta_mlp(meta_features)          # (B, meta_hidden)
@@ -491,12 +520,7 @@ class GatedAspectSemanticMetaFusionACSAModel(nn.Module):
         overall_labels: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
     ):
-        bert_out = self.bert(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-            return_dict=True,
-        )
+        bert_out = _safe_bert_forward(self.bert, input_ids, attention_mask, output_attentions)
         text_vec = bert_out.last_hidden_state[:, 0, :]
         B = text_vec.size(0)
         meta_tokens = self.meta_tokenizer(meta_features)
@@ -589,10 +613,7 @@ class BertACSAModel(nn.Module):
     def forward(self, input_ids, attention_mask,
                 labels: Optional[torch.Tensor] = None,
                 output_attentions: bool = False):
-        bert_out = self.bert(
-            input_ids=input_ids, attention_mask=attention_mask,
-            output_attentions=output_attentions, return_dict=True,
-        )
+        bert_out = _safe_bert_forward(self.bert, input_ids, attention_mask, output_attentions)
         text_vec = bert_out.last_hidden_state[:, 0, :]
         logits = self.heads(text_vec)
         loss = _aspect_loss(logits, labels, self.class_weights) if labels is not None else None
@@ -622,8 +643,7 @@ class BertOverallModel(nn.Module):
         self.num_classes = num_classes
 
     def forward(self, input_ids, attention_mask, labels=None, output_attentions=False):
-        out = self.bert(input_ids=input_ids, attention_mask=attention_mask,
-                        output_attentions=output_attentions, return_dict=True)
+        out = _safe_bert_forward(self.bert, input_ids, attention_mask, output_attentions)
         logits = self.classifier(out.last_hidden_state[:, 0, :])
         loss = F.cross_entropy(logits, labels) if labels is not None else None
         return {