pip install -U pymilvus
pip install "pymilvus[model]"

conda install pytorch::pytorch torchvision torchaudio -c pytorch

import torch
print(torch.__version__)
# 2.5.1

from pymilvus import Collection

# 创建一个名为 "example_collection" 的集合
collection = Collection("example_collection")

from pymilvus import MilvusClient, DataType, Collection

# Step 1️⃣ 创建 Milvus 客户端并初始化 Collection
# MilvusClient 是与 Milvus 服务进行交互的客户端，连接到指定的数据库。
client = MilvusClient("milvus_rag_zh_demo_2.db")  # 假设使用一个本地数据库文件作为存储

# COLLECTION_NAME 是我们想要操作的集合名
COLLECTION_NAME = "zh_rag_demo_2"

# Step 2️⃣ 删除已有的 Collection（如果存在）
# 如果该集合已经存在，我们删除旧的集合以重新创建
if client.has_collection(COLLECTION_NAME):
    print(f"Collection '{COLLECTION_NAME}' already exists. Dropping it.")
    client.drop_collection(COLLECTION_NAME)  # 删除集合

# Step 3️⃣ 定义 Collection 的 Schema
# Collection Schema 用于定义集合中的字段，包括数据类型、维度等。
schema = MilvusClient.create_schema(
    auto_id=False,  # 不自动生成 ID（我们手动指定 ID）
    enable_dynamic_field=True,  # 启用动态字段支持
)

# Step 4️⃣ 添加字段到 Schema
# 每个字段需要指定字段名称、数据类型，以及其他一些属性（如是否是主键）
schema.add_field(field_name="my_id", datatype=DataType.INT64, is_primary=True)  # 主键字段
schema.add_field(field_name="my_vector", datatype=DataType.FLOAT_VECTOR, dim=5)  # 向量字段，维度为5
schema.add_field(field_name="my_varchar", datatype=DataType.VARCHAR, max_length=512)  # 字符串字段，最大长度为512

# 从 pymilvus 中引入用于 embedding 的模型工具
from pymilvus import model

# 如果访问 huggingface 官方地址失败（如网络受限），可以取消下面的注释
# 设置环境变量，切换为 huggingface 镜像地址（例如清华、阿里等镜像）
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 初始化一个默认的 embedding 函数，会自动下载一个小型的句子向量模型：
# paraphrase-albert-small-v2（约 50MB），适合快速测试或演示用途。
embedding_fn = model.DefaultEmbeddingFunction()

# 定义待向量化的文本列表，每条文本都是一个独立的文档或知识片段。
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

# 使用 embedding 函数将文本转化为向量（embedding）
# 得到的结果是一个 list，每条向量为一个 numpy 数组，维度为 768
vectors = embedding_fn.encode_documents(docs)

# 输出向量
print("vectors:", vectors)

# 输出向量维度信息（用于对齐 Milvus 中 Collection 的字段设计）
# embedding_fn.dim 表示当前模型生成的向量维度
# vectors[0].shape 表示第一条向量的维度
print("Dim:", embedding_fn.dim, vectors[0].shape)  # 预期输出 Dim: 768 (768,)

# 组织向量数据为结构化的格式，准备后续插入 Milvus：
# 每条记录包含：
# - id：唯一编号
# - vector：生成的 embedding 向量
# - text：原始文本内容（可用于展示/检索返回）
# - subject：人为指定的分类标签（这里用作 metadata 过滤的演示）
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

# 打印生成数据的基本信息：
# - 实体数量
# - 每条记录包含哪些字段
# - 向量维度是多少
print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))

# 从 pymilvus 中引入用于 embedding 的模型工具
from pymilvus import model

# 如果访问 huggingface 官方地址失败（如网络受限），可以取消下面的注释
# 设置环境变量，切换为 huggingface 镜像地址（例如清华、阿里等镜像）
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 初始化一个默认的 embedding 函数，会自动下载一个小型的句子向量模型：
# paraphrase-albert-small-v2（约 50MB），适合快速测试或演示用途。
# embedding_fn = model.DefaultEmbeddingFunction()

