# bge-m3模型
def dm01():
    from milvus_model.hybrid import BGEM3EmbeddingFunction

    m3_path = "../rag_qa/models/bge-m3"
    embedding_function = BGEM3EmbeddingFunction(model_name_or_path=m3_path, use_fp16=False, device='cpu')

    # 对文本进行Embedding
    embed = embedding_function(["黑马JAVA", "黑马大模型"])

    # 打印Embedding后的结果
    print("embed-->", embed)
    print("embed.dense-->", embed["dense"])
    print("embed.sparse-->", embed["sparse"])

    row = embed["dense"][0]
    print("row-->", row)


# reranker 模型
def dm02():
    from sentence_transformers import CrossEncoder
    reranker_path = "../rag_qa/models/bge-reranker-large"
    reranker = CrossEncoder(reranker_path)

    queries = ["什么是机器学习？"]
    documents = [
        "深度学习是机器学习的一个分支。",
        "今天的天气很好。",
        "机器学习是一种让计算机从数据中自动学习规律的方法。",
    ]

    pairs = [[queries[0], doc] for doc in documents]
    print("pairs-->", pairs)
    scores = reranker.predict(pairs)
    print(f'scores-->', scores)

    # 输出排序结果
    for doc, score in sorted(zip(documents, scores), key=lambda x: x[1], reverse=True):
        print(f"{score:.4f} - {doc}")


if __name__ == '__main__':
    # dm01()
    dm02()