import os

import torch

os.environ["TF_ENABLE_ONEDNN_OPTS"]="0" # 数字必须变成字符串

import jieba
from keras.src.legacy.preprocessing.text import Tokenizer
import torch.nn as nn
# 需要安装 pip install tensorboard
from torch.utils.tensorboard import SummaryWriter   # 词向量可视化工具

if __name__ == '__main__':
    # 1- 准备文本内容
    sentence1 = '传智教育是一家上市公司，旗下有黑马程序员品牌。我是在黑马这里学习人工智能'
    sentence2 = "我爱自然语言处理"
    sentence_list = [sentence1,sentence2]

    # 2- 分词处理
    # word_list最终的结构：[[词11, 词12...], [词21, 词22...]]
    # [['传智', '教育', '是', '一家', '上市公司', '，', '旗下', '有', '黑马', '程序员', '品牌', '。', '我', '是', '在', '黑马', '这里', '学习', '人工智能'], ['我', '爱', '自然语言', '处理']]
    word_list = []
    for sen in sentence_list:
        # lcut不会去重
        word_list.append(jieba.lcut(sen))

    # print(word_list)

    # 3- 训练词汇映射器：去重得到词汇表
    # 3.1- 创建词汇映射器
    my_tokenizer = Tokenizer()
    # 3.2- 训练
    my_tokenizer.fit_on_texts(word_list)
    # 3.3- 获得训练后的结果：词索引与词内容的映射字典
    index_word = my_tokenizer.index_word
    print(index_word)

    # 4- 【掌握】创建词嵌入层word Embedding
    """
        原理：句子 -> 分词 -> 去重得到词汇表 -> 得到词索引 -> Embedding -> 词向量
        
        参数解释：
            num_embeddings：词汇表大小
            embedding_dim：词向量维度。也就是一个词使用一个向量来表示，向量中元素的个数。该值一般是2的n次方，例如：128、256、512....
    """
    ebd = nn.Embedding(num_embeddings=len(index_word),embedding_dim=4)

    # 5- 获得每个词的词向量
    for i in range(len(index_word)):
        # 5.1- 根据词索引得到词向量
        # 注意：传递进去的需要是张量数据类型
        word_vector = ebd(torch.tensor(i))

        # 5.2- 输出
        print(f"词索引{i+1}，词的内容 {index_word[i+1]}，对应的词向量{word_vector}")

    # 6- 【了解】可视化展示
    summary_obj = SummaryWriter(log_dir="../../runs")  # 注意：runs目录不要手动创建；runs的绝对路径中不要有中文
    summary_obj.add_embedding(ebd.weight.data, index_word.values())
    summary_obj.close()