# sentence transformers 模型
embedding_fn = model.dense.SentenceTransformerEmbeddingFunction(
    model_name='all-MiniLM-L6-v2',  # Specify the model name
    device='cpu'  # Specify the device to use, e.g., 'cpu' or 'cuda:0'
)

# # BGE M3 模型
# embedding_fn = model.hybrid.BGEM3EmbeddingFunction(
#     model_name='BAAI/bge-m3',  # Specify the model name
#     device='cpu',  # Specify the device to use, e.g., 'cpu' or 'cuda:0'
#     use_fp16=False  # Specify whether to use fp16. Set to `False` if `device` is `cpu`.
# )

# # splade 模型
# embedding_fn = model.sparse.SpladeEmbeddingFunction(
#     model_name="naver/splade-cocondenser-selfdistil",
#     device="cpu"
# )

# 定义待向量化的文本列表，每条文本都是一个独立的文档或知识片段。
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

# 使用 embedding 函数将文本转化为向量（embedding）
# 得到的结果是一个 list，每条向量为一个 numpy 数组，维度为 768
vectors = embedding_fn.encode_documents(docs)

# 输出向量
print("vectors:", vectors)

# 输出向量维度信息（用于对齐 Milvus 中 Collection 的字段设计）
# embedding_fn.dim 表示当前模型生成的向量维度
# vectors[0].shape 表示第一条向量的维度
print("Dim:", embedding_fn.dim, vectors[0].shape)  # 预期输出 Dim: 768 (768,)

# 组织向量数据为结构化的格式，准备后续插入 Milvus：
# 每条记录包含：
# - id：唯一编号
# - vector：生成的 embedding 向量
# - text：原始文本内容（可用于展示/检索返回）
# - subject：人为指定的分类标签（这里用作 metadata 过滤的演示）
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

# 打印生成数据的基本信息：
# - 实体数量
# - 每条记录包含哪些字段
# - 向量维度是多少
print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))

# 引入必要的库：pymilvus 中的模型工具和客户端类
from pymilvus import model
from pymilvus import MilvusClient

# 初始化一个 Milvus 客户端，连接到本地或指定路径的数据库文件
client = MilvusClient("milvus_demo.db")

# 如果已存在名为 "demo_collection" 的集合，先删除（避免重复创建时报错）
if client.has_collection(collection_name="demo_collection"):
    client.drop_collection(collection_name="demo_collection")

# 创建一个新的集合（Collection），用于存储向量数据
# 参数说明：
# - collection_name：集合名称
# - dimension：向量的维度（本例使用的是 768 维）
client.create_collection(
    collection_name="demo_collection",
    dimension=768,
)

# ⚠️ 如果无法连接 huggingface 官方模型仓库（比如被墙），可以使用国内镜像站点：
# 设置 HF_ENDPOINT 环境变量来更改默认的 huggingface 地址，例如使用清华镜像：
# import os
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

# 初始化一个默认的 embedding 函数，会自动下载 paraphrase-albert-small-v2 模型
# 模型体积小，加载快，适合用于快速测试/演示
embedding_fn = model.DefaultEmbeddingFunction()

# （可选）你也可以使用更强的模型，例如 SentenceTransformer、BGE M3、Splade 等
# 下面是一些替代方案（按需取消注释）：

# 使用 sentence-transformers 模型
# embedding_fn = model.dense.SentenceTransformerEmbeddingFunction(
#     model_name='all-MiniLM-L6-v2',
#     device='cpu'  # 或使用 'cuda:0' 来使用 GPU
# )

# 使用 BGE M3 多任务嵌入模型
# embedding_fn = model.hybrid.BGEM3EmbeddingFunction(
#     model_name='BAAI/bge-m3',
#     device='cpu',
#     use_fp16=False
# )

