import torch
import torch.nn as nn


def dm01():
    # 按照Bert模型的维度：768维，在这里，我们使用8维度演示
    # num_embeddings：20，词表大小
    # embedding_dim：8，词的维度
    # 把每个token都使用8维向量表示，一共是20个token
    embedded = nn.Embedding(21, 8)

    # 准备一句话，shape（1, 5），1句话，5个token
    input = torch.tensor([[1, 2, 5, 7, 9]])

    # 把这句话送给模型
    result = embedded(input)
    print("result1-->", result)
    # 问题一：这个结果是几维？3维。
    # 问题二：维度分别是多少？【1,5,8】

    # 问题三：embedding的加强
    # 准备一句话，shape（1, 6），1句话，6个token
    input = torch.tensor([[1, 2, 5, 7, 9, 20]])
    result = embedded(input)
    print("result2-->", result)  # 报错，index out of range in self，越界了
    # 解决方案
    # 方式一：只需要把输入数据index调小，最大值是19
    # input = torch.tensor([[1, 2, 5, 7, 9, 19]])
    # 方式二：不需要修改输入数据，只需要调整词表大小
    # embedded = nn.Embedding(21, 8)