对句子分词，找到对应词的腾讯词向量模型并使用Python进行faiss检索

zkkkkkkkkkkkkk

已于 2022-07-04 13:53:39 修改

阅读量1.7k

点赞数 2

分类专栏： python faiss检索

于 2022-06-28 14:38:19 首次发布

本文链接：https://blog.csdn.net/zkkkkkkkkkkkkk/article/details/125501250

版权

python 全文检索

python 同时被 2 个专栏收录

61 篇文章 5 订阅

订阅专栏

faiss检索

1 篇文章 0 订阅

订阅专栏

本文主要是将句子分词转向量，再加总词向量求平均变为句子向量。接着再存储到faiss中。等待新句子到来，同样按照上述方法处理。达到在faiss能检索出相似的向量。

一、下载腾讯的词向量

下载后放到一个地方，待会代码部分需要使用。下载地址： Embedding Dataset -- NLP Center, Tencent AI Lab

二、停用词

可以上网查找一些停用词表，或者自己定义一个stop_words.txt。同样代码部分需要使用。

三、代码部分

3.1、代码思想

1）首先加载腾讯的词向量模型

2）把要初始化的句子和要检索的句子

3）对句子分词，找出词在腾讯词向量模型里面的向量

4）每个句子的所有词向量加总取平均做为句子向量

5）存入faiss，等待检索

6）输入待检索句子向量，输出相似句子向量

from gensim.models import Word2Vec
import os
import jieba.posseg as pseg
import faiss
from gensim.models import KeyedVectors
import pandas as pd
import numpy as np


# 获取腾讯词向量
file = r'D:\Gitlab\extract_key\baidu_word_model\tencent-ailab-embedding-zh-d100-v0.2.0_small.txt'
wv_from_text = KeyedVectors.load_word2vec_format(file, binary=False)  # 加载时间比较长
wv_from_text.init_sims(replace=True)



def jieba_postag(text):
    words = pseg.cut(text)
    return words


def defined_stop_words():
    all_stop_words = []
    for i, file in enumerate(os.listdir(r'D:\Gitlab\extract_key\stop_words')):
        # 读取图片
        filepath = fr'D:\Gitlab\extract_key\stop_words\{file}'
        with open(filepath, 'r', encoding='utf-8') as fp:
            all_line = fp.readlines()
            for line in all_line:
                all_stop_words.append(line.replace('\n',''))
    return all_stop_words


def load_baidu_model(sentence_list):
    sen_list = []
    vector_list = []
    # 外层列表 多个句子的分词结果
    for sentence in sentence_list:
        count = 0
        avg_list = []
        # 内层，每个句子的分词结果
        for word in sentence:
            # 找出对应词的百度词向量
            if word in wv_from_text.index_to_key:
                count += 1
                # 取出对应词向量
                vector = wv_from_text[word]
                # 存储起来
                avg_list.append(vector)

        # 一句话结束后，把所有词向量相加求平均
        sen_vector = np.vstack(avg_list) / count
        # 添加返回信息
        sen_list.append(sentence)
        vector_list.append(sen_vector)

    #        句子     对应句子向量
    return sen_list,vector_list



def word2vec(content):
    # 获取停用词
    stop_words = defined_stop_words()

    # jieba分词
    sentence = []
    for con in content:
        sen = []
        word = jieba_postag(con)
        for w in word:
            # 去停用词
            if w.word not in stop_words:
                sen.append(w.word)
        sentence.append(sen)
    # jieba 分词结果
    print(sentence)


    # 句子列表  句子向量
    sen_list,vector_list = load_baidu_model(sentence)


    return vector_list


def faiss_calculate(word_vec,test_word_vec,near_size):
    vec_shape = 100
    index = faiss.IndexFlatL2(vec_shape)

    # # 有多少条数据训练
    # print(index.ntotal)
    # # 是否已经训练
    # print(index.is_trained)

    # 添加向量
    for i,vec in enumerate(word_vec):
        print(f"添加第{i}个向量")
        index.add(vec)
    # print(index.ntotal)
    print("向量添加成功\n")

    # search vectors
    k = near_size
    for i,test_vec in enumerate(test_word_vec):
        print(f"检索第{i}个向量")
        I,D = index.search(test_vec,k)
        print(I)
        print(D)




