Mistral 模型嵌入层的输入要求与 tokenizer 张量返回配置详解-Python教程-PHP中文网

Mistral 等基于 Transformer 的大语言模型，其 embed_tokens 层仅接受 torch.Tensor 类型的 token ID 序列作为输入；直接传入 tokenizer() 返回的 BatchEncoding 或 Python 列表会报错，需显式指定 return_tensors='pt' 或手动转换。

mistral 等基于 transformer 的大语言模型，其 `embed_tokens` 层仅接受 `torch.tensor` 类型的 token id 序列作为输入；直接传入 `tokenizer()` 返回的 `batchencoding` 或 python 列表会报错，需显式指定 `return_tensors='pt'` 或手动转换。

在使用 Mistral（如 mistralai/Mistral-7B-v0.1）进行底层嵌入操作时，开发者常需访问 model.model.embed_tokens（即词嵌入层），以获取原始 token ID 对应的稠密向量表示。该层本质是 torch.nn.Embedding(vocab_size, hidden_size)，例如输出中显示的 Embedding(32000, 4096)，表明其期望输入为形状为 (batch_size, seq_len) 的 整数型 torch.Tensor（dtype 通常为 torch.long），每个元素代表词汇表中的 token ID 索引。

然而，AutoTokenizer 默认返回的是 BatchEncoding 对象（一个字典式容器），其 .input_ids 字段为 Python list，而非张量：

enc = tokenizer("Hello world")
print(type(enc.input_ids))  # <class>
print(enc.input_ids)        # [1, 8754, 743, 2]</class>

因此以下调用均会失败：

model.model.embed_tokens(enc) → ❌ BatchEncoding 不可直接索引；
model.model.embed_tokens(enc.input_ids) → ❌ list 不满足 torch.nn.Embedding 对 indices 参数的 Tensor 类型要求。

✅ 正确做法是：在调用 tokenizer() 时显式指定 return_tensors="pt"，让 tokenizer 直接返回 PyTorch 张量：

# ✅ 推荐：一步到位，返回 torch.Tensor
input_ids = tokenizer("Hello world", return_tensors="pt").input_ids
embedded = model.model.embed_tokens(input_ids)  # shape: [1, seq_len, hidden_size]

print(f"Input IDs shape: {input_ids.shape}")        # e.g., torch.Size([1, 4])
print(f"Embedding shape: {embedded.shape}")         # e.g., torch.Size([1, 4, 4096])
print(f"Data type: {input_ids.dtype}, {embedded.dtype}")  # torch.int64, torch.bfloat16

⚠️ 注意事项：

return_tensors="pt" 是最简洁、安全的方式，避免手动 torch.tensor(...) 转换可能引入的设备（CPU/GPU）和 dtype（如 torch.long）不匹配风险；
若批量处理多条文本，return_tensors="pt" 自动完成 padding 和堆叠，返回二维张量（如 [batch_size, max_seq_len]），可直接送入嵌入层；
确保 input_ids 在与模型相同设备上（如 model.device）。若 tokenizer 在 CPU 上生成张量而模型在 GPU，需显式移动：
```
input_ids = input_ids.to(model.device)
embedded = model.model.embed_tokens(input_ids)
```
不要忽略 tokenizer.pad_token 和 padding_side 配置——虽然嵌入层本身不依赖 padding，但后续训练或推理中若涉及变长序列，统一 padding 策略对 batch 兼容性至关重要。

? 总结：model.model.embed_tokens 的输入契约明确且严格——仅接受 torch.long 类型的 token ID 张量。始终使用 tokenizer(..., return_tensors="pt") 是最佳实践，它既符合类型要求，又保障了设备一致性与批处理友好性，是构建可靠 LLM 微调或特征提取流程的基础环节。