# 使用稀疏向量模型 Splade（适用于稀疏向量检索场景）
# embedding_fn = model.sparse.SpladeEmbeddingFunction(
#     model_name="naver/splade-cocondenser-selfdistil",
#     device="cpu"
# )

# 准备文本数据：这些将作为示例文档被嵌入为向量
docs = [
    "Artificial intelligence was founded as an academic discipline in 1956.",
    "Alan Turing was the first person to conduct substantial research in AI.",
    "Born in Maida Vale, London, Turing was raised in southern England.",
]

# 使用嵌入函数将文本转换为向量
# 返回值为一个列表，每个元素是一个 numpy 数组（维度为 768）
vectors = embedding_fn.encode_documents(docs)

# 打印生成的向量（可以注释掉以避免打印过多内容）
print("vectors:", vectors)

# 打印嵌入函数维度和实际向量维度（用于确认维度一致性）
print("Dim:", embedding_fn.dim, vectors[0].shape)

# 将向量数据组织为结构化格式，便于后续插入到 Milvus 中
# 每个字典代表一条记录，包括：
# - id：唯一编号
# - vector：对应的嵌入向量
# - text：原始文本
# - subject：可选分类标签（可用于后续过滤查询）
data = [
    {"id": i, "vector": vectors[i], "text": docs[i], "subject": "history"}
    for i in range(len(vectors))
]

# 打印数据结构信息：
print("Data has", len(data), "entities, each with fields: ", data[0].keys())
print("Vector dim:", len(data[0]["vector"]))

# 将数据插入到 Milvus 中的 demo_collection 集合
res = client.insert(collection_name="demo_collection", data=data)

# 打印插入结果（通常包含成功插入的数量等信息）
print("res:", res)

# 接下来执行向量搜索（即：语义检索）
# 使用同样的嵌入函数对查询语句进行编码，得到向量
query_vectors = embedding_fn.encode_queries(["Who is Alan Turing?"])

# ⚠️ 如果没有 embedding 函数，也可以使用随机向量模拟查询（仅限演示用）
# import random
# query_vectors = [ [ random.uniform(-1, 1) for _ in range(768) ] ]

# 在 demo_collection 中执行向量检索
res_search = client.search(
    collection_name="demo_collection",  # 检索目标集合
    data=query_vectors,  # 查询向量
    limit=2,  # 返回最相似的 2 条记录
    output_fields=["text", "subject"],  # 指定返回的字段（原始文本和标签）
)

# 打印检索结果（包括匹配文本和相似度等信息）
print("res_search:", res_search)

# 提取并打印搜索结果中的具体条目
print(list(res_search[0])[0])  # 第一个查询返回的第一条结果
print(list(res_search)[0][1])  # 第一个查询返回的第二条结果

from pymilvus import MilvusClient
client = MilvusClient("milvus_demo.db")

# Drop collection
client.drop_collection(collection_name="demo_collection")

# 从 pymilvus.model.reranker 模块中引入 BGE 重排序模型函数
from pymilvus.model.reranker import BGERerankFunction

# ----------------------------------------------------------
# 初始化 Reranker（重排序器）：
# 使用 BAAI（智源研究院）提供的 BGE-Reranker-v2-M3 模型。
# 该模型基于 Cross-Encoder 架构，通过对 query 和文档对进行语义交互建模，
# 输出相关性打分，用于对初始检索结果进行排序提升精度。
# ----------------------------------------------------------
bge_rf = BGERerankFunction(
    model_name="BAAI/bge-reranker-v2-m3",  # 模型名称，默认即为该模型
    device="cpu"  # 计算设备，可改为 'cuda:0' 使用 GPU 加速
)

# ----------------------------------------------------------
# 定义一个查询（query），用于检索历史相关信息。
# ----------------------------------------------------------
query = "What event in 1956 marked the official birth of artificial intelligence as a discipline?"