if __name__ == '__main__':
    content = ["长江是中国第一大河，干流全长6397公里（以沱沱河为源），一般称其6300公里。流域总面积一百八十余万平方公里，年平均入海量约九千六百余亿立方米。以干流长度和入海水量论，长江居世界第三位。",
               "长江干流宜昌以上为上游，长4504公里，流域面积100万平方公里，其中直门达至宜宾称金沙江，长3364公里；宜宾至宜昌河段习称川江，长1040公里；宜昌至湖口为中游，长955公里，流域面积68万平方公里。湖口以下为下游，长938公里，流域面积12万平方公里。"]
    search_content = ["长江干流宜昌以上为上游，长4504公里，流域面积100万平方公里，其中直门达至宜宾称金沙江，长3364公里；宜宾至宜昌河段习称川江，长1040公里；宜昌至湖口为中游，长955公里，流域面积68万平方公里。湖口以下为下游，长938公里，流域面积12万平方公里。"]
    word_vec = word2vec(content)
    test_word_vec = word2vec(search_content)
    faiss_calculate(word_vec,test_word_vec,5)

四、输出结果

前面两个列表分别是content和search_content分词后的结果。往下看就是faiss的代码部分。最后的是向量和对应下标内容。

C:\Users\dell\Anaconda3\python.exe D:/Gitlab/extract_key/faiss_demo.py
Building prefix dict from the default dictionary ...
Loading model from cache C:\Users\dell\AppData\Local\Temp\jieba.cache
[['长江', '中国', '第一', '大河', '干流', '全长', '6397', '公里', '沱沱河', '源', '称', '6300', '公里', '流域', '总面积', '一百八十', '余万平方公里', '年', '平均', '入海量', '约', '九千', '六百余', '亿立方米', '干流', '长度', '入海', '水量', '长江', '居', '世界', '第三位'], ['长江', '干流', '宜昌', '上游', '长', '4504', '公里', '流域面积', '100', '万平方公里', '直门', '达至', '宜宾', '称', '金沙江', '长', '3364', '公里', '宜宾', '宜昌', '河段', '习称', '川江', '长', '1040', '公里', '宜昌', '湖口', '中游', '长', '955', '公里', '流域面积', '68', '万平方公里', '湖口', '下游', '长', '938', '公里', '流域面积', '12', '万平方公里']]
Loading model cost 0.467 seconds.
Prefix dict has been built successfully.
[['长江', '干流', '宜昌', '上游', '长', '4504', '公里', '流域面积', '100', '万平方公里', '直门', '达至', '宜宾', '称', '金沙江', '长', '3364', '公里', '宜宾', '宜昌', '河段', '习称', '川江', '长', '1040', '公里', '宜昌', '湖口', '中游', '长', '955', '公里', '流域面积', '68', '万平方公里', '湖口', '下游', '长', '938', '公里', '流域面积', '12', '万平方公里']]
添加第0个向量
添加第1个向量
向量添加成功

