| """嵌入层：Token Embedding 词嵌入""" | |
| # 2026-01-23 | |
| import torch.nn as nn | |
| class TokenEmbedding(nn.Module): | |
| """词嵌入""" | |
| def __init__(self, vocab_size, hidden_size): | |
| """ | |
| 初始化词嵌入层 | |
| 参数: | |
| vocab_size: 词汇表大小 | |
| hidden_size: 隐藏层维度 | |
| """ | |
| super().__init__() | |
| self.vocab_size = vocab_size | |
| self.hidden_size = hidden_size | |
| # 词嵌入层：将 token ID 映射为向量 | |
| self.embedding = nn.Embedding(vocab_size, hidden_size) | |
| def forward(self, input_ids): | |
| """ | |
| 前向传播 | |
| 参数: | |
| input_ids: Token ID 张量，形状为 (batch_size, seq_len) | |
| 返回: | |
| 嵌入向量，形状为 (batch_size, seq_len, hidden_size) | |
| """ | |
| return self.embedding(input_ids) | |