# ----------------------------------------------------------
# 定义候选文档列表（documents）：
# 模拟从 Milvus 检索返回的初步候选文本片段（Top-K），
# 接下来将使用 Reranker 进一步对它们进行精排。
# 文档集合里面有四个元素，从0到3编号。
# ----------------------------------------------------------
documents = [
    "In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.",
    "The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.",
    "In 1951, British mathematician and computer scientist Alan Turing also developed the first program designed to play chess, demonstrating an early example of AI in game strategy.",
    "The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems."
]

# ----------------------------------------------------------
# 执行重排序操作：
# 将 query 与每个文档组合进行相关性评分，输出按分值降序排列的 Top-K 文档。
# top_k 参数指定只保留得分最高的前 K 条（这里设为 3）。
# 返回结果为一个包含 RerankResult 对象的列表，每个对象包含：
# - index：原文档在输入列表中的索引位置
# - score：query 与该文档的语义相关性打分（越高越相关）
# - text：文档原文内容
# ----------------------------------------------------------
results = bge_rf(
    query=query,
    documents=documents,
    top_k=3,  # 返回得分最高的前 3 条
)

# ----------------------------------------------------------
# 遍历输出精排结果：
# 展示每条候选文档的原始索引、得分（保留 6 位小数）、文本内容。
# 注意结果已按 score 排序，score 越高代表与 query 越匹配。
# ----------------------------------------------------------
for result in results:
    print(f"Index: {result.index}")  # 文本在原始 documents 中的位置
    print(f"Score: {result.score:.6f}")  # 重排序得分
    print(f"Text: {result.text}\n")  # 文本内容

query = "1956年标志着人工智能正式诞生的事件是什么？"

documents = [
    "1950年，图灵发表了著名论文《计算机器与智能》，提出了图灵测试，这是人工智能哲学的重要基石。",
    "1956年达特茅斯会议被广泛认为是人工智能作为一个学科诞生的标志，在这次会议上，“人工智能”这一术语首次被提出，并确立了研究目标。",
    "1951年，英国数学家图灵开发了一个可以下棋的程序，展示了早期的人工智能应用。",
    "1955年，艾伦·纽厄尔与赫伯特·西蒙等人发明了“逻辑理论家”程序，这是第一个能够自动证明定理的人工智能系统。"
]

from sentence_transformers import CrossEncoder

model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2')
scores = model.predict([
    ("How many people live in Berlin?", "Berlin had a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers."),
    ("How many people live in Berlin?", "Berlin is well known for its museums."),
])
print(scores)

# [ 8.607141 -4.320079]

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from sentence_transformers import CrossEncoder

# 方法一：
model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L12-v2')
scores = model.predict([
    ("How many people live in Berlin?", "Berlin had a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers."),
    ("How many people live in Berlin?", "Berlin is well known for its museums."),
])
print(scores)

# 方法二：
model = AutoModelForSequenceClassification.from_pretrained('cross-encoder/ms-marco-MiniLM-L12-v2')
tokenizer = AutoTokenizer.from_pretrained('cross-encoder/ms-marco-MiniLM-L12-v2')

features = tokenizer(['How many people live in Berlin?', 'How many people live in Berlin?'], ['Berlin had a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers.', 'Berlin is well known for its museums.'],  padding=True, truncation=True, return_tensors="pt")

model.eval()
with torch.no_grad():
    scores = model(**features).logits
    print(scores)

# 从 pymilvus.model.reranker 模块中导入交叉编码器重排序函数
from pymilvus.model.reranker import CrossEncoderRerankFunction

# ✅ 步骤 1：定义交叉编码器（Cross Encoder）重排序函数
# 该函数内部会自动下载 Hugging Face 上的模型并用于 rerank
ce_rf = CrossEncoderRerankFunction(
    model_name="cross-encoder/ms-marco-MiniLM-L12-v2",  # 使用指定的 cross-encoder 模型（支持语义匹配任务）
    device="cpu"  # 指定模型运行设备，如 'cpu' 或 'cuda:0'（GPU）
)

