429 lines
15 KiB
Python
429 lines
15 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from modelscope import AutoModel
|
||
|
||
|
||
# ---------------------------- 注意力池化模块----------------------------
|
||
class AttentionPooling(nn.Module):
|
||
def __init__(self, hidden_size):
|
||
super().__init__()
|
||
self.attn = nn.Linear(hidden_size, 1)
|
||
# 三个可学习偏置：文本、拼音、个性化
|
||
self.bias = nn.Parameter(torch.zeros(3)) # [text_bias, pinyin_bias, user_bias]
|
||
|
||
def forward(self, x, mask=None, token_type_ids=None):
|
||
scores = self.attn(x).squeeze(-1) # [batch, seq_len]
|
||
if token_type_ids is not None:
|
||
# 根据 token_type_ids 添加对应偏置
|
||
# bias 形状 [3]，通过索引扩展为 [batch, seq_len]
|
||
bias_per_token = self.bias[token_type_ids] # [batch, seq_len]
|
||
scores = scores + bias_per_token
|
||
if mask is not None:
|
||
scores = scores.masked_fill(mask == 0, -1e9)
|
||
weights = torch.softmax(scores, dim=-1)
|
||
pooled = torch.sum(weights.unsqueeze(-1) * x, dim=1)
|
||
return pooled
|
||
|
||
|
||
# ---------------------------- 拼音LSTM编码器 ----------------------------
|
||
class PinyinLSTMEncoder(nn.Module):
|
||
"""
|
||
拼音序列编码器，返回每位置的拼音编码。
|
||
使用双向LSTM，每个位置都能看到前后文信息。
|
||
"""
|
||
|
||
def __init__(self, input_dim, hidden_dim=None, num_layers=2, dropout=0.2):
|
||
super().__init__()
|
||
self.input_dim = input_dim
|
||
self.hidden_dim = hidden_dim if hidden_dim is not None else input_dim // 2
|
||
self.num_layers = num_layers
|
||
self.dropout = dropout
|
||
|
||
self.lstm = nn.LSTM(
|
||
input_size=input_dim,
|
||
hidden_size=self.hidden_dim,
|
||
num_layers=num_layers,
|
||
bidirectional=True,
|
||
batch_first=True,
|
||
dropout=dropout if num_layers > 1 else 0.0,
|
||
)
|
||
|
||
self.proj = nn.Linear(self.hidden_dim * 2, input_dim)
|
||
self.layer_norm = nn.LayerNorm(input_dim)
|
||
|
||
def forward(self, x, mask=None):
|
||
"""
|
||
Args:
|
||
x: [batch, seq_len, input_dim] pinyin embeddings
|
||
mask: [batch, seq_len] optional padding mask (True for valid, False for padding)
|
||
Returns:
|
||
output: [batch, seq_len, input_dim] 每位置的拼音编码
|
||
"""
|
||
total_len = x.size(1)
|
||
|
||
if mask is not None:
|
||
lengths = mask.sum(dim=1).cpu().clamp(min=1)
|
||
packed = nn.utils.rnn.pack_padded_sequence(
|
||
x, lengths, batch_first=True, enforce_sorted=False
|
||
)
|
||
packed_out, (hidden, cell) = self.lstm(packed)
|
||
output, _ = nn.utils.rnn.pad_packed_sequence(
|
||
packed_out, batch_first=True, total_length=total_len
|
||
)
|
||
else:
|
||
output, (hidden, cell) = self.lstm(x)
|
||
|
||
projected = self.proj(output)
|
||
return self.layer_norm(projected)
|
||
|
||
|
||
# ---------------------------- 残差块 ----------------------------
|
||
class ResidualBlock(nn.Module):
|
||
def __init__(self, dim, dropout_prob=0.3):
|
||
super().__init__()
|
||
self.linear1 = nn.Linear(dim, dim)
|
||
self.ln1 = nn.LayerNorm(dim)
|
||
self.linear2 = nn.Linear(dim, dim)
|
||
self.ln2 = nn.LayerNorm(dim)
|
||
self.gelu = nn.GELU()
|
||
self.dropout = nn.Dropout(dropout_prob)
|
||
|
||
def forward(self, x):
|
||
residual = x
|
||
# 修复：使用 self.gelu 而不是未定义的 self.relu
|
||
x = self.gelu(self.linear1(x))
|
||
x = self.ln1(x)
|
||
x = self.linear2(x)
|
||
x = self.ln2(x)
|
||
x = self.dropout(x)
|
||
x = x + residual
|
||
return self.gelu(x)
|
||
|
||
|
||
# ---------------------------- 专家网络 ----------------------------
|
||
class Expert(nn.Module):
|
||
def __init__(
|
||
self,
|
||
input_dim,
|
||
d_model=512,
|
||
num_resblocks=4,
|
||
output_multiplier=1,
|
||
dropout_prob=0.3,
|
||
):
|
||
super().__init__()
|
||
self.output_dim = input_dim * output_multiplier
|
||
self.linear_in = nn.Linear(input_dim, d_model)
|
||
self.res_blocks = nn.ModuleList(
|
||
[ResidualBlock(d_model, dropout_prob) for _ in range(num_resblocks)]
|
||
)
|
||
self.output = nn.Sequential(
|
||
nn.Linear(d_model, d_model),
|
||
nn.GELU(),
|
||
nn.Dropout(dropout_prob),
|
||
nn.Linear(d_model, self.output_dim),
|
||
)
|
||
|
||
def forward(self, x):
|
||
x = self.linear_in(x)
|
||
for block in self.res_blocks:
|
||
x = block(x)
|
||
return self.output(x)
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 1. 上下文编码器 (Context Encoder)
|
||
# 对应 README 4.1: 4层 Transformer, 512维, 输出 H [1]
|
||
# ------------------------------------------------------------------
|
||
class ContextEncoder(nn.Module):
|
||
def __init__(
|
||
self, vocab_size, pinyin_vocab_size, dim=512, n_layers=4, n_heads=4, max_len=128
|
||
):
|
||
super().__init__()
|
||
self.dim = dim
|
||
|
||
# Embeddings
|
||
self.text_emb = AutoModel.from_pretrained(
|
||
"iic/nlp_structbert_backbone_lite_std"
|
||
).embeddings
|
||
self.pinyin_emb = nn.Embedding(pinyin_vocab_size, dim)
|
||
self.pos_emb = nn.Embedding(max_len, dim)
|
||
self.pinyin_pooling = PinyinLSTMEncoder(dim)
|
||
|
||
# Transformer Encoder (4 layers, 4 heads) [1]
|
||
encoder_layer = nn.TransformerEncoderLayer(
|
||
d_model=dim,
|
||
nhead=n_heads,
|
||
dim_feedforward=dim * 4,
|
||
dropout=0.1,
|
||
batch_first=True, # 方便处理 [B, L, D]
|
||
)
|
||
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
|
||
|
||
# LayerNorm for stability
|
||
self.ln = nn.LayerNorm(dim)
|
||
|
||
def forward(self, text_ids, pinyin_ids, mask=None):
|
||
"""
|
||
Args:
|
||
text_ids: [batch, seq_len] 文本 token ids
|
||
pinyin_ids: [batch, pinyin_len] 拼音 token ids
|
||
mask: [batch, seq_len] optional padding mask (1 for valid, 0 for padding)
|
||
Returns:
|
||
H: [batch, seq_len, dim] 文本上下文编码
|
||
P: [batch, pinyin_len, dim] 拼音序列编码（每位置）
|
||
"""
|
||
text_emb = self.text_emb(text_ids) # [B, seq_len, dim]
|
||
|
||
seq_len = text_emb.size(1)
|
||
pos_ids = torch.arange(seq_len, device=text_ids.device).unsqueeze(0)
|
||
x = text_emb + self.pos_emb(pos_ids)
|
||
|
||
if mask is not None:
|
||
src_mask = mask == 0
|
||
else:
|
||
src_mask = None
|
||
|
||
H = self.transformer(x, src_key_padding_mask=src_mask)
|
||
H = self.ln(H)
|
||
|
||
pinyin_emb = self.pinyin_emb(pinyin_ids) # [B, pinyin_len, dim]
|
||
pinyin_mask = pinyin_ids != 0
|
||
P = self.pinyin_pooling(pinyin_emb, mask=pinyin_mask) # [B, pinyin_len, dim]
|
||
|
||
return H, P
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 2. 槽位记忆模块 (Slot Memory)
|
||
# 对应 README 4.2: 8个槽位, 每槽3步, 拼接+位置编码 [1]
|
||
# ------------------------------------------------------------------
|
||
class SlotMemory(nn.Module):
|
||
def __init__(self, vocab_size, max_slots=8, steps_per_slot=3, dim=512):
|
||
super().__init__()
|
||
self.max_slots = max_slots
|
||
self.steps_per_slot = steps_per_slot
|
||
self.total_steps = max_slots * steps_per_slot # 24 steps [1]
|
||
|
||
# Shared embedding layer for history tokens [1]
|
||
self.emb = nn.Embedding(vocab_size, dim)
|
||
|
||
# Learnable positional embeddings for the flattened sequence [1]
|
||
self.pos_emb = nn.Embedding(self.total_steps, dim)
|
||
|
||
# Start token embedding for empty slots [1]
|
||
self.start_emb = nn.Parameter(torch.randn(1, 1, dim))
|
||
|
||
def forward(self, history_ids):
|
||
"""
|
||
Args:
|
||
history_ids: [batch, total_steps]
|
||
Flattened sequence of history tokens.
|
||
Zero positions use start_emb as learnable query.
|
||
Returns:
|
||
S: [batch, total_steps, 512] Slot sequence representation [1]
|
||
"""
|
||
S = self.emb(history_ids) # [B, num_slots, dim]
|
||
|
||
zero_mask = (history_ids == 0).unsqueeze(-1).float() # [B, num_slots, 1]
|
||
S = S * (1 - zero_mask) + self.start_emb * zero_mask
|
||
|
||
pos_ids = (
|
||
torch.arange(S.size(1), device=S.device).unsqueeze(0).expand_as(history_ids)
|
||
)
|
||
S += self.pos_emb(pos_ids)
|
||
|
||
return S
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 3. 交叉注意力融合 (Cross-Attention Fusion)
|
||
# 槽位同时查询文本上下文和拼音序列
|
||
# ------------------------------------------------------------------
|
||
class CrossAttentionFusion(nn.Module):
|
||
"""
|
||
双路交叉注意力：
|
||
- 槽位查询文本上下文 (H)
|
||
- 槽位查询拼音序列 (P)
|
||
通过注意力机制，start_emb 自动学会关注"当前应该预测的拼音位置"
|
||
"""
|
||
|
||
def __init__(self, dim=512, n_heads=4):
|
||
super().__init__()
|
||
self.dim = dim
|
||
self.n_heads = n_heads
|
||
self.head_dim = dim // n_heads
|
||
assert self.head_dim * n_heads == dim
|
||
|
||
self.q_proj = nn.Linear(dim, dim, bias=False)
|
||
self.k_text_proj = nn.Linear(dim, dim, bias=False)
|
||
self.v_text_proj = nn.Linear(dim, dim, bias=False)
|
||
self.k_pinyin_proj = nn.Linear(dim, dim, bias=False)
|
||
self.v_pinyin_proj = nn.Linear(dim, dim, bias=False)
|
||
|
||
self.out_proj = nn.Linear(dim, dim, bias=False)
|
||
self.ln = nn.LayerNorm(dim)
|
||
|
||
def forward(
|
||
self, slots_S, context_H, pinyin_P, context_mask=None, pinyin_mask=None
|
||
):
|
||
"""
|
||
Args:
|
||
slots_S: [batch, num_slots, dim] 槽位编码
|
||
context_H: [batch, ctx_len, dim] 文本上下文编码
|
||
pinyin_P: [batch, pinyin_len, dim] 拼音序列编码
|
||
context_mask: [batch, ctx_len] 文本 padding mask (True for padding)
|
||
pinyin_mask: [batch, pinyin_len] 拼音 padding mask (True for padding)
|
||
Returns:
|
||
fused: [batch, num_slots, dim]
|
||
"""
|
||
batch_size, num_slots, _ = slots_S.shape
|
||
ctx_len = context_H.size(1)
|
||
pinyin_len = pinyin_P.size(1)
|
||
|
||
Q = self.q_proj(slots_S)
|
||
K_text = self.k_text_proj(context_H)
|
||
V_text = self.v_text_proj(context_H)
|
||
K_pinyin = self.k_pinyin_proj(pinyin_P)
|
||
V_pinyin = self.v_pinyin_proj(pinyin_P)
|
||
|
||
Q = Q.view(batch_size, num_slots, self.n_heads, self.head_dim).transpose(1, 2)
|
||
K_text = K_text.view(
|
||
batch_size, ctx_len, self.n_heads, self.head_dim
|
||
).transpose(1, 2)
|
||
V_text = V_text.view(
|
||
batch_size, ctx_len, self.n_heads, self.head_dim
|
||
).transpose(1, 2)
|
||
K_pinyin = K_pinyin.view(
|
||
batch_size, pinyin_len, self.n_heads, self.head_dim
|
||
).transpose(1, 2)
|
||
V_pinyin = V_pinyin.view(
|
||
batch_size, pinyin_len, self.n_heads, self.head_dim
|
||
).transpose(1, 2)
|
||
|
||
text_attn_mask = None
|
||
if context_mask is not None:
|
||
text_attn_mask = (
|
||
context_mask[:, None, None, :]
|
||
.float()
|
||
.masked_fill(context_mask[:, None, None, :], -1e9)
|
||
)
|
||
|
||
pinyin_attn_mask = None
|
||
if pinyin_mask is not None:
|
||
pinyin_attn_mask = (
|
||
pinyin_mask[:, None, None, :]
|
||
.float()
|
||
.masked_fill(pinyin_mask[:, None, None, :], -1e9)
|
||
)
|
||
|
||
text_attn = F.scaled_dot_product_attention(
|
||
Q, K_text, V_text, attn_mask=text_attn_mask
|
||
)
|
||
pinyin_attn = F.scaled_dot_product_attention(
|
||
Q, K_pinyin, V_pinyin, attn_mask=pinyin_attn_mask
|
||
)
|
||
|
||
combined_attn = text_attn + pinyin_attn
|
||
|
||
combined_attn = (
|
||
combined_attn.transpose(1, 2)
|
||
.contiguous()
|
||
.view(batch_size, num_slots, self.dim)
|
||
)
|
||
fused = self.out_proj(combined_attn)
|
||
fused = self.ln(fused + slots_S)
|
||
|
||
return fused
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 4. 专家混合层 (MoE Layer)
|
||
# 对应 README: 20个专家 [1], 使用 components.py 中的 Expert 类
|
||
# ------------------------------------------------------------------
|
||
|
||
|
||
@torch.compiler.allow_in_graph
|
||
def _sparse_moe_dispatch(x_flat, experts, topk_indices, topk_weights, num_experts):
|
||
output = torch.zeros_like(x_flat)
|
||
for e in range(num_experts):
|
||
mask = topk_indices == e
|
||
idx, k_idx = mask.nonzero(as_tuple=True)
|
||
if idx.numel() > 0:
|
||
w = topk_weights[idx, k_idx].unsqueeze(-1)
|
||
output.index_add_(0, idx, (w * experts[e](x_flat[idx])).to(output.dtype))
|
||
return output
|
||
|
||
|
||
class MoELayer(nn.Module):
|
||
"""
|
||
moe_mode 支持三种策略:
|
||
- "all": 计算全部专家，torch.compile 不断裂 (当前默认)
|
||
- "sparse": 只计算被路由到的专家 (产生 graph break)
|
||
- "sparse_allow_graph": 稀疏 MoE，通过 allow_in_graph 避免 graph break
|
||
"""
|
||
|
||
def __init__(
|
||
self, dim=512, num_experts=10, top_k=3, num_resblocks=8, moe_mode="all"
|
||
):
|
||
super().__init__()
|
||
self.num_experts = num_experts
|
||
self.top_k = top_k
|
||
self.dim = dim
|
||
self.moe_mode = moe_mode
|
||
|
||
self.experts = nn.ModuleList(
|
||
[
|
||
Expert(
|
||
input_dim=dim,
|
||
d_model=dim,
|
||
num_resblocks=num_resblocks,
|
||
output_multiplier=1,
|
||
)
|
||
for _ in range(num_experts)
|
||
]
|
||
)
|
||
|
||
self.gate = nn.Linear(dim, num_experts)
|
||
|
||
def forward(self, x):
|
||
B, L, D = x.shape
|
||
num_tokens = B * L
|
||
x_flat = x.view(num_tokens, D)
|
||
|
||
gates = self.gate(x_flat)
|
||
topk_weights, topk_indices = torch.topk(gates, self.top_k, dim=-1)
|
||
topk_weights = F.softmax(topk_weights, dim=-1)
|
||
|
||
if self.moe_mode == "all":
|
||
expert_outputs = torch.stack(
|
||
[expert(x_flat) for expert in self.experts], dim=1
|
||
)
|
||
indices_expanded = topk_indices.unsqueeze(-1).expand(-1, -1, D)
|
||
selected_outputs = torch.gather(expert_outputs, 1, indices_expanded)
|
||
weighted_outputs = selected_outputs * topk_weights.unsqueeze(-1)
|
||
out_flat = weighted_outputs.sum(dim=1)
|
||
|
||
elif self.moe_mode == "sparse":
|
||
out_flat = torch.zeros_like(x_flat)
|
||
for e in range(self.num_experts):
|
||
mask = topk_indices == e
|
||
idx, k_idx = mask.nonzero(as_tuple=True)
|
||
if idx.numel() > 0:
|
||
w = topk_weights[idx, k_idx].unsqueeze(-1)
|
||
out_flat.index_add_(
|
||
0,
|
||
idx,
|
||
(w * self.experts[e](x_flat[idx])).to(out_flat.dtype),
|
||
)
|
||
|
||
elif self.moe_mode == "sparse_allow_graph":
|
||
out_flat = _sparse_moe_dispatch(
|
||
x_flat, self.experts, topk_indices, topk_weights, self.num_experts
|
||
)
|
||
|
||
else:
|
||
raise ValueError(f"Unknown moe_mode: {self.moe_mode}")
|
||
|
||
return out_flat.view(B, L, D)
|