Elasticsearch: 基于Text Embedding的文本相似性搜索

xxxxx一行医学相关的文本，百姓所说的感冒是指“普通感冒”，又称“伤风”、急性鼻炎或上呼吸道感染。感冒是一种常见的急性上呼吸道病毒性感染性疾病，多由鼻病毒、副流感病毒、呼吸道合胞病毒、埃可病毒、柯萨奇病毒、冠状病毒、腺病毒等引起。临床表现为鼻塞、喷嚏、流涕、发热、咳嗽、头痛等，多呈自限性。大多散发，冬春季节多发，但不会出现大流行。 
一行医学相关的文本xxxxx 
一行医学相关的文本xxxxx

[
    {
        "id": "1",
        "name": "肝功能异常",
        "department": "消化科",
        "feature": "消化功能xxxxxxxxxxxxxxxxxxxxxxx\n"
    },
    {
        "id": "2",
        "name": "反胃",
        "department": "消化科",
        "feature": "xxxxxxxxxxxxxx为主要表现。\n"
    }
]

两者 
个
个别
临
为
为了
为什么

# 停用词
stopwords = [line.strip() for line in open('./data/ChineseStopWords.txt', encoding='UTF-8').readlines()]


def segment(sentence: str):
    """
    结巴分词，并去除停用词
    """
    resp = []
    sentence_depart = jieba.cut(sentence.strip())
    for word in sentence_depart:
        if word not in stopwords:
            if word != "":
                resp.append(word)
    return resp


def read_corpus(f_name):
    """
    读数据
    """
    with open(f_name, encoding="utf-8") as f:
        for i, line in enumerate(f):
            yield gensim.models.doc2vec.TaggedDocument(segment(line), [i])


def train():
    """
    训练 Doc2Vec 模型
    """
    train_file = "./data/train_data.txt"
    train_corpus = list(read_corpus(train_file))
    model = gensim.models.doc2vec.Doc2Vec(vector_size=300, min_count=2, epochs=10)
    print(len(train_corpus))
    model.build_vocab(train_corpus)
    model.train(train_corpus, total_examples=model.corpus_count, epochs=model.epochs)
    model.save("doc2vec.model")

def create_index():
    print("begin create index")
    setting = {
        "settings": {
            "number_of_replicas": 0,
            "number_of_shards": 2
        },
        "mappings": {
            "properties": {
                "name": {
                    "type": "keyword"
                },
                "department": {
                    "type": "keyword"
                },
                "feature": {
                    "type": "text"
                },
                "feature_vector": {
                    "type": "dense_vector",
                    "dims": 300
                }
            }
        }
    }
    get_es_client().indices.create(index=indexName, body=setting)
    print("end create index")

def embed_text(sentences):
    """
    将所有的句子转化为向量
    """
    model = doc2vec.Doc2Vec.load("doc2vec.model")
    resp = []
    for s in sentences:
        resp.append(model.infer_vector(segment(s)).tolist())
    return resp

def bulk_index_data():
    """
    将数据索引到es中，且其中包含描述的特征向量字段
    """
    print("begin embed index data to vector")
    with open("./data/data.json") as file:
        load_dict = json.load(file)
    features = [doc["feature"] for doc in load_dict]
    print("number of lines to embed:", len(features))
    features_vectors = embed_text(features)
    print("begin index data to es")
    requests = []
    for i, doc in enumerate(load_dict):
        request = {'_op_type': 'index',  # 操作 index update create delete  
                   '_index': indexName,  # index
                   '_id': doc["id"],
                   '_source':
                       {
                           'name': doc["name"],
                           'department': doc["department"],
                           'feature': doc["feature"],
                           'feature_vector': features_vectors[i],
                       }
                   }
        requests.append(request)
    bulk(get_es_client(), requests)
    print("end index data to es")

def test():
    model = doc2vec.Doc2Vec.load("doc2vec.model")
    es = get_es_client()
    while True:
        try:
            query = input("Enter query: ")
            input_vector = model.infer_vector(segment(query)).tolist()
            resp = es.search(index=indexName, body={
                "_source": ["name", "feature"],
                "query": {
                    "script_score": {
                        "query": {
                            "match_all": {}
                        },
                        "script": {
                            "source": "cosineSimilarity(params.queryVector, doc['feature_vector'])+1",
                            "params": {
                                "queryVector": input_vector
                            }
                        }
                    }
                }
            })
            print("可能获得的疾病是：", end=" ")
            for hit in resp["hits"]["hits"]:
                print(hit["_source"]["name"], end="\t")
            print("\n")
        except KeyboardInterrupt:
            return

Enter query: 我眼睛充血，怎么办？
可能获得的疾病是： 红眼病 眼角膜发炎 外伤


Enter query: 呼吸不畅，咳嗽，胸闷是怎么回事？
可能获得的疾病是： 肺炎  上呼吸道感染   支气管炎