如何实现RNN数据集嵌入器并正确声明rnn_embedder函数参数？

你的函数存在参数与注释不匹配、核心依赖缺失的问题，下面是修正后的实现，同时完成「词索引张量→单词→嵌入向量」的完整流程：

修正后的函数实现 #

import torch
import numpy as np

def rnn_embedder(tensor, vocab, embedding_dim):
    '''
    将词索引张量转换为对应单词及嵌入向量
    参数:
        tensor (torch.Tensor): 由词汇表索引表示的单词张量
        vocab (dict): 词汇映射表，支持两种格式：
            - 单词→索引（如{"hello":0, "world":1}）
            - 索引→单词（如{0:"hello", 1:"world"}）
        embedding_dim (int): 嵌入向量的维度
    返回:
        tuple: (嵌入向量numpy数组, 对应单词列表)
    '''
    # 确保张量为长整型，适配Embedding层输入要求
    tensor = tensor.long()
    
    # 1. 从词索引张量映射到单词
    # 统一为索引→单词的映射格式
    if all(isinstance(k, str) for k in vocab.keys()):
        idx_to_word = {v: k for k, v in vocab.items()}
    else:
        idx_to_word = vocab
    
    words = [idx_to_word[idx.item()] for idx in tensor.flatten()]
    
    # 2. 初始化嵌入层（可替换为预训练嵌入权重）
    vocab_size = len(idx_to_word)
    embedding_layer = torch.nn.Embedding(vocab_size, embedding_dim)
    
    # 3. 生成嵌入向量
    embedding_tensor = embedding_layer(tensor)
    # 转为numpy数组以便存储
    embedding_np = embedding_tensor.detach().numpy()
    
    return embedding_np, words

关键说明 #

• 参数修正：原函数缺少「词汇表映射（vocab）」这个核心参数，没有它无法完成索引到单词的转换；同时将embedding_length更名为embedding_dim，更贴合行业术语
• 流程拆解：
  • 词索引→单词：通过词汇表的反向映射完成索引到单词的转换
  • 单词→标记：这里的「标记」就是输入的词索引，是模型可直接处理的数值形式
  • 标记→嵌入向量：利用PyTorch内置的Embedding层实现，支持随机初始化或加载预训练嵌入权重
• 存储嵌入向量：返回的embedding_np是numpy数组，可直接用np.save()保存；若需保存嵌入层权重，可使用torch.save(embedding_layer.state_dict(), "embedding_weights.pth")

使用示例 #

# 示例词汇表
vocab = {"hello": 0, "world": 1, "rnn": 2}
# 输入词索引张量
input_tensor = torch.tensor([[0, 1], [2, 0]])
# 获取嵌入向量和对应单词
embeddings, words = rnn_embedder(input_tensor, vocab, 10)

print("对应单词:", words)
print("嵌入向量形状:", embeddings.shape)

内容的提问来源于stack exchange，提问作者Elda