ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(transformer): 调整注意力掩码处理逻辑

This commit is contained in:
ViperEkura 2025-09-29 11:31:42 +08:00
parent 816bc78894
commit 8206c7855e
1 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
khaosz/core/transformer.py
View File

@ -91,12 +91,12 @@ def process_attention_mask(
Tensor: The attention mask tensor. Tensor: The attention mask tensor.
""" """
if start_pos != 0 and seq_mask is None:
# for single prompt chat
seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
if seq_mask is None: if seq_mask is None:
return None if start_pos != 0:
# for single prompt chat
seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
else:
return None
if seq_mask.dim() > 2: if seq_mask.dim() > 2:
# shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos) # shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)