Harden inference token ids

src/inference.py

@@ -173,6 +173,21 @@ class AspectPredictor:
         self.model, self.tokenizer, self.meta_encoder, self.device = load_meta_acsa(
             checkpoint_dir=checkpoint_dir, device=device,
         )
+        self._align_tokenizer_and_embeddings()
+
+    def _align_tokenizer_and_embeddings(self) -> None:
+        bert = getattr(self.model, "bert", None)
+        embeddings = getattr(bert, "embeddings", None)
+        word_embeddings = getattr(embeddings, "word_embeddings", None)
+        if bert is None or word_embeddings is None:
+            return
+        try:
+            tokenizer_size = len(self.tokenizer)
+        except Exception:
+            tokenizer_size = 0
+        vocab_size = int(getattr(word_embeddings, "num_embeddings", 0) or 0)
+        if tokenizer_size > vocab_size > 0 and hasattr(bert, "resize_token_embeddings"):
+            bert.resize_token_embeddings(tokenizer_size)
 
     def _prepare_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
         bert = getattr(self.model, "bert", None)
@@ -183,8 +198,7 @@ class AspectPredictor:
         vocab_size = int(word_embeddings.num_embeddings)
         if vocab_size <= 0:
             return input_ids
-        unk_id = int(getattr(self.tokenizer, "unk_token_id", 0) or 0)
-        return torch.where(input_ids >= vocab_size, torch.full_like(input_ids, unk_id), input_ids)
+        return input_ids.clamp(min=0, max=vocab_size - 1)
 
     def _max_inference_length(self) -> int:
         bert = getattr(self.model, "bert", None)

src/models.py

@@ -61,6 +61,12 @@ def _load_bert_model(bert_name: str):
 
 
 def _safe_bert_forward(bert, input_ids, attention_mask, output_attentions=False):
+    input_ids = input_ids.long()
+    if attention_mask is None:
+        attention_mask = torch.ones_like(input_ids)
+    elif attention_mask.size(1) != input_ids.size(1):
+        attention_mask = attention_mask[:, :input_ids.size(1)]
+
     embeddings = getattr(bert, "embeddings", None)
     word_embeddings = getattr(embeddings, "word_embeddings", None)
     position_embeddings = getattr(embeddings, "position_embeddings", None)