Langchain系列[16]下载和使用Rerank模型

git clone https://www.modelscope.cn/maidalun/bce-embedding-base_v1.git

git clone https://www.modelscope.cn/maidalun/bce-reranker-base_v1.git

pip install BCEmbedding

[1] 我最喜欢的电影是《阿甘正传》。        Score:0.6105921907574262
[2] 我最喜欢的电影类型是恐怖片。          Score:0.5381861816884799
[3] 我最喜欢的电影导演是斯皮尔伯格。       Score:0.4905046919817628
[4] 我最喜欢的书是《哈利波特》。          Score:0.4652198126696215
[5] 我不喜欢看电影。                     Score:0.44154025662995605
[6] 我喜欢看科幻电影，尤其是《星际穿越》。 Score:0.4249389530579074
[7] 我最喜欢的音乐是爵士乐。               Score:0.39165986473334957
[8] 昨天晚上我看了一部很好的电影。         Score:0.35591581297344105
[9] 我最喜欢的运动是篮球。                 Score:0.3554765233404198
[10] 我喜欢吃披萨。                       Score:0.34989368275228544

我最喜欢的电影是《阿甘正传》。 [1]
我最喜欢的电影类型是恐怖片 [2]
我喜欢看科幻电影，尤其是《星际穿越》。 [6]
我最喜欢的电影导演是斯皮尔伯格。 [3]
昨天晚上我看了一部很好的电影。 [8]
我最喜欢的书是《哈利波特》。 [4]
我不喜欢看电影。 [5]
我最喜欢的音乐是爵士乐。 [7]
我喜欢吃披萨。 [10]
我最喜欢的运动是篮球。 [9]

# 初始化本地的 embedding 模块
embedding_model_name = 'D:\LLM\\bce_modesl\\bce-embedding-base_v1'
# 使用 GPU
# embedding_model_kwargs = {'device': 'cuda:0'}

# 使用 CPU
embedding_model_kwargs = {'device': 'cpu'}

embedding_encode_kwargs = {'batch_size': 32, 'normalize_embeddings': True, }

# 初始化
embed_model = HuggingFaceEmbeddings(
    model_name=embedding_model_name,
    model_kwargs=embedding_model_kwargs,
    encode_kwargs=embedding_encode_kwargs
)
# 小测试
query = '你好呀兄弟！'
passages = ['天真热', '嗯呢']
query_embedding = embed_model.embed_query(query)
passages_embeddings = embed_model.embed_documents(passages)

query = '你最喜欢的电影是什么？'
passages = [
    '我最喜欢的电影是《阿甘正传》。',
    '我最喜欢的电影类型是恐怖片',
    '我最喜欢的电影导演是斯皮尔伯格。',
    '我最喜欢的书是《哈利波特》。',
    '我不喜欢看电影。',
    '我喜欢看科幻电影，尤其是《星际穿越》。',
    '我最喜欢的音乐是爵士乐。',
    '昨天晚上我看了一部很好的电影。',
    '我最喜欢的运动是篮球。',
    '我喜欢吃披萨。'
]

# construct sentence pairs
sentence_pairs = [[query, passage] for passage in passages]

# rerank 模型所在目录
reranker_model = 'D:\LLM\\bce_modesl\\bce-reranker-base_v1'

# 初始化
model = RerankerModel(model_name_or_path=reranker_model)

# 得到模型评估的分数
scores = model.compute_score(sentence_pairs)

# rerank 过滤
rerank_results = model.rerank(query, passages)
doc_list = rerank_results['rerank_passages']
for line in doc_list:
    print(line)
pass

# 定义 reranker
reranker_args = {'model': 'D:\LLM\\bce_modesl\\bce-reranker-base_v1', 'top_n': 10, 'device': 'cuda:0'}
reranker = BCERerank(**reranker_args)

# 生成一个langchain 的语境压缩检索器
compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker, base_retriever=retriever
)
# 搜索
docs = compression_retriever.get_relevant_documents('你最喜欢的电影是什么?')

