文档进行embedding,Faiss向量检索

这里采用Langchain的HuggingFaceEmbeddings

参照博主,改了一些东西,因为Langchain0.3在0.2的基础上进行了一定的修改

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.vectorstores import Chroma

filepath = 'files/txt/res.txt'
raw_documents = TextLoader(filepath, encoding='utf8').load()

# 按行分割块
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=100,
    chunk_overlap=20,
    separators=[  # 分隔符字符串数组
        "\n\n",
        "\n",
        " ",
        # 因为这里是中文文档,加上这个会影响小数点".",
        ",",
        "\u200B",  # Zero-width space
        "\uff0c",  # Fullwidth comma
        "\u3001",  # Ideographic comma
        "\uff0e",  # Fullwidth full stop
        "\u3002",  # Ideographic full stop
        ""
    ],
    length_function=len,
    is_separator_regex=True,
)
documents = text_splitter.split_documents(raw_documents)
print(len(documents))

# 词嵌入模型
EMBEDDING_DEVICE = "cpu"

# 加载本地 embedding 模型
embedding = HuggingFaceEmbeddings(model_name='m3e-base', model_kwargs={'device': EMBEDDING_DEVICE})
# 创建向量数据库
db = FAISS.from_documents(documents, embedding)
# 保存
db.save_local("./faiss_index")
'''
如果已经创建好了,可以直接读取
db = FAISS.load_local("./faiss_index", embeddings)
'''

# 直接传入文本
query = "嘻嘻嘻嘻嘻嘻嘻嘻嘻"
docs = db.similarity_search(query, k=3)
# docs = db.similarity_search_with_score(query, k=3)  # 带分数的
print(docs[0].page_content)

print("分割线~~~~~~~~~~~~~~~~~~··")
# 传入向量去搜索
embedding_vector = embedding.embed_query(query)
docs = db.similarity_search_by_vector(embedding_vector, k=3)
print(docs[0].page_content)

接下来是运行所需的pip ,我这里Embedding_model选的cpu,因为还没配cuda,且这个规模的文档,笔记本cpu也支持。

首先,git自己需要的Embedding模型,我这里用的比较熟悉的m3e-base,采用了魔搭的镜像

m3e-base: git clone https://www.modelscope.cn/Jerry0/m3e-base.git

pip install sentence_transformers
pip install -U langchain-huggingface   #因为langchain升级引起的更新
 conda install -c pytorch faiss-cpu    #这里有个坑,用这个命令解决的(同样用CPU的可以用这个,跟我对齐)
 

你可能感兴趣的:(embedding,faiss)