检索第0个向量
[[2.3283064e-10 7.9978607e-05 7.9978607e-05 2.4693029e-04 3.2110960e-04]
 [1.1641532e-10 7.9978839e-05 7.9978839e-05 1.4126929e-04 2.5593187e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 1.6702199e-04 1.6702199e-04]
 [0.0000000e+00 5.6024292e-05 2.0133099e-04 3.0256179e-04 3.0258833e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
 [0.0000000e+00 2.0129327e-04 3.3610640e-04 4.1364867e-04 4.2342267e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 1.1641532e-10]
 [1.1641532e-10 1.1641532e-10 3.4924597e-10 2.6041776e-04 2.9876560e-04]
 [0.0000000e+00 2.3221667e-04 2.6178535e-04 5.0974241e-04 5.1002507e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 2.9876560e-04 2.9876560e-04]
 [1.1641532e-10 1.1641532e-10 1.6702199e-04 1.6702199e-04 1.6702199e-04]
 [1.1641532e-10 7.9978723e-05 3.5184540e-04 5.9754209e-04 5.9754209e-04]
 [0.0000000e+00 2.5593187e-04 3.1231623e-04 3.4288963e-04 3.6646723e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
 [0.0000000e+00 2.0129327e-04 3.5257672e-04 4.4940610e-04 4.7621201e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 1.1641532e-10]
 [1.1641532e-10 1.1641532e-10 1.6702199e-04 1.6702199e-04 1.6702199e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 1.6702199e-04 1.6702199e-04]
 [0.0000000e+00 1.4126929e-04 2.7304713e-04 2.7304713e-04 3.0256179e-04]
 [1.1641532e-10 3.5184540e-04 5.6083419e-04 7.7203778e-04 7.8827370e-04]
 [0.0000000e+00 2.4693029e-04 3.1231623e-04 4.1363190e-04 4.1363190e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
 [0.0000000e+00 2.9964186e-04 3.3610640e-04 3.4230750e-04 3.5257672e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 1.1641532e-10]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 1.6702199e-04 1.6702199e-04]
 [0.0000000e+00 0.0000000e+00 3.7042430e-04 3.9958593e-04 3.9958593e-04]
 [1.1641532e-10 2.0133099e-04 2.1574745e-04 4.6068663e-04 4.7935866e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
 [1.1641532e-10 2.0490785e-04 2.9964186e-04 4.1288522e-04 4.1364867e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 1.1641532e-10]
 [1.1641532e-10 1.1641532e-10 3.4924597e-10 2.6041776e-04 2.9876560e-04]
 [0.0000000e+00 2.3221667e-04 2.4346728e-04 4.1288522e-04 4.3369649e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 2.9876560e-04 2.9876560e-04]
 [0.0000000e+00 0.0000000e+00 3.7042430e-04 3.9958593e-04 3.9958593e-04]
 [0.0000000e+00 5.6024292e-05 2.1574751e-04 3.3985404e-04 3.4211029e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00]
 [2.3283064e-10 2.0490785e-04 3.4230767e-04 4.2342278e-04 4.3369655e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 0.0000000e+00 1.1641532e-10]
 [1.1641532e-10 1.1641532e-10 3.4924597e-10 2.6041776e-04 2.9876560e-04]
 [0.0000000e+00 2.4346705e-04 2.6178546e-04 5.1894749e-04 5.1894749e-04]
 [0.0000000e+00 0.0000000e+00 0.0000000e+00 2.9876566e-04 2.9876566e-04]]
[[30 26  0 50 33]
 [31  4 22 48 42]
 [47 32 54 40 46]
 [33 64 56 48 31]
 [34 51 57 65 43]
 [35 44 52 58 66]
 [45 36 59 53 67]
 [60 37 68 31 62]
 [38 61 69 52 58]
 [62 70 39 60 37]
 [46 40 47 54 32]
 [41 10 49 43 34]
 [42 31 50 48 30]
 [34 51 57 65 43]
 [44 35 52 66 58]
 [45 36 59 53 67]
 [46 40 47 54 32]
 [47 32 54 40 46]
 [48 31  4 22 33]
 [49 41 10 64 30]
 [50 30 42 63 55]
 [34 51 57 65 43]
 [52 58 35 66 44]
 [45 36 59 53 67]
 [47 32 54 40 46]
 [63 55 30 54 47]
 [56 33 64 30 31]
 [34 51 57 65 43]
 [58 66 52 61 35]
 [45 36 59 53 67]
 [60 37 68 31 62]
 [61 38 69 58 66]
 [62 70 39 60 37]
 [63 55 30 54 47]
 [64 33 56 48 31]
 [34 51 57 65 43]
 [66 58 52 35 61]
 [45 36 59 53 67]
 [60 37 68 31 62]
 [69 61 38 43 57]
 [62 70 39 68 60]]

Process finished with exit code 0