doc/word2vec.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e1f1f263",
   "metadata": {},
   "source": [
    "# 中文wiki"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eae26a03",
   "metadata": {},
   "source": [
    "## 数据获取"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d4502cb3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/homebrew/Caskroom/miniconda/base/lib/python3.11/site-packages/gensim/utils.py:1333: UserWarning: detected OSX with python3.8+; aliasing chunkize to chunkize_serial\n",
      "  warnings.warn(\"detected %s; aliasing chunkize to chunkize_serial\" % entity)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save10000articles\n",
      "save20000articles\n",
      "save30000articles\n",
      "save40000articles\n",
      "save50000articles\n",
      "save60000articles\n",
      "save70000articles\n",
      "save80000articles\n",
      "save90000articles\n",
      "save100000articles\n",
      "save110000articles\n",
      "save120000articles\n",
      "save130000articles\n",
      "save140000articles\n",
      "save150000articles\n",
      "save160000articles\n",
      "save170000articles\n",
      "save180000articles\n",
      "save190000articles\n",
      "save200000articles\n",
      "save210000articles\n",
      "save220000articles\n",
      "save230000articles\n",
      "save240000articles\n",
      "save250000articles\n",
      "save260000articles\n",
      "save270000articles\n",
      "save280000articles\n",
      "save290000articles\n",
      "save300000articles\n",
      "save310000articles\n",
      "save320000articles\n",
      "save330000articles\n",
      "save340000articles\n",
      "save350000articles\n",
      "save360000articles\n",
      "save370000articles\n",
      "save380000articles\n",
      "save390000articles\n",
      "save400000articles\n",
      "save410000articles\n",
      "save420000articles\n",
      "save430000articles\n",
      "save440000articles\n",
      "save450000articles\n",
      "save460000articles\n",
      "finish saved\n"
     ]
    }
   ],
   "source": [
    "from gensim.corpora import WikiCorpus\n",
    "import os\n",
    "\n",
    "def parse_corpora():\n",
    "    space = ''\n",
    "    idx = 0\n",
    "    dir_path = \"./data/zhwiki-pages-articles\"\n",
    "    # 检查你的目录是否存在，如果不存在，创建它\n",
    "    if not os.path.exists(dir_path):\n",
    "        os.makedirs(dir_path)\n",
    "\n",
    "    output = open(dir_path + \"/zhwiki_{}.txt\".format(idx),'w',encoding='utf-8')\n",
    "    wiki = WikiCorpus(\"./data/zhwiki-latest-pages-articles.xml.bz2\",dictionary={})\n",
    "    for text in wiki.get_texts():\n",
    "        output.write(space.join(text) + '\\n')\n",
    "        idx += 1\n",
    "        if(idx % 10000 == 0):\n",
    "            print(\"save\" + str(idx) + \"articles\")\n",
    "            output = open(dir_path + \"/zhwiki_{}.txt\".format(int(idx / 10000)),'w',encoding='utf-8')\n",
    "    output.close()\n",
    "    print(\"finish saved\")\n",
    "    \n",
    "parse_corpora()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3efaeccc",
   "metadata": {},
   "source": [
    "## 数据处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3c318366",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "第0个txt文件提取汉字成功\n",
      "第0个txt文件繁体汉字转化简体汉字成功\n",
      "第0个txt文件分词成功\n",
      "第0个txt文件去除停用词成功\n",
      "第1个txt文件提取汉字成功\n",
      "第1个txt文件繁体汉字转化简体汉字成功\n",
      "第1个txt文件分词成功\n",
      "第1个txt文件去除停用词成功\n",
      "第2个txt文件提取汉字成功\n",
      "第2个txt文件繁体汉字转化简体汉字成功\n",
      "第2个txt文件分词成功\n",
      "第2个txt文件去除停用词成功\n",
      "第3个txt文件提取汉字成功\n",
      "第3个txt文件繁体汉字转化简体汉字成功\n",
      "第3个txt文件分词成功\n",
      "第3个txt文件去除停用词成功\n",
      "第4个txt文件提取汉字成功\n",
      "第4个txt文件繁体汉字转化简体汉字成功\n",
      "第4个txt文件分词成功\n",
      "第4个txt文件去除停用词成功\n",
      "第5个txt文件提取汉字成功\n",
      "第5个txt文件繁体汉字转化简体汉字成功\n",
      "第5个txt文件分词成功\n",
      "第5个txt文件去除停用词成功\n",
      "第6个txt文件提取汉字成功\n",
      "第6个txt文件繁体汉字转化简体汉字成功\n",
      "第6个txt文件分词成功\n",
      "第6个txt文件去除停用词成功\n",
      "第7个txt文件提取汉字成功\n",
      "第7个txt文件繁体汉字转化简体汉字成功\n",
      "第7个txt文件分词成功\n",
      "第7个txt文件去除停用词成功\n",
      "第8个txt文件提取汉字成功\n",
      "第8个txt文件繁体汉字转化简体汉字成功\n",
      "第8个txt文件分词成功\n",
      "第8个txt文件去除停用词成功\n",
      "第9个txt文件提取汉字成功\n",
      "第9个txt文件繁体汉字转化简体汉字成功\n",
      "第9个txt文件分词成功\n",
      "第9个txt文件去除停用词成功\n",
      "第10个txt文件提取汉字成功\n",
      "第10个txt文件繁体汉字转化简体汉字成功\n",
      "第10个txt文件分词成功\n",
      "第10个txt文件去除停用词成功\n",
      "第11个txt文件提取汉字成功\n",
      "第11个txt文件繁体汉字转化简体汉字成功\n",
      "第11个txt文件分词成功\n",
      "第11个txt文件去除停用词成功\n",
      "第12个txt文件提取汉字成功\n",
      "第12个txt文件繁体汉字转化简体汉字成功\n",
      "第12个txt文件分词成功\n",
      "第12个txt文件去除停用词成功\n",
      "第13个txt文件提取汉字成功\n",
      "第13个txt文件繁体汉字转化简体汉字成功\n",
      "第13个txt文件分词成功\n",
      "第13个txt文件去除停用词成功\n",
      "第14个txt文件提取汉字成功\n",
      "第14个txt文件繁体汉字转化简体汉字成功\n",
      "第14个txt文件分词成功\n",
      "第14个txt文件去除停用词成功\n",
      "第15个txt文件提取汉字成功\n",
      "第15个txt文件繁体汉字转化简体汉字成功\n",
      "第15个txt文件分词成功\n",
      "第15个txt文件去除停用词成功\n",
      "第16个txt文件提取汉字成功\n",
      "第16个txt文件繁体汉字转化简体汉字成功\n",
      "第16个txt文件分词成功\n",
      "第16个txt文件去除停用词成功\n",
      "第17个txt文件提取汉字成功\n",
      "第17个txt文件繁体汉字转化简体汉字成功\n",
      "第17个txt文件分词成功\n",
      "第17个txt文件去除停用词成功\n",
      "第18个txt文件提取汉字成功\n",
      "第18个txt文件繁体汉字转化简体汉字成功\n",
      "第18个txt文件分词成功\n",
      "第18个txt文件去除停用词成功\n",
      "第19个txt文件提取汉字成功\n",
      "第19个txt文件繁体汉字转化简体汉字成功\n",
      "第19个txt文件分词成功\n",
      "第19个txt文件去除停用词成功\n",
      "第20个txt文件提取汉字成功\n",
      "第20个txt文件繁体汉字转化简体汉字成功\n",
      "第20个txt文件分词成功\n",
      "第20个txt文件去除停用词成功\n",
      "第21个txt文件提取汉字成功\n",
      "第21个txt文件繁体汉字转化简体汉字成功\n",
      "第21个txt文件分词成功\n",
      "第21个txt文件去除停用词成功\n",
      "第22个txt文件提取汉字成功\n",
      "第22个txt文件繁体汉字转化简体汉字成功\n",
      "第22个txt文件分词成功\n",
      "第22个txt文件去除停用词成功\n",
      "第23个txt文件提取汉字成功\n",
      "第23个txt文件繁体汉字转化简体汉字成功\n",
      "第23个txt文件分词成功\n",
      "第23个txt文件去除停用词成功\n",
      "第24个txt文件提取汉字成功\n",
      "第24个txt文件繁体汉字转化简体汉字成功\n",
      "第24个txt文件分词成功\n",
      "第24个txt文件去除停用词成功\n",
      "第25个txt文件提取汉字成功\n",
      "第25个txt文件繁体汉字转化简体汉字成功\n",
      "第25个txt文件分词成功\n",
      "第25个txt文件去除停用词成功\n",
      "第26个txt文件提取汉字成功\n",
      "第26个txt文件繁体汉字转化简体汉字成功\n",
      "第26个txt文件分词成功\n",
      "第26个txt文件去除停用词成功\n",
      "第27个txt文件提取汉字成功\n",
      "第27个txt文件繁体汉字转化简体汉字成功\n",
      "第27个txt文件分词成功\n",
      "第27个txt文件去除停用词成功\n",
      "第28个txt文件提取汉字成功\n",
      "第28个txt文件繁体汉字转化简体汉字成功\n",
      "第28个txt文件分词成功\n",
      "第28个txt文件去除停用词成功\n",
      "第29个txt文件提取汉字成功\n",
      "第29个txt文件繁体汉字转化简体汉字成功\n",
      "第29个txt文件分词成功\n",
      "第29个txt文件去除停用词成功\n",
      "第30个txt文件提取汉字成功\n",
      "第30个txt文件繁体汉字转化简体汉字成功\n",
      "第30个txt文件分词成功\n",
      "第30个txt文件去除停用词成功\n",
      "第31个txt文件提取汉字成功\n",
      "第31个txt文件繁体汉字转化简体汉字成功\n",
      "第31个txt文件分词成功\n",
      "第31个txt文件去除停用词成功\n",
      "第32个txt文件提取汉字成功\n",
      "第32个txt文件繁体汉字转化简体汉字成功\n",
      "第32个txt文件分词成功\n",
      "第32个txt文件去除停用词成功\n",
      "第33个txt文件提取汉字成功\n",
      "第33个txt文件繁体汉字转化简体汉字成功\n",
      "第33个txt文件分词成功\n",
      "第33个txt文件去除停用词成功\n",
      "第34个txt文件提取汉字成功\n",
      "第34个txt文件繁体汉字转化简体汉字成功\n",
      "第34个txt文件分词成功\n",
      "第34个txt文件去除停用词成功\n",
      "第35个txt文件提取汉字成功\n",
      "第35个txt文件繁体汉字转化简体汉字成功\n",
      "第35个txt文件分词成功\n",
      "第35个txt文件去除停用词成功\n",
      "第36个txt文件提取汉字成功\n",
      "第36个txt文件繁体汉字转化简体汉字成功\n",
      "第36个txt文件分词成功\n",
      "第36个txt文件去除停用词成功\n",
      "第37个txt文件提取汉字成功\n",
      "第37个txt文件繁体汉字转化简体汉字成功\n",
      "第37个txt文件分词成功\n",
      "第37个txt文件去除停用词成功\n",
      "第38个txt文件提取汉字成功\n",
      "第38个txt文件繁体汉字转化简体汉字成功\n",
      "第38个txt文件分词成功\n",
      "第38个txt文件去除停用词成功\n",
      "第39个txt文件提取汉字成功\n",
      "第39个txt文件繁体汉字转化简体汉字成功\n",
      "第39个txt文件分词成功\n",
      "第39个txt文件去除停用词成功\n",
      "第40个txt文件提取汉字成功\n",
      "第40个txt文件繁体汉字转化简体汉字成功\n",
      "第40个txt文件分词成功\n",
      "第40个txt文件去除停用词成功\n",
      "第41个txt文件提取汉字成功\n",
      "第41个txt文件繁体汉字转化简体汉字成功\n",
      "第41个txt文件分词成功\n",
      "第41个txt文件去除停用词成功\n",
      "第42个txt文件提取汉字成功\n",
      "第42个txt文件繁体汉字转化简体汉字成功\n",
      "第42个txt文件分词成功\n",
      "第42个txt文件去除停用词成功\n",
      "第43个txt文件提取汉字成功\n",
      "第43个txt文件繁体汉字转化简体汉字成功\n",
      "第43个txt文件分词成功\n",
      "第43个txt文件去除停用词成功\n",
      "第44个txt文件提取汉字成功\n",
      "第44个txt文件繁体汉字转化简体汉字成功\n",
      "第44个txt文件分词成功\n",
      "第44个txt文件去除停用词成功\n",
      "第45个txt文件提取汉字成功\n",
      "第45个txt文件繁体汉字转化简体汉字成功\n",
      "第45个txt文件分词成功\n",
      "第45个txt文件去除停用词成功\n",
      "第46个txt文件提取汉字成功\n",
      "第46个txt文件繁体汉字转化简体汉字成功\n",
      "第46个txt文件分词成功\n",
      "第46个txt文件去除停用词成功\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "import opencc\n",
    "import jieba\n",
    "import os\n",
    "\n",
    "zhwiki_dir_path = \"./data/zhwiki-pages-articles\"\n",
    "stopword_dir_path = zhwiki_dir_path + '/stopword_txt'\n",
    "seg_dir_path = zhwiki_dir_path + '/parse_txt'\n",
    "stopwords_table_path = './data/四川大学机器智能实验室停用词库.txt'\n",
    "\n",
    "cc = opencc.OpenCC('t2s')\n",
    "\n",
    "\n",
    "# 读取停词表，并使用set来存储\n",
    "with open(stopwords_table_path, 'r', encoding='utf-8') as file:\n",
    "    stopwords_set = set(line.strip() for line in file)\n",
    "\n",
    "    \n",
    "def checkDir():\n",
    "    # 检查你的目录是否存在，如果不存在，创建它\n",
    "    if not os.path.exists(zhwiki_dir_path):\n",
    "        os.makedirs(zhwiki_dir_path)\n",
    "    if not os.path.exists(stopword_dir_path):\n",
    "        os.makedirs(stopword_dir_path)\n",
    "    if not os.path.exists(seg_dir_path):\n",
    "        os.makedirs(seg_dir_path)\n",
    "\n",
    "def simplified_Chinese(txt):\n",
    "    txt_sim = []\n",
    "    for sentence in txt.split('\\n'):\n",
    "        txt_sim.append(cc.convert(sentence) + '\\n')\n",
    "        #print(\"第{}句话转化成简体成功\".format(i))\n",
    "    txt = ''.join(txt_sim)\n",
    "    return txt_sim\n",
    "\n",
    "\n",
    "def remove_stopwords(idx):\n",
    "    in_path = seg_dir_path\n",
    "    out_path = stopword_dir_path\n",
    "    \n",
    "    with open(f'{in_path}/zhwikiSegDone_{idx}.txt', 'r', encoding='utf-8') as file_in,\\\n",
    "         open(f'{out_path}/zhwikiStopWord_{idx}.txt', 'w', encoding='utf-8') as file_out:\n",
    "        \n",
    "        # 读取所有行，并将每一行分割成单词\n",
    "        lines = file_in.readlines()\n",
    "        sentence_list = [line.split(' ') for line in lines]\n",
    "        \n",
    "        # 对每一行的每一个单词，如果它不在停词表中，就保留\n",
    "        result = []\n",
    "        for words in sentence_list:\n",
    "            result.append(' '.join(word for word in words if word not in stopwords_set))\n",
    "\n",
    "        # 将结果写入文件\n",
    "        file_out.write('\\n'.join(result))\n",
    "\n",
    "\n",
    "\n",
    "def seg_done(idx, txt):\n",
    "    # 以下为分词部分\n",
    "    out_path = seg_dir_path\n",
    "    file = open(out_path + '/zhwikiSegDone_{}.txt'.format(idx),'w',encoding='utf-8')\n",
    "    file.write(' '.join(jieba.cut(txt, cut_all=False)).replace(' \\n ', '\\n'))\n",
    "    file.close()\n",
    "\n",
    "\n",
    "def parse_txt():\n",
    "    in_path = zhwiki_dir_path\n",
    "    for i in range(0, 47):      # 理论上应该是从0至47\n",
    "        file = open(in_path+'/zhwiki_{}.txt'.format(i),'r',encoding='utf-8')\n",
    "        txt = file.read()\n",
    "        file.close()\n",
    "        # 1. 提取汉字\n",
    "        txt = ''.join(re.findall('[\\u4e00-\\u9fa5|\\n]',txt))      # 只保留汉字,如果其后有空格则保留\n",
    "        print('第' + str(i) + '个txt文件提取汉字成功')\n",
    "        # 2. 简化汉字\n",
    "        txt = simplified_Chinese(txt)\n",
    "        print('第' + str(i) + '个txt文件繁体汉字转化简体汉字成功')\n",
    "        # 3. 汉字分词\n",
    "        seg_done(i, txt)\n",
    "        print('第' + str(i) + '个txt文件分词成功')\n",
    "        # 4. 去除停用词\n",
    "        remove_stopwords(i)\n",
    "        print('第' + str(i) + '个txt文件去除停用词成功')\n",
    "\n",
    "checkDir()\n",
    "parse_txt()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ed3af6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "5658ca84",
   "metadata": {},
   "source": [
    "# 清华大学自然语言处理实验室数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e691f9eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "处理类型：体育\n",
      "  已处理体育：0\n",
      "  已处理体育：10000\n",
      "  已处理体育：20000\n",
      "  已处理体育：30000\n",
      "  已处理体育：40000\n",
      "  已处理体育：50000\n",
      "  已处理体育：60000\n",
      "  已处理体育：70000\n",
      "  已处理体育：80000\n",
      "  已处理体育：90000\n",
      "  已处理体育：100000\n",
      "  已处理体育：110000\n",
      "  已处理体育：120000\n",
      "  已处理体育：130000\n",
      "处理类型：娱乐\n",
      "  已处理娱乐：0\n",
      "  已处理娱乐：10000\n",
      "  已处理娱乐：20000\n",
      "  已处理娱乐：30000\n",
      "  已处理娱乐：40000\n",
      "  已处理娱乐：50000\n",
      "  已处理娱乐：60000\n",
      "  已处理娱乐：70000\n",
      "  已处理娱乐：80000\n",
      "  已处理娱乐：90000\n",
      "处理类型：家居\n",
      "  已处理家居：0\n",
      "  已处理家居：10000\n",
      "  已处理家居：20000\n",
      "  已处理家居：30000\n",
      "处理类型：彩票\n",
      "  已处理彩票：0\n",
      "处理类型：房产\n",
      "  已处理房产：0\n",
      "  已处理房产：10000\n",
      "  已处理房产：20000\n",
      "处理类型：教育\n",
      "  已处理教育：0\n",
      "  已处理教育：10000\n",
      "  已处理教育：20000\n",
      "  已处理教育：30000\n",
      "  已处理教育：40000\n",
      "处理类型：时尚\n",
      "  已处理时尚：0\n",
      "  已处理时尚：10000\n",
      "处理类型：时政\n",
      "  已处理时政：0\n",
      "  已处理时政：10000\n",
      "  已处理时政：20000\n",
      "  已处理时政：30000\n",
      "  已处理时政：40000\n",
      "  已处理时政：50000\n",
      "  已处理时政：60000\n",
      "处理类型：星座\n",
      "  已处理星座：0\n",
      "处理类型：游戏\n",
      "  已处理游戏：0\n",
      "  已处理游戏：10000\n",
      "  已处理游戏：20000\n",
      "处理类型：社会\n",
      "  已处理社会：0\n",
      "  已处理社会：10000\n",
      "  已处理社会：20000\n",
      "  已处理社会：30000\n",
      "  已处理社会：40000\n",
      "  已处理社会：50000\n",
      "处理类型：科技\n",
      "  已处理科技：0\n",
      "  已处理科技：10000\n",
      "  已处理科技：20000\n",
      "  已处理科技：30000\n",
      "  已处理科技：40000\n",
      "  已处理科技：50000\n",
      "  已处理科技：60000\n",
      "  已处理科技：70000\n",
      "  已处理科技：80000\n",
      "  已处理科技：90000\n",
      "  已处理科技：100000\n",
      "  已处理科技：110000\n",
      "  已处理科技：120000\n",
      "  已处理科技：130000\n",
      "  已处理科技：140000\n",
      "  已处理科技：150000\n",
      "  已处理科技：160000\n",
      "处理类型：股票\n",
      "  已处理股票：0\n",
      "  已处理股票：10000\n",
      "  已处理股票：20000\n",
      "  已处理股票：30000\n",
      "  已处理股票：40000\n",
      "  已处理股票：50000\n",
      "  已处理股票：60000\n",
      "  已处理股票：70000\n",
      "  已处理股票：80000\n",
      "  已处理股票：90000\n",
      "  已处理股票：100000\n",
      "  已处理股票：110000\n",
      "  已处理股票：120000\n",
      "  已处理股票：130000\n",
      "  已处理股票：140000\n",
      "  已处理股票：150000\n",
      "处理类型：财经\n",
      "  已处理财经：0\n",
      "  已处理财经：10000\n",
      "  已处理财经：20000\n",
      "  已处理财经：30000\n",
      "财经文本处理完毕\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "import jieba\n",
    "\n",
    "stopwords_table_path = './data/四川大学机器智能实验室停用词库.txt'\n",
    "\n",
    "# 读取停词表，并使用set来存储\n",
    "with open(stopwords_table_path, 'r', encoding='utf-8') as file:\n",
    "    stopwords_set = set(line.strip() for line in file)\n",
    "\n",
    "NewsCatalog = ['体育','娱乐','家居','彩票','房产','教育','时尚','时政','星座','游戏','社会','科技','股票','财经']\n",
    "\n",
    "dir_path = './data/THUCNews'\n",
    "\n",
    "def list_all_files(path):\n",
    "    result = []\n",
    "    for root, dirs, files in os.walk(path):\n",
    "        for file in files:\n",
    "            if file != '.DS_Store':  # 添加此行来跳过.DS_Store文件\n",
    "                result.append(os.path.join(root, file))\n",
    "    return result\n",
    "\n",
    "for category in NewsCatalog:\n",
    "    categorya_dir_path = dir_path + \"/\" + category\n",
    "    if not os.path.exists(categorya_dir_path):\n",
    "        os.makedirs(categorya_dir_path)\n",
    "        \n",
    "    combine = open(dir_path + '/' + '{}.txt'.format(category), 'w', encoding='utf-8')\n",
    "    sentence = []\n",
    "    idx = 0\n",
    "    print(\"处理类型：{}\".format(category))\n",
    "    for file_path in list_all_files(categorya_dir_path):\n",
    "        if idx % 10000 == 0:\n",
    "            print(\"  已处理{}：{}\".format(category,idx))\n",
    "        file = open(file_path, 'r', encoding='utf-8')\n",
    "        txt = file.read().replace('\\n　　',' ')      # 一篇文章为一排\n",
    "        file.close()\n",
    "        # 提取中文\n",
    "        txt = ''.join(re.findall('[\\u4e00-\\u9fa5| |]', txt))\n",
    "        # 分词\n",
    "        txt = ' '.join(jieba.cut(txt, cut_all=False)).replace('   ',' ')\n",
    "        # 删除停用词\n",
    "        for word in txt.split(' '):\n",
    "            if word in stopwords_set:\n",
    "                txt = txt.replace(word+' ','')\n",
    "        sentence.append(txt+'\\n')\n",
    "        idx += 1 \n",
    "    combine.write(''.join(sentence))\n",
    "    print(\"总和：{}\".formato(idx))\n",
    "print('文本处理完毕')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d245fe28",
   "metadata": {},
   "source": [
    "# 合并数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "30f3dd92",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "中文wiki百科文件合并完成\n",
      "THU数据集合并完成\n"
     ]
    }
   ],
   "source": [
    "NewsCatalog = ['体育','娱乐','家居','彩票','房产','教育','时尚','时政','星座','游戏','社会','科技','股票','财经']\n",
    "\n",
    "path = './data'\n",
    "wiki_path = './data/zhwiki-pages-articles/stopword_txt'\n",
    "THUCNews_path = './data/THUCNews'\n",
    "\n",
    "Data = open(path + '/data.txt', 'a', encoding='utf-8')\n",
    "\n",
    "for i in range(47):      # 合并中文wiki百科文件\n",
    "    file = open(wiki_path + '/zhwikiStopWord_{}.txt'.format(i), 'r', encoding='utf-8')\n",
    "    txt = file.read().strip('\\n').strip(' ')\n",
    "    Data.write(txt + '\\n')\n",
    "    file.close()\n",
    "\n",
    "print('中文wiki百科文件合并完成')\n",
    "\n",
    "for item in NewsCatalog:        # 合并THU数据集\n",
    "    file = open(THUCNews_path + '/{}.txt'.format(item), 'r', encoding='utf-8')\n",
    "    txt = file.read().strip('\\n').strip(' ')\n",
    "    Data.write(txt + '\\n')\n",
    "    file.close()\n",
    "\n",
    "print('THU数据集合并完成')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ddda36e",
   "metadata": {},
   "source": [
    "# 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "d99cd8d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-11 19:57:36.722 | INFO     | __main__:<module>:28 - Starting to load sentences from %s\n",
      "2024-01-11 19:57:36.724 | INFO     | __main__:<module>:30 - Finished loading sentences\n",
      "2024-01-11 19:57:36.725 | INFO     | __main__:<module>:43 - Starting to build Word2Vec model\n",
      "2024-01-11 19:59:11.095 | INFO     | __main__:on_epoch_begin:14 - Epoch #0 start\n",
      "2024-01-11 20:07:37.577 | INFO     | __main__:on_epoch_end:19 - Epoch #0 end, loss: 80493176.0, duration: 600.8519821166992\n",
      "2024-01-11 20:07:37.579 | INFO     | __main__:on_epoch_begin:14 - Epoch #1 start\n",
      "2024-01-11 21:56:16.634 | INFO     | __main__:on_epoch_end:19 - Epoch #1 end, loss: 109677968.0, duration: 6519.055454015732\n",
      "2024-01-11 21:56:16.636 | INFO     | __main__:on_epoch_begin:14 - Epoch #2 start\n",
      "2024-01-12 00:45:29.837 | INFO     | __main__:on_epoch_end:19 - Epoch #2 end, loss: 134217728.0, duration: 10153.201048135757\n",
      "2024-01-12 00:45:29.838 | INFO     | __main__:on_epoch_begin:14 - Epoch #3 start\n",
      "2024-01-12 03:17:13.699 | INFO     | __main__:on_epoch_end:19 - Epoch #3 end, loss: 134217728.0, duration: 9103.861315011978\n",
      "2024-01-12 03:17:13.701 | INFO     | __main__:on_epoch_begin:14 - Epoch #4 start\n",
      "2024-01-12 05:56:07.317 | INFO     | __main__:on_epoch_end:19 - Epoch #4 end, loss: 134217728.0, duration: 9533.615617990494\n",
      "2024-01-12 05:56:07.318 | INFO     | __main__:on_epoch_begin:14 - Epoch #5 start\n",
      "2024-01-12 08:17:16.458 | INFO     | __main__:on_epoch_end:19 - Epoch #5 end, loss: 134217728.0, duration: 8469.140832901001\n",
      "2024-01-12 08:17:16.460 | INFO     | __main__:on_epoch_begin:14 - Epoch #6 start\n",
      "2024-01-12 11:06:56.453 | INFO     | __main__:on_epoch_end:19 - Epoch #6 end, loss: 134217728.0, duration: 10179.993557929993\n",
      "2024-01-12 11:06:56.455 | INFO     | __main__:on_epoch_begin:14 - Epoch #7 start\n",
      "2024-01-12 11:15:15.114 | INFO     | __main__:on_epoch_end:19 - Epoch #7 end, loss: 134217728.0, duration: 498.6591639518738\n",
      "2024-01-12 11:15:15.115 | INFO     | __main__:on_epoch_begin:14 - Epoch #8 start\n",
      "2024-01-12 11:23:41.653 | INFO     | __main__:on_epoch_end:19 - Epoch #8 end, loss: 134217728.0, duration: 506.53724789619446\n",
      "2024-01-12 11:23:41.654 | INFO     | __main__:on_epoch_begin:14 - Epoch #9 start\n",
      "2024-01-12 11:31:59.078 | INFO     | __main__:on_epoch_end:19 - Epoch #9 end, loss: 134217728.0, duration: 497.4236319065094\n",
      "2024-01-12 11:31:59.080 | INFO     | __main__:<module>:45 - Finished building Word2Vec model\n",
      "2024-01-12 11:31:59.081 | INFO     | __main__:<module>:47 - Starting to save model to %s\n",
      "2024-01-12 11:32:00.169 | INFO     | __main__:<module>:49 - Finished saving model\n"
     ]
    }
   ],
   "source": [
    "from gensim.models import Word2Vec\n",
    "from gensim.models import word2vec\n",
    "from loguru import logger\n",
    "import time\n",
    "from gensim.models.callbacks import CallbackAny2Vec\n",
    "\n",
    "# 定义回调函数类，用于在每个epoch结束时记录训练信息\n",
    "class EpochLogger(CallbackAny2Vec):\n",
    "    def __init__(self):\n",
    "        self.epoch = 0\n",
    "        self.time = time.time()\n",
    "\n",
    "    def on_epoch_begin(self, model):\n",
    "        logger.info(\"Epoch #{} start\".format(self.epoch))\n",
    "\n",
    "    def on_epoch_end(self, model):\n",
    "        loss = model.get_latest_training_loss()\n",
    "        duration = time.time() - self.time\n",
    "        logger.info(\"Epoch #{} end, loss: {}, duration: {}\".format(self.epoch, loss, duration))\n",
    "        model.save('word2vec_model_{}.model'.format(self.epoch))  # 保存模型，模型名中包含epoch数\n",
    "        self.epoch += 1\n",
    "        self.time = time.time()\n",
    "\n",
    "# 打印日志\n",
    "logger.add(\"word2vec_training.log\")\n",
    "\n",
    "path = './data/data.txt'\n",
    "\n",
    "logger.info('Starting to load sentences from %s', path)\n",
    "sentences = word2vec.LineSentence(path)\n",
    "logger.info('Finished loading sentences')\n",
    "\n",
    "# 提前定义callback\n",
    "epoch_logger = EpochLogger()\n",
    "\n",
    "# sg——word2vec两个模型的选择。如果是0， 则是CBOW模型，是1则是Skip-Gram模型，默认是0即CBOW模型\n",
    "# hs——word2vec两个解法的选择，如果是0， 则是Negative Sampling，是1的话并且负采样个数negative大于0， 则是Hierarchical Softmax。默认是0即Negative Sampling\n",
    "# negative——即使用Negative Sampling时负采样的个数，默认是5。推荐在[3,10]之间\n",
    "# min_count——需要计算词向量的最小词频。这个值可以去掉一些很生僻的低频词，默认是5。如果是小语料，可以调低这个值\n",
    "# iter——随机梯度下降法中迭代的最大次数，默认是5。对于大语料，可以增大这个值\n",
    "# alpha——在随机梯度下降法中迭代的初始步长。算法原理篇中标记为η，默认是0.025\n",
    "# min_alpha——由于算法支持在迭代的过程中逐渐减小步长，min_alpha给出了最小的迭代步长值\n",
    "# 训练模型\n",
    "logger.info('Starting to build Word2Vec model')\n",
    "model = Word2Vec(sentences, vector_size=300, window=5, epochs=10, compute_loss=True, callbacks=[epoch_logger])\n",
    "logger.info('Finished building Word2Vec model')\n",
    "model_path = 'word2vec.model'\n",
    "logger.info('Starting to save model to %s', model_path)\n",
    "model.save(model_path)\n",
    "logger.info('Finished saving model')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db11f9c9",
   "metadata": {},
   "source": [
    "# 模型使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "460b3028",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "300\n",
      "4965.289603250014\n",
      "[('狗', 0.7110657691955566), ('猫咪', 0.671169638633728), ('小猫', 0.6650978326797485), ('兔子', 0.6501124501228333), ('小狗', 0.6325607895851135), ('小猫咪', 0.6306896805763245), ('犬', 0.6204975843429565), ('宠物猫', 0.6035280227661133), ('吉娃娃', 0.5858094096183777), ('宠物狗', 0.5799086093902588)]\n",
      "[('东北师范大学', 0.7164520621299744), ('大连理工大学', 0.6689789295196533), ('哈尔滨工程大学', 0.6448072791099548), ('哈尔滨工业大学', 0.6404638886451721), ('西北农林科技大学', 0.637593150138855), ('东北财经大学', 0.6347334384918213), ('中国医科大学', 0.630100429058075), ('杭州大学', 0.6278619766235352), ('华东政法学院', 0.6263515949249268), ('白求恩医科大学', 0.6260219216346741)]\n",
      "[ 0.1132143  -0.97853965  0.05373773  0.14764664  0.18700294  1.2022227\n",
      " -0.21538667  0.02402738 -1.06705     0.11823492  0.10939491  0.8852046\n",
      " -0.46388403 -0.47683364 -0.40410247  0.6492659   1.0039027   0.62877876\n",
      " -0.4860649   1.5025599  -0.4394662   0.2118384  -0.09238251 -0.44268852\n",
      " -0.50969356 -1.0903822   0.17130436 -0.10265341 -0.80796087  0.664532\n",
      " -1.0039085   0.25025678  0.04414725 -0.01792347  0.7919607   0.41984817\n",
      " -0.3575323   0.18697622 -0.07712696 -0.10364122 -0.40916306  1.0299633\n",
      " -0.56382954  0.27914467  0.9686123   0.8203663  -0.46985823 -0.93142617\n",
      "  0.10248896  0.7577041   0.47676733  0.5174965   0.65654194  0.28826612\n",
      " -0.3585663   0.7022999   0.8263731  -0.04216862  0.02500334 -0.7070341\n",
      "  0.25324202  0.649409    1.240935    0.3266476   0.630639   -1.0397396\n",
      "  0.3424891  -0.7646398  -0.19346936  0.2914219   0.98589504 -0.7022151\n",
      " -0.27375183 -0.1522235  -0.6156613  -0.72736305  0.07530206 -0.24211873\n",
      "  0.64059955  1.2007562   0.6956061   0.32782248 -0.73003155 -0.45203686\n",
      " -0.12535138 -1.1661445   0.568099    0.22218777  0.48054948  0.20383002\n",
      " -0.49496123 -0.24454357  1.1155365  -0.5667631   0.1183758  -0.20017321\n",
      "  0.61268234 -0.7823204   0.5108612   0.14948742  1.1731504   0.61714953\n",
      "  0.52926683 -0.3168325   0.22383472  1.3645821   0.48143372  0.4869122\n",
      " -1.1978815   1.3120998  -0.25365075 -0.7019392  -0.28617772 -1.1574032\n",
      " -0.5979565  -0.7429307   0.6758254  -0.31816763 -0.5543492  -0.8312415\n",
      "  0.12124222  0.34394652  0.3149824  -0.60783565  0.32489538 -0.29193026\n",
      "  1.1530893  -0.4766924  -0.67198545 -0.09657984  0.7584408  -0.21362124\n",
      " -0.01491908 -0.31326288  0.16493037 -0.37198356 -0.40049285 -0.67375666\n",
      " -0.3091159  -0.54321146  1.1786495  -0.5601435   0.3725502  -0.14363445\n",
      "  0.07417241  0.0858407  -0.09605137 -0.3465434  -0.8076209  -0.47772712\n",
      "  0.02103014  0.10443293  1.0851771  -0.3794516  -0.22299889  0.4570437\n",
      " -0.5327905  -0.08009802  0.23383555 -0.33382177  0.81951284 -0.44892916\n",
      "  0.3442512  -1.167091   -1.1099758  -0.8648395  -0.55939156 -0.42185715\n",
      "  0.09658866 -0.34348848 -1.1694944  -1.023533    0.3585494   0.5340697\n",
      "  0.62883943  0.38963985  0.5679375   0.17269118 -0.0363089   0.08183577\n",
      " -0.81524056 -0.5103102  -0.30726776  0.3712215  -0.1217225  -0.5629984\n",
      "  0.2826522  -0.62297547  0.6323471   0.5705954  -0.9228098   0.15771797\n",
      " -0.06817795 -1.0059818  -0.7157485   0.01633879 -1.3137162  -0.8303017\n",
      " -0.3240206   1.036065    0.24214198  0.18876825 -0.14052728  0.7027756\n",
      "  0.728728   -0.2808199   0.61719835 -0.8220316  -0.3327694   0.63173074\n",
      "  0.4710167   0.49089044 -0.5921744  -0.19216907  0.1612914  -0.17338698\n",
      "  0.549266   -0.7131234  -1.5237281  -0.40163565 -1.0692923   0.67473805\n",
      "  0.15759072 -0.39806348 -0.54464644 -0.18551251 -0.5637427  -1.30043\n",
      " -0.33715084  0.5585721   0.4028558  -0.2981486  -0.96702737  1.0499699\n",
      " -1.1605064  -0.8771676  -0.2903848   0.42213959  0.53413737  0.20702627\n",
      "  0.3255487  -0.8504953  -0.7580894   0.48683533  0.17126118 -0.42370135\n",
      "  0.22639593 -0.31918603  0.47024855 -0.1242829  -0.08245605 -0.52308935\n",
      " -0.7984175  -0.30186662 -0.18843462  0.2574621  -0.52304614 -0.5546153\n",
      " -0.17558096  0.34805897 -0.6047662   0.11681713  0.38815215  0.84626627\n",
      "  0.4415602   0.9326162  -0.02561045 -0.140011    0.01838834  0.4470796\n",
      "  0.7339259   0.16596353  0.66341037  0.6561227  -0.39067554 -0.9236066\n",
      "  0.8950667   1.5514098   0.08265802 -0.21772468  0.977433   -0.58401597\n",
      " -0.6834202   0.1735257  -1.1483457   0.19879417 -0.30087522 -0.8386484\n",
      " -0.1907756   0.93102986 -0.54182345 -0.31680775 -0.06809943 -0.6149766\n",
      "  0.7648053   0.26587275  0.22695826  0.11216222 -0.29015964 -0.8306872 ]\n"
     ]
    }
   ],
   "source": [
    "import jieba\n",
    "\n",
    "model = Word2Vec.load('word2vec.model')\n",
    "\n",
    "# 加载模型\n",
    "def load_word2vec_model(w2v_path):\n",
    "    model = word2vec.Word2Vec.load(w2v_path)\n",
    "    return model\n",
    "\n",
    "\n",
    "# 计算词语最相似的词\n",
    "def calculate_most_similar(self, word):\n",
    "    similar_words = self.wv.most_similar(word)\n",
    "    print(word)\n",
    "    for term in similar_words:\n",
    "        print(term[0], term[1])\n",
    "\n",
    "\n",
    "# 计算两个词相似度\n",
    "def calculate_words_similar(self, word1, word2):\n",
    "    print(self.wv.similarity(word1, word2))\n",
    "\n",
    "\n",
    "# 找出不合群的词\n",
    "def find_word_dismatch(self, list):\n",
    "    print(self.wv.doesnt_match(list))\n",
    "    \n",
    "def sentence_to_vector(sentence, model):\n",
    "    words = jieba.cut(sentence, cut_all=False)# 分词\n",
    "    word_vectors = [model.wv[word] for word in words if word in model.wv.key_to_index]  # 获取句子中每个词的词向量\n",
    "    if len(word_vectors) == 0:\n",
    "        return None  # 当句子中的词都不在词向量模型的词汇表中时返回None\n",
    "    sentence_vector = sum(word_vectors) / len(word_vectors)  # 平均词向量作为句子向量\n",
    "    return sentence_vector\n",
    "\n",
    "\n",
    "print(model.vector_size)\n",
    "# print(model.accuracy)\n",
    "print(model.total_train_time)\n",
    "# print(model.wv)\n",
    "print(model.wv.most_similar('猫'))\n",
    "print(model.wv.most_similar('吉林大学'))\n",
    "print(sentence_to_vector('你有没有听说过侠客行的故事', model))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a673ff82",
   "metadata": {},
   "source": [
    "# 模型封装整理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f473dd7b",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-16 14:35:27.858 | DEBUG    | __main__:__init__:42 - Load w2v from ./word2vec.model, spend 0.00 sec\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([-0.00971168, -0.04826424, -0.00908102, -0.0255302 , -0.0049236 ,\n",
       "       -0.01438301,  0.04413503, -0.02452322, -0.01416836,  0.00758989,\n",
       "       -0.02274896, -0.01674865, -0.00534592, -0.04227909, -0.04475795,\n",
       "        0.05243133,  0.01878846,  0.01741701, -0.01213578,  0.03605546,\n",
       "       -0.06731813,  0.00357951, -0.06001818,  0.0485325 ,  0.02999079,\n",
       "       -0.02109329, -0.01103431, -0.0232751 , -0.03735034,  0.00409066,\n",
       "       -0.04482819, -0.00952696,  0.01776599, -0.03094684,  0.00796766,\n",
       "       -0.02564649,  0.02180307,  0.00201033, -0.00390506, -0.0262962 ,\n",
       "        0.01786203,  0.0261272 ,  0.00363575, -0.05368124, -0.00087189,\n",
       "       -0.01614523,  0.00941788, -0.00405458, -0.0009126 ,  0.04293044,\n",
       "        0.01348076,  0.02482552, -0.02495426,  0.08352873,  0.05218577,\n",
       "        0.02552636,  0.04842734, -0.01041524, -0.00184335, -0.04186204,\n",
       "        0.00382873,  0.04991162,  0.05825759, -0.00431897,  0.01830918,\n",
       "       -0.0179254 ,  0.01311377, -0.08098135,  0.0094429 ,  0.05149109,\n",
       "        0.01464755, -0.00563559, -0.00633816,  0.02260829, -0.01177813,\n",
       "       -0.03842177, -0.01145056,  0.03803317,  0.02654381,  0.04959752,\n",
       "        0.02055019,  0.04525859, -0.04733358, -0.00848596,  0.03011006,\n",
       "       -0.08951913,  0.00810363, -0.02332886,  0.0036681 , -0.00620404,\n",
       "       -0.02755947, -0.0062119 ,  0.0484091 , -0.02453121,  0.0136365 ,\n",
       "       -0.01656763, -0.01803135,  0.02424384,  0.03932329,  0.02536655,\n",
       "        0.0156375 ,  0.05189383,  0.00368451, -0.04472817,  0.01919022,\n",
       "        0.09611236, -0.03537383,  0.03146692, -0.02483856, -0.01417733,\n",
       "       -0.00242565,  0.00469232,  0.02028014, -0.04565012,  0.02033007,\n",
       "       -0.03828707,  0.01145232, -0.00703325, -0.0325814 , -0.00908688,\n",
       "       -0.01221111,  0.01725346, -0.03927354,  0.02376855, -0.06051896,\n",
       "        0.05104158,  0.00277323, -0.04824781,  0.02521803,  0.01341386,\n",
       "        0.02878465, -0.0005022 , -0.00520125,  0.00250364,  0.00092492,\n",
       "        0.00195595,  0.03273427, -0.01070663, -0.06290315, -0.00911458,\n",
       "        0.03185628,  0.0259452 ,  0.01755554,  0.00542306, -0.00106431,\n",
       "       -0.02150373, -0.0097809 , -0.02054908, -0.02723214, -0.01719017,\n",
       "       -0.0191807 ,  0.02319624,  0.00617398, -0.0153762 ,  0.04146733,\n",
       "       -0.00146514, -0.00761952,  0.04216683,  0.00610778,  0.0189975 ,\n",
       "        0.04436166, -0.0294796 ,  0.053834  , -0.05160212,  0.00279633,\n",
       "       -0.00203449,  0.02838268,  0.01490771, -0.02814383,  0.00612153,\n",
       "       -0.0505361 , -0.0607308 ,  0.01684631,  0.01262455,  0.04369979,\n",
       "        0.05679374,  0.00121475,  0.00345395,  0.02656786,  0.01353406,\n",
       "       -0.00122587, -0.03182663, -0.03802079,  0.06843756,  0.02408378,\n",
       "        0.01018853,  0.01771774, -0.02847864,  0.00637929,  0.00385173,\n",
       "        0.00375585,  0.00820121,  0.04383601, -0.03017748, -0.03137831,\n",
       "       -0.01879181,  0.01071443,  0.01721148,  0.00732773,  0.04019905,\n",
       "        0.00775822, -0.0551675 ,  0.02872854,  0.06168383, -0.03300812,\n",
       "       -0.02469163,  0.10567348, -0.05612998, -0.01086914, -0.02480209,\n",
       "        0.03762581, -0.00201162, -0.07710073, -0.04326778, -0.01478149,\n",
       "        0.02889041, -0.01204577,  0.06261262, -0.06758116, -0.04837571,\n",
       "       -0.05493668,  0.07846469,  0.05927523, -0.02695504, -0.02852703,\n",
       "       -0.0122641 , -0.05349978, -0.03053444, -0.01343611,  0.02056761,\n",
       "       -0.00373374, -0.03360544, -0.00562876,  0.0340185 , -0.03100684,\n",
       "       -0.08553706, -0.01329791,  0.02840758,  0.02002341,  0.06147271,\n",
       "       -0.00837766, -0.02349422, -0.05712864,  0.00449986,  0.01649672,\n",
       "       -0.02094095,  0.05695444, -0.04894469,  0.05436326, -0.05972929,\n",
       "       -0.01016123,  0.01153997, -0.01293384,  0.01186467, -0.02139057,\n",
       "        0.00218127, -0.05465185,  0.03206271, -0.0077856 , -0.02034308,\n",
       "       -0.02975587, -0.02727948, -0.02746982,  0.01703248,  0.05459865,\n",
       "       -0.03147747, -0.02781112, -0.04131823, -0.05516142,  0.01829861,\n",
       "        0.01475928,  0.01657524,  0.00498256,  0.03687517, -0.04392608,\n",
       "        0.00966434,  0.02465582,  0.03151698,  0.00556606, -0.06580236,\n",
       "        0.06928153, -0.01945015,  0.00552959,  0.02488908, -0.09265882,\n",
       "        0.02932162, -0.04511765, -0.03362798, -0.01427745,  0.05174087,\n",
       "        0.01171231, -0.05505661, -0.01201095, -0.07107122, -0.01121976,\n",
       "       -0.01977549,  0.02873673,  0.03292044, -0.04271439, -0.00314295])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "from typing import List, Union, Dict\n",
    "\n",
    "from gensim.models import Word2Vec\n",
    "import jieba\n",
    "import numpy as np\n",
    "from loguru import logger\n",
    "from numpy import ndarray\n",
    "from tqdm import tqdm\n",
    "\n",
    "stopwords_table_path = './data/四川大学机器智能实验室停用词库.txt'\n",
    "\n",
    "\n",
    "def load_stopwords(file_path):\n",
    "    # 读取停词表，并使用set来存储\n",
    "    with open(stopwords_table_path, 'r', encoding='utf-8') as file:\n",
    "        stopwords_set = set(line.strip() for line in file)\n",
    "        return stopwords_set\n",
    "\n",
    "\n",
    "class Word2VecManager:\n",
    "    \"\"\"Pre-trained word2vec embedding\"\"\"\n",
    "    def __init__(self, model_name_or_path: str = './word2vec.model',\n",
    "                 w2v_kwargs: Dict = None,\n",
    "                 stopwords: List[str] = None):\n",
    "        \"\"\"\n",
    "        Init word2vec model\n",
    "\n",
    "        Args:\n",
    "            model_name_or_path: word2vec file path\n",
    "            w2v_kwargs: dict, params pass to the ``load_word2vec_format()`` function of ``gensim.models.KeyedVectors`` -\n",
    "                https://radimrehurek.com/gensim/models/keyedvectors.html#module-gensim.models.keyedvectors\n",
    "            stopwords: list, stopwords\n",
    "        \"\"\"\n",
    "        from gensim.models import KeyedVectors  # noqa\n",
    "\n",
    "        self.w2v_kwargs = w2v_kwargs if w2v_kwargs is not None else {}\n",
    "\n",
    "        t0 = time.time()\n",
    "        # w2v.init_sims(replace=True)\n",
    "        logger.debug('Load w2v from {}, spend {:.2f} sec'.format(model_name_or_path, time.time() - t0))\n",
    "        self.stopwords = stopwords if stopwords else load_stopwords(default_stopwords_file)\n",
    "        self.model = Word2Vec.load(model_name_or_path)\n",
    "        self.w2v = self.model.wv\n",
    "        self.jieba = jieba\n",
    "        self.model_name_or_path = model_name_or_path\n",
    "\n",
    "    def __str__(self):\n",
    "        return f\"<Word2Vec, word count: {len(self.w2v.key_to_index)}, emb size: {self.w2v.vector_size}, \" \\\n",
    "               f\"stopwords count: {len(self.stopwords)}>\"\n",
    "\n",
    "    def encode(self, sentences: Union[List[str], str], show_progress_bar: bool = False) -> ndarray:\n",
    "        \"\"\"\n",
    "        Encode sentences to vectors\n",
    "        \"\"\"\n",
    "        if self.w2v is None:\n",
    "            raise ValueError('No model for embed sentence')\n",
    "\n",
    "        input_is_string = False\n",
    "        if isinstance(sentences, str) or not hasattr(sentences, '__len__'):\n",
    "            sentences = [sentences]\n",
    "            input_is_string = True\n",
    "\n",
    "        all_embeddings = []\n",
    "        for sentence in tqdm(sentences, desc='Word2Vec Embeddings', disable=not show_progress_bar):\n",
    "            emb = []\n",
    "            count = 0\n",
    "            for word in sentence:\n",
    "                # 过滤停用词\n",
    "                if word in self.stopwords:\n",
    "                    continue\n",
    "                # 调用词向量\n",
    "                if word in self.w2v.key_to_index:\n",
    "                    emb.append(self.w2v.get_vector(word, norm=True))\n",
    "                    count += 1\n",
    "                else:\n",
    "                    if len(word) == 1:\n",
    "                        continue\n",
    "                    # 再切分\n",
    "                    ws = self.jieba.lcut(word, cut_all=True, HMM=True)\n",
    "                    for w in ws:\n",
    "                        if w in self.w2v.key_to_index:\n",
    "                            emb.append(self.w2v.get_vector(w, norm=True))\n",
    "                            count += 1\n",
    "            tensor_x = np.array(emb).sum(axis=0)  # 纵轴相加\n",
    "            if count > 0:\n",
    "                avg_tensor_x = np.divide(tensor_x, count)\n",
    "            else:\n",
    "                avg_tensor_x = np.zeros(self.w2v.vector_size, dtype=float)\n",
    "            all_embeddings.append(avg_tensor_x)\n",
    "        all_embeddings = np.array(all_embeddings, dtype=float)\n",
    "        if input_is_string:\n",
    "            all_embeddings = all_embeddings[0]\n",
    "\n",
    "        return all_embeddings\n",
    "\n",
    "stop_word_set = load_stopwords(stopwords_table_path)              \n",
    "model = Word2VecManager(stopwords=stop_word_set)\n",
    "model.encode(\"你好我是中国人\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "402afc46",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-01-16 15:23:11.756 | DEBUG    | __main__:__init__:42 - Load w2v from ./word2vec.model, spend 0.00 sec\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.791279037506656\n",
      "0.27267503283759903\n"
     ]
    }
   ],
   "source": [
    "import numpy\n",
    "\n",
    "stop_word_set = load_stopwords(stopwords_table_path)              \n",
    "model = Word2VecManager(stopwords=stop_word_set)\n",
    "\n",
    "sentence1 = \"我喜欢小狗。\"\n",
    "sentence2 = \"我喜欢小动物。\"\n",
    "sentence3 = \"我今天心情很差。\"\n",
    "\n",
    "embedding1 = model.encode(sentence1)\n",
    "embedding2 = model.encode(sentence2)\n",
    "embedding3 = model.encode(sentence3)\n",
    "\n",
    "\n",
    "def cosine_similarity(vec1, vec2):\n",
    "    norm_vec1 = numpy.linalg.norm(vec1)\n",
    "    norm_vec2 = numpy.linalg.norm(vec2)\n",
    "    if norm_vec1 == 0 or norm_vec2 == 0:\n",
    "        return 0\n",
    "    else:\n",
    "        dot_product = numpy.dot(vec1, vec2)\n",
    "        return dot_product / (norm_vec1 * norm_vec2)\n",
    "\n",
    "print(cosine_similarity(embedding1, embedding2))\n",
    "print(cosine_similarity(embedding1, embedding3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "2d3c2586",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['欢迎使用结巴中文分词', '请问有什么可以帮助您的吗']"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def split_sentences(text):\n",
    "    sent_delimiters = ['。', '？', '！', '?', '!', '.']\n",
    "    for delimiter in sent_delimiters:\n",
    "        text = text.replace(delimiter, '\\n')\n",
    "    sentences = text.split('\\n')\n",
    "    sentences = [sent for sent in sentences if sent.strip()]\n",
    "    return sentences\n",
    "\n",
    "text = '欢迎使用结巴中文分词！请问有什么可以帮助您的吗？'\n",
    "split_sentences(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "179bc343",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "问题: 《铃芽户缔》剧情有几条线？ \n",
      "最优前5条:\n",
      "《铃芽户缔》有三条线，分别是：1、铃芽自我治愈的成长之旅 2、铃芽和草太的爱情故事 3、铃芽和草太关闭“往门”的旅途故事[1] (0.75)\n",
      "如果要让我一句话介绍《铃芽户缔》再讲什么，那我一定会引用主线3 (0.61)\n",
      "在剧中所有人看来这是铃芽青春期拒绝沟通、离家出走的叛逆 (0.60)\n",
      "这是事实，但在《铃芽户缔》中，虽然是第一次尝试公路片，我认为新海诚对故事节奏的把握进步了很多 (0.59)\n",
      "《铃芽户缔》并不是一部以恋爱为主题的电影，或者说恋爱并不是位于第一位的线 (0.59)\n"
     ]
    }
   ],
   "source": [
    "# 以下是系统的知识库\n",
    "def create_corpus(doc):\n",
    "    sents = split_sentences(doc)\n",
    "    corpus = []\n",
    "    corpus_embeddings = []\n",
    "    for sent in sents:\n",
    "        vc = model.encode(sent)\n",
    "        corpus.append(sent)\n",
    "        corpus_embeddings.append(vc)\n",
    "    return corpus, corpus_embeddings\n",
    "\n",
    "def create_corpus_from_file(filepath):\n",
    "    with open(filepath, 'r', encoding='utf-8') as file:\n",
    "        doc = file.read()\n",
    "    return create_corpus(doc)\n",
    "\n",
    "def semantic_search(query_embedding, corpus_embeddings, top_k=5):\n",
    "    similarities = [cosine_similarity(query_embedding, doc_embedding) for doc_embedding in corpus_embeddings]\n",
    "    sorted_indexes = np.argsort(similarities)[::-1]\n",
    "    top_k_indexes = sorted_indexes[:top_k]\n",
    "    return top_k_indexes, np.array(similarities)[top_k_indexes]\n",
    "    \n",
    "\n",
    "# 以下是用户的问题\n",
    "queries = [\"《铃芽户缔》剧情有几条线？\"]\n",
    "corpus, corpus_embeddings = create_corpus_from_file(\"./data/suzume.md\")\n",
    "for query in queries:\n",
    "    query_embedding = model.encode(query)\n",
    "    # 将问题通过模型从知识库匹配，取前3条\n",
    "    hit_indexes, hit_similarities = semantic_search(query_embedding, corpus_embeddings, top_k=5)\n",
    "    print(\"\\n问题:\", query, \"\\n最优前5条:\")\n",
    "    for hit_index, hit_similarity in zip(hit_indexes, hit_similarities):\n",
    "        print(corpus[hit_index], \"({:.2f})\".format(hit_similarity))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.7"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}