Update modeling_transhla2.py

modeling_transhla2.py

@@ -27,6 +27,7 @@ class TransHLA2Config(PretrainedConfig):
         lora_inference_mode=False,
         target_modules=None,
         return_prob=True,  # 是否在 forward 返回概率（softmax），否则返回 logits
+        pad_token_id=1,     # ESM 默认 pad id
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -50,6 +51,7 @@ class TransHLA2Config(PretrainedConfig):
         self.target_modules = target_modules or ['query', 'out_proj', 'value', 'key', 'dense', 'regression']
 
         self.return_prob = return_prob
+        self.pad_token_id = pad_token_id
 
 
 class TransHLA2(PreTrainedModel):
@@ -165,11 +167,24 @@ class TransHLA2(PreTrainedModel):
         x = px + x
         return x
 
+    def _ensure_mapping_input(self, x):
+        # 允许两种输入形式：
+        # 1) 字典: {"input_ids": ..., "attention_mask": ...}
+        # 2) 直接的 input_ids 张量: (B, L)
+        if isinstance(x, torch.Tensor):
+            # 仅用 input_ids；如需自动构造 attention_mask，可解除注释：
+            # pad_id = self.config.pad_token_id
+            # return {"input_ids": x, "attention_mask": (x != pad_id).long()}
+            return {"input_ids": x}
+        elif isinstance(x, dict):
+            return x
+        else:
+            raise TypeError(f"Unsupported input type: {type(x)}; expected Tensor or dict.")
+
     def forward(self, epitope_in, hla_in, return_dict=None):
-        # epitope_in, hla_in: 输入应为 ESM 的输入字典或张量（通常是 input_ids/attention_mask）
-        # 这里假定传入的是 ESM 的标准输入字典，例如:
-        #   epitope_in = {"input_ids": ..., "attention_mask": ...}
-        #   hla_in     = {"input_ids": ..., "attention_mask": ...}
+        # 兼容张量或字典输入
+        epitope_in = self._ensure_mapping_input(epitope_in)
+        hla_in = self._ensure_mapping_input(hla_in)
 
         epitope_outputs = self.epitope_lora(**epitope_in)
         hla_outputs = self.hla_lora(**hla_in)