Upload modeling_qwen3.py with huggingface_hub

modeling_qwen3.py

@@ -578,7 +578,7 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin):
         # [ADD] Custom Logit Masking Logic (Inference Only)
         # =================================================================
         # 仅在非训练模式 (self.training == False) 且 input_ids 存在时执行
-        if not self.training and input_ids is not None:
+        if not self.training and input_ids is not None and labels is None:
             # 1. 判断 Mask 触发条件
             # input_ids shape: (batch_size, seq_len)
             
@@ -591,29 +591,29 @@ class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin):
             if seq_len == 1:
                 # 检查 input_ids 是否在 [172207, 180398] 区间内
                 in_safe_range = ((input_ids >= 172207) & (input_ids <= 180398)).any(dim=-1)
-                # 如果不在安全区间，则满足条件 B
                 cond_b = ~in_safe_range
             else:
-                # 长度不为1，条件 B 必定不满足
                 cond_b = torch.zeros_like(has_trigger_token, dtype=torch.bool)
             
-            # 综合条件: 满足 A 或 满足 B
+            # 综合条件
             rows_to_mask = has_trigger_token | cond_b
 
             # 2. 执行 Mask 操作
             if rows_to_mask.any():
-                # 离散 Token 列表
+                # 修改点 2: 获取行索引，解决广播报错的关键
+                row_idxs = torch.nonzero(rows_to_mask, as_tuple=True)[0]
+                
                 target_discrete_tokens = [151691, 151692, 151693, 151695, 151696, 151697, 151698]
                 mask_indices = torch.tensor(target_discrete_tokens, device=logits.device)
                 neg = torch.finfo(logits.dtype).min
+                
                 # (1) Mask 离散 Token
-                # logits[rows_to_mask] 选取需要 mask 的 batch 行
-                # [:, mask_indices] 选取特定的 token ID 列
-                logits[rows_to_mask, :, mask_indices] = neg
+                # 修改点 3: 使用 row_idxs[:, None] 将形状变为 [N, 1]，使其能与 mask_indices [7] 广播
+                if mask_indices.numel() > 0:
+                    logits[row_idxs[:, None], :, mask_indices] = neg
                 
                 # (2) Mask 连续区间 [172206, 180398]
-                # 注意：Python 切片右边界是开区间，所以要写到 180399
-                logits[rows_to_mask, :, 151727:180399] = neg
+                logits[row_idxs, :, 151727:180399] = neg
         # =================================================================
 
         loss = None