# ✅ 步骤 2：定义用户查询（Query）
query = "What event in 1956 marked the official birth of artificial intelligence as a discipline?"
# 中文翻译：1956年哪一事件标志着人工智能作为一门学科的正式诞生？

# ✅ 步骤 3：准备待重排序的候选文档（Documents）
# 每个字符串都是一个候选答案，Cross Encoder 会将它们与 Query 组合成一个个句对进行评分
documents = [
    "In 1950, Alan Turing published his seminal paper, 'Computing Machinery and Intelligence,' proposing the Turing Test as a criterion of intelligence, a foundational concept in the philosophy and development of artificial intelligence.",
    # 图灵1950年的论文，提出图灵测试，为AI发展奠定哲学基础

    "The Dartmouth Conference in 1956 is considered the birthplace of artificial intelligence as a field; here, John McCarthy and others coined the term 'artificial intelligence' and laid out its basic goals.",
    # 1956年达特茅斯会议，被广泛认为是AI的诞生标志

    "In 1951, British mathematician and computer scientist Alan Turing also developed the first program designed to play chess, demonstrating an early example of AI in game strategy.",
    # 图灵在1951年开发了下棋程序，展示早期AI在博弈中的应用

    "The invention of the Logic Theorist by Allen Newell, Herbert A. Simon, and Cliff Shaw in 1955 marked the creation of the first true AI program, which was capable of solving logic problems, akin to proving mathematical theorems."
    # 1955年“逻辑理论家”程序是首个能解逻辑问题的AI程序
]

# ✅ 步骤 4：调用重排序函数进行语义匹配排序（Reranking）
# Cross Encoder 会对 (query, document) 成对输入进行语义评分，返回得分最高的 top_k 条
results = ce_rf(
    query=query,
    documents=documents,
    top_k=3,  # 返回得分最高的前 3 个文档
)

# ✅ 步骤 5：遍历结果，输出每条结果的索引、分数、文本内容
# Cross Encoder 输出的是基于语义匹配的相关性打分，越高越相关
for result in results:
    print(f"Index: {result.index}")       # 文档在原始列表中的索引
    print(f"Score: {result.score:.6f}")   # Cross Encoder 计算出的相关性得分
    print(f"Text: {result.text}\n")       # 对应的文档内容

from pymilvus import MilvusClient
from pymilvus import model
from pymilvus.model.reranker import CrossEncoderRerankFunction
from sentence_transformers import SentenceTransformer
import ollama
import time

# Step 1️⃣ 创建 Milvus 客户端并初始化 Collection
client = MilvusClient("milvus_rag_zh.db")
COLLECTION_NAME = "zh_rag_demo"

# 若已存在，删除原 Collection
if client.has_collection(COLLECTION_NAME):
    client.drop_collection(COLLECTION_NAME)

# 创建新的向量集合，维度要和选用的 embedding 模型一致
client.create_collection(
    collection_name=COLLECTION_NAME,
    dimension=768,  # 使用默认模型生成的是 768 维向量
)

# Step 2️⃣ 中文句向量模型加载，用于对文档和查询编码
embedding_model = model.DefaultEmbeddingFunction()