# ------- 一条龙处理 

   # 文档分割器
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

    # 向量过滤
    embedding_filter = EmbeddingsFilter(embeddings=embed_model, similarity_threshold=0.6)

    # 文档去重
    redundant_filter = EmbeddingsRedundantFilter(embeddings=embed_model)

    # 内容抽取
    llm_extractor = LLMChainExtractor.from_llm(qianfan_chat)

    pipeline_compressor = DocumentCompressorPipeline(
        transformers=[text_splitter, embedding_filter, redundant_filter, reranker, llm_extractor]
    )
    pipeline_retriever = ContextualCompressionRetriever(
        base_compressor=pipeline_compressor, base_retriever=retriever
    )
    docs = pipeline_retriever.get_relevant_documents(query)

import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.chat_models import QianfanChatEndpoint
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain_community.embeddings import QianfanEmbeddingsEndpoint, HuggingFaceEmbeddings
from langchain import hub
from langchain_community.vectorstores.elasticsearch import ElasticsearchStore
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

import numpy as np


def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    return dot_product / (norm_vec1 * norm_vec2)


if __name__ == '__main__':

    # 使用本地网易有道模型  BCE
    embedding_model_name = 'D:\LLM\\bce_modesl\\bce-embedding-base_v1'
    embedding_model_kwargs = {'device': 'cuda:0'}
    embedding_encode_kwargs = {'batch_size': 32, 'normalize_embeddings': True, }

    embed_model = HuggingFaceEmbeddings(
        model_name=embedding_model_name,
        model_kwargs=embedding_model_kwargs,
        encode_kwargs=embedding_encode_kwargs
    )

    doc_list = [
        '我最喜欢的电影是《阿甘正传》。',
        '我最喜欢的电影类型是恐怖片',
        '我最喜欢的电影导演是斯皮尔伯格。',
        '我最喜欢的书是《哈利波特》。',
        '我不喜欢看电影。',
        '我喜欢看科幻电影，尤其是《星际穿越》。',
        '我最喜欢的音乐是爵士乐。',
        '昨天晚上我看了一部很好的电影。',
        '我最喜欢的运动是篮球。',
        '我喜欢吃披萨。'
    ]
    docs = []
    index =1
    for line in doc_list:
        tmp_doc = Document(page_content=line)
        tmp_doc.metadata['num'] = index
        index += 1
        docs.append(tmp_doc)
    pass

    vectorstore = ElasticsearchStore(
        es_url=os.environ['ELASTIC_HOST_HTTP'],
        index_name="index_ex_768_vectors",
        embedding=embed_model,
        es_user="elastic",
        vector_query_field='question_vectors',
        es_password=os.environ['ELASTIC_ACCESS_PASSWORD']
    )
    # 将本文加入数据库（注意，只需要运行一次；多次运行会有冗余数据）
    vectorstore.add_documents(docs)
    print('Done!')

import os
from uuid import uuid4

from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import DocumentCompressorPipeline, EmbeddingsFilter, LLMChainExtractor
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.chat_models import QianfanChatEndpoint
from langchain_community.document_loaders import TextLoader, DirectoryLoader
from langchain_community.document_transformers import EmbeddingsRedundantFilter
from langchain_community.embeddings import QianfanEmbeddingsEndpoint, HuggingFaceEmbeddings
from langchain import hub
from langchain_community.vectorstores.elasticsearch import ElasticsearchStore
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from BCEmbedding.tools.langchain import BCERerank
import numpy as np

os.environ["QIANFAN_ACCESS_KEY"] = os.getenv('MY_QIANFAN_ACCESS_KEY')
os.environ["QIANFAN_SECRET_KEY"] = os.getenv('MY_QIANFAN_SECRET_KEY')

# 千帆 chatModel
qianfan_chat = QianfanChatEndpoint(
    model="ERNIE-Bot-4"
)

# Langsmith 配置，不用可注掉
unique_id = uuid4().hex[0:8]
os.environ["LANGCHAIN_PROJECT"] = f" [一条龙] qianfan Tracing Walkthrough - {unique_id}"
# os.environ["LANGCHAIN_TRACING_V2"] = 'true'
os.environ["LANGCHAIN_API_KEY"] = os.getenv('MY_LANGCHAIN_API_KEY')

