langchain实现文档向量检索

symptoms.txt数据如下：

09 视觉系统疾病 > 眼附属器或眼眶疾患 > 眼睑或眼周区疾患 > 9A03 获得性眼睑错位 > 9A03.0 上睑下垂 > 9A03.00 下颌瞬目综合征::::上睑下垂|患侧上睑与下颌的联带运动|吸吮上睑提高|闭口时患侧上睑回到原来位置|张口上睑提高|咀嚼上睑提高|下颌朝向对侧方向运动时上睑提高
01 某些感染性疾病或寄生虫病 > 分枝杆菌病 > 结核病 > 1B10 呼吸系统结核::::喘息症状|胸痛|咳嗽|咳痰|痰中带血|咯血|呼吸困难|发热
LA14.04 先天性上睑下垂::::上睑抬起困难|下颌上抬|患侧眉毛高竖|屈光不正|弱视|斜视|额部皮肤有明显的横行皱纹
09 视觉系统疾病 > 眼附属器或眼眶疾患 > 眼睑或眼周区疾患 > 9A03 获得性眼睑错位 > 9A03.0 上睑下垂 > 9A03.01 机械性睑下垂::::上睑遮盖部分角膜|沙眼性睑板浸润|外伤性瘢痕增厚|睑神经纤维瘤|血管瘤|上睑遮盖全部角膜
09 视觉系统疾病 > 眼附属器或眼眶疾患 > 眼睑或眼周区疾患 > 9A03 获得性眼睑错位 > 9A03.0 上睑下垂 > 9A03.03 麻痹性上睑下垂::::瞳孔扩大|眼球向内运动受限|向上运动受限|向下运动受限|瞳孔固定|单眼上睑下垂|眼球运动受限|复视
09 视觉系统疾病 > 眼附属器或眼眶疾患 > 眼睑或眼周区疾患 > 9A03 获得性眼睑错位 > 9A03.1 睑内翻::::畏光|流泪|疼痛|异物感|视力下降
09 视觉系统疾病 > 眼附属器或眼眶疾患 > 眼睑或眼周区疾患 > 9A03 获得性眼睑错位 > 9A03.1 睑内翻 > 9A03.10 瘢痕性睑内翻::::畏光|流泪|疼痛|异物感|视力下降
01 某些感染性疾病或寄生虫病 > 分枝杆菌病 > 结核病 > 1B12 其他系统和器官的结核 > 1B12.6 周围淋巴结结核病::::低热|眼痒|畏光|红肿|流泪|红色或者紫红色结节|盗汗|皮肤苍白|消瘦|全身无力
01 某些感染性疾病或寄生虫病 > 分枝杆菌病 > 结核病 > 1B11 神经系统结核 > 1B11.0 结核性脑脊膜炎::::低热|高热|食欲缺乏|恶心|头晕|头痛|咳嗽|呕吐|咳痰|便秘|易激动|畏光|盗汗|乏力
01 某些感染性疾病或寄生虫病 > 分枝杆菌病 > 1B21 非结核分枝杆菌所致感染 > 1B21.2 皮肤的非结核分枝杆菌感染 > 1B21.20 溃疡分枝杆菌感染::::结节|水疱|黄色坏死物|瘢痕|坏死性溃疡|典型的 Buruli 溃疡以慢性，坏死性，进行性，无痛性，无炎症性，潜行性，难治性为其临床特征|溃疡愈合时，常有瘢痕组织挛缩，导致肢体畸形，活动障碍

代码如下：

# -*- coding: utf-8 -*-

from langchain.embeddings import ModelScopeEmbeddings
from langchain.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma


def split_document(filepath='./symptoms.txt'):
    # Load the document, split it into chunks
    raw_documents = TextLoader(filepath, encoding='utf-8').load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1, chunk_overlap=0, separators=["\n"])
    documents = text_splitter.split_documents(raw_documents)
    return documents


def load_embedding_function():
    # create the open-source embedding function
    embedding_function = ModelScopeEmbeddings(model_id="damo/nlp_corom_sentence-embedding_chinese-base-medical")
    return embedding_function


# save chroma db of embeds to disk
def save_document_embedding_db(documents, embedding_function, db_filepath='./chroma_db'):
    chroma_db = Chroma.from_documents(documents, embedding_function, persist_directory=db_filepath)
    return chroma_db


# load chroma db of embeds from disk
def load_document_embedding_db(db_filepath='./chroma_db'):
    # load from disk
    chroma_db = Chroma(persist_directory=db_filepath, embedding_function=embedding_function)
    return chroma_db


def document_query(query, embed_db, k=10):
    results = embed_db.similarity_search(query, k=k)
    return results


if __name__ == '__main__':
    document_filepath = './disease_search/symptoms.txt'
    db_filepath = './disease_search/chroma_db'

    documents = split_document(document_filepath)
    embedding_function = load_embedding_function()
    # save_document_embedding_db(documents, embedding_function, db_filepath='./chroma_db')
    embed_db = load_document_embedding_db(db_filepath)
    print("向量检索库构造完成！开始进行文本相似度检索...")

    query = "患者，男，26岁，湖南籍防空兵中士，因“双侧腋窝多汗3年，发现左侧腋毛起皮疹2天”就诊。患者3年前无诱因双侧腋窝出现多汗，训练及加班后加重。近1年因工作压力大、加班多自觉出汗增多明显，偶可沿腋毛滴下，个别时有轻微异味，2天前体检时发现左侧腋窝腋毛干上附着白色细小结节，特来就诊。患者既往体健，无系统性疾病史，无药物过敏史，有掌跖部及会阴部多汗史，军队计划免疫。患者所在单位既往曾有“腋毛癣”患者。体格检查：一般情况好，心、肺及腹部查体未见异常，全身淋巴结未触及肿大。皮肤科情况：双侧腋窝腋毛浓密，局部多汗，左侧腋窝部分腋毛上附着冰凌状白色小结节，与毛干紧密黏着，无自觉症状。相邻淋巴结无肿大。实验室及辅助检查：血、尿、便常规、胸片未见异常。心电图示窦性心律，正常心电图。"
    query_results = document_query(query, embed_db, k=10)

    for result in query_results:
        print(result.page_content)
        print('-' * 30)