# 准备中文知识库文档（可视为 RAG 的数据库），伪造的数据，如有雷同，纯属巧合。
docs = [
    "A公司由五个一级部门组成：战略规划部、产品研发部、市场销售部、人力资源部与财务管理部。",
    "战略规划部负责公司整体战略制定与年度业务目标设定，由高级副总裁王蕾领导。",
    "产品研发部下设三个小组：智能算法组、后端架构组、移动客户端组，分别由赵新宇、韩飞和李倩倩负责。",
    "人力资源部目前有6名成员，涵盖招聘、培训、绩效、薪酬等模块，负责人为HR总监陈蓉。",
    "市场销售部主要负责B端客户开拓与服务，目前设有北京、上海、深圳三个销售大区。",
    "财务管理部由CFO黄志明牵头，下设预算组、税务组与资金组，分别管理不同的财务职责。",
    "A公司CEO为沈立峰，CTO为冯启航，两人每季度召开一次高管战略会，协调部门合作。",
    "智能算法组正在负责天眼项目的模型迭代，预计Q3完成V2版本部署。",
    "后端架构组正在推动服务中台的容器化改造，目前已完成40%的微服务迁移。",
    "移动客户端组刚完成v6.2.0版本的迭代，新增了消息推送与报表订阅功能。"
]

# 将中文文档编码为向量（embedding）
vectors = embedding_model.encode_documents(docs)

# 组织结构化数据，插入到 Milvus 中
data_to_insert = [
    {"id": i, "vector": vectors[i], "text": docs[i], "tag": "公司组织架构"}
    for i in range(len(docs))
]

client.insert(collection_name=COLLECTION_NAME, data=data_to_insert)

# Step 3️⃣ 用户输入查询（中文）
query = "谁是A公司的CTO？"

# 查询向量化
query_vector = embedding_model.encode_documents([query])

# Step 4️⃣ 从 Milvus 中基于向量检索 Top-K 最相关文档
search_results = client.search(
    collection_name=COLLECTION_NAME,
    data=query_vector,
    limit=5,
    output_fields=["text"]
)

# print(search_results)
# print(list(search_results[0]))

# 提取原始文本内容
retrieved_texts = [doc["entity"]["text"] for doc in list(search_results[0])]

# Step 5️⃣ 使用 CrossEncoder 进行重排序，提升相关性排名
reranker = CrossEncoderRerankFunction(
    model_name="cross-encoder/ms-marco-MiniLM-L12-v2",
    device="cpu"
)

reranked = reranker(
    query=query,
    documents=retrieved_texts,
    top_k=3  # 可以根据需要调整返回数量
)

# 拼接重排序后的文档，作为 LLM 的上下文
context_text = "\n".join([r.text for r in reranked])

# Step 6️⃣ 构造 prompt，准备传入 LLM 生成回答
prompt = f"""
请根据以下背景知识回答用户提出的问题：

【背景知识】
{context_text}

【问题】
{query}

【回答】
"""

# Step 7️⃣ 使用 Ollama + DeepSeek 模型生成回答（流式输出）
print("🤖 正在生成回答...\n")

# 调用 Ollama SDK，使用 streaming=True 获取流式输出
stream = ollama.chat(
    model="deepseek-r1:32b",
    messages=[{"role": "user", "content": prompt}],
    stream=True
)

# Step 8️⃣ 打印流式输出结果
response_buffer = ""
for chunk in stream:
    delta = chunk.get("message", {}).get("content", "")
    print(delta, end="", flush=True)  # 实时输出
    response_buffer += delta

from pymilvus import MilvusClient

# 创建 Milvus 客户端并连接数据库
client = MilvusClient("milvus_rag_zh.db")
COLLECTION_NAME = "zh_rag_demo"

# 获取集合的collection信息。要查看数据库中的所有集合，你可以使用 list_collections() 方法，它将返回当前数据库中的所有集合名称。
collection_list = client.list_collections()
print("List of Collections:", collection_list)

# 查看集合的统计信息，要查看集合中包含的实体数量，你可以使用 get_collection_stats() 方法。它将返回集合的统计信息，其中包括存储的向量数量（即实体数量）。
collection_stats = client.get_collection_stats(COLLECTION_NAME)
print("Collection Stats:", collection_stats)

# 还可以获取特定 Collection 的详细信息
desc_collection = client.describe_collection(collection_name=COLLECTION_NAME)
print("Describe Collection:", desc_collection)