if __name__ == '__main__':

    # 使用本地网易有道模型  BCE， 使用CPU
    embedding_model_name = 'D:\LLM\\bce_modesl\\bce-embedding-base_v1'
    embedding_model_kwargs = {'device': 'cpu'}
    embedding_encode_kwargs = {'batch_size': 32, 'normalize_embeddings': True, }

    embed_model = HuggingFaceEmbeddings(
        model_name=embedding_model_name,
        model_kwargs=embedding_model_kwargs,
        encode_kwargs=embedding_encode_kwargs
    )

    vectorstore = ElasticsearchStore(
        es_url=os.environ['ELASTIC_HOST_HTTP'],
        index_name="index_ex_768_vectors",
        embedding=embed_model,
        es_user="elastic",
        vector_query_field='question_vectors',
        es_password=os.environ['ELASTIC_ACCESS_PASSWORD']
    )

    retriever = vectorstore.as_retriever(search_kwargs={"k": 20})

    query = '粥余知识库'
    # 使用GPU
    reranker_args = {'model': 'D:\LLM\\bce_modesl\\bce-reranker-base_v1', 'top_n': 10, 'device': 'cuda:0'}
    reranker = BCERerank(**reranker_args)

    compression_retriever = ContextualCompressionRetriever(
        base_compressor=reranker, base_retriever=retriever
    )

    # rerank 过滤文档
    # docs = compression_retriever.get_relevant_documents(query)

    # 文档分割器
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=400)

    # 向量过滤
    embedding_filter = EmbeddingsFilter(embeddings=embed_model, similarity_threshold=0.6)

    # 文档去重
    redundant_filter = EmbeddingsRedundantFilter(embeddings=embed_model)

    # 内容抽取
    llm_extractor = LLMChainExtractor.from_llm(qianfan_chat)

    pipeline_compressor = DocumentCompressorPipeline(
        transformers=[text_splitter, embedding_filter, redundant_filter, reranker, llm_extractor]
    )
    pipeline_retriever = ContextualCompressionRetriever(
        base_compressor=pipeline_compressor, base_retriever=retriever
    )
    docs = pipeline_retriever.get_relevant_documents(query)
    pass

from BCEmbedding import RerankerModel, EmbeddingModel

if __name__ == '__main__':

    # 测试嵌入模块
    embedding_model_name = 'D:\LLM\\bce_modesl\\bce-embedding-base_v1'

    # list of sentences
    sentences = ['哈喽啊兄弟', '明天去上学了']

    # init embedding model
    model = EmbeddingModel(model_name_or_path=embedding_model_name)

    # extract embeddings
    # embeddings = model.encode(sentences)
    # pass

    # your query and corresponding passages
    query = '你最喜欢的电影是什么？'
    passages = [
        '我最喜欢的电影是《阿甘正传》。',
        '我最喜欢的电影类型是恐怖片',
        '我最喜欢的电影导演是斯皮尔伯格。',
        '我最喜欢的书是《哈利波特》。',
        '我不喜欢看电影。',
        '我喜欢看科幻电影，尤其是《星际穿越》。',
        '我最喜欢的音乐是爵士乐。',
        '昨天晚上我看了一部很好的电影。',
        '我最喜欢的运动是篮球。',
        '我喜欢吃披萨。'
    ]

    # construct sentence pairs
    sentence_pairs = [[query, passage] for passage in passages]
    reranker_model = 'D:\LLM\\bce_modesl\\bce-reranker-base_v1'
    # init reranker model
    model = RerankerModel(model_name_or_path=reranker_model)

    # method 0: calculate scores of sentence pairs
    scores = model.compute_score(sentence_pairs)

    # method 1: rerank passages
    rerank_results = model.rerank(query, passages)
    doc_list = rerank_results['rerank_passages']
    for line in doc_list:
        print(line)
    pass