"""
向量化和 Chroma 存储过程详解
从切割后的文档到向量数据库的完整流程
"""

print("=" * 80)
print("向量化和 Chroma 存储过程详解")
print("=" * 80)

# ============================================================================
# Part 1: 完整流程概览
# ============================================================================
print("\n" + "=" * 80)
print("📊 Part 1: 完整流程概览")
print("=" * 80)

print("""
从文档切割到向量数据库的完整流程：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Step 1: 文档切割
原始文档 → RecursiveCharacterTextSplitter → 20 个 chunks
        (5000 tokens)                         (每个 250 tokens)

Step 2: 向量化 (Embedding)
每个 chunk → HuggingFace 模型 → 向量 (384维)
   "人工智能是..."  →  [0.12, -0.34, 0.56, ...]

Step 3: 存入 Chroma
向量 + 原文 + 元数据 → Chroma 数据库
                      └─ 持久化存储

Step 4: 构建索引
Chroma → HNSW 索引 → 快速近似检索
        (层次化图结构)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
""")


# ============================================================================
# Part 2: Embedding 模型详解
# ============================================================================
print("\n" + "=" * 80)
print("🤖 Part 2: Embedding 模型 - HuggingFaceEmbeddings")
print("=" * 80)

print("""
你的项目配置：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

self.embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={'device': device},  # CPU 或 GPU
    encode_kwargs={'normalize_embeddings': True}  # 归一化
)

模型说明：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

模型名称: all-MiniLM-L6-v2
├─ 类型: Sentence-BERT (双编码器)
├─ 参数量: 22M (轻量级)
├─ 输出维度: 384 维向量
├─ 训练数据: 10亿+ 句子对
└─ 特点: 快速、准确、适合语义检索

工作原理：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

输入文本: "人工智能是计算机科学的一个分支"
    ↓
Tokenization (分词)
    ↓
Token IDs: [101, 782, 1435, 1819, 2510, 3221, ...]
    ↓
BERT Encoder (6 层 Transformer)
    ↓
[CLS] Token 的向量表示
    ↓
384 维向量: [0.123, -0.456, 0.789, ...]
    ↓
L2 归一化 (normalize_embeddings=True)
    ↓
最终向量: ||v|| = 1 (单位向量)
""")


# ============================================================================
# Part 3: 向量化过程分步解析
# ============================================================================
print("\n" + "=" * 80)
print("🔍 Part 3: 向量化过程 - 逐步解析")
print("=" * 80)

print("""
假设我们有 3 个 chunks：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Chunk 1: "人工智能是计算机科学的一个分支。它致力于..."
Chunk 2: "机器学习是人工智能的子领域。它使计算机..."
Chunk 3: "深度学习使用多层神经网络来处理复杂的..."

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

向量化过程（批量处理）：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

embeddings.embed_documents([chunk1, chunk2, chunk3])
                    ↓
┌──────────────────────────────────────────────────────────┐
│         HuggingFace Embedding 模型                        │
│         (sentence-transformers/all-MiniLM-L6-v2)         │
└──────────────────────────────────────────────────────────┘
                    ↓
        内部处理（每个 chunk）：
                    ↓
    ┌───────────────────────────────────────┐
    │ Step 1: Tokenization                  │
    │ "人工智能..." → [101, 782, 1435, ...] │
    └───────────────────────────────────────┘
                    ↓
    ┌───────────────────────────────────────┐
    │ Step 2: 转换为 Token Embeddings       │
    │ Token IDs → 初始向量表示               │
    └───────────────────────────────────────┘
                    ↓
    ┌───────────────────────────────────────┐
    │ Step 3: BERT Encoder (6 层)          │
    │ Self-Attention + Feed Forward         │
    │ 每层提取更深层的语义                  │
    └───────────────────────────────────────┘
                    ↓
    ┌───────────────────────────────────────┐
    │ Step 4: Mean Pooling                  │
    │ 所有 token 向量的平均 → 句子向量      │
    └───────────────────────────────────────┘
                    ↓
    ┌───────────────────────────────────────┐
    │ Step 5: L2 Normalization              │
    │ 向量归一化到单位长度                  │
    └───────────────────────────────────────┘
                    ↓
        输出：3 个向量
                    ↓
┌─────────────────────────────────────────────────────────┐
│ Vector 1: [0.123, -0.456, 0.789, ..., 0.321]  (384维)  │
│ Vector 2: [0.234, 0.567, -0.890, ..., 0.432]  (384维)  │
│ Vector 3: [-0.345, 0.678, 0.901, ..., -0.543] (384维)  │
└─────────────────────────────────────────────────────────┘

关键点：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
✅ 每个 chunk → 1 个固定维度的向量 (384维)
✅ 语义相似的文本 → 向量距离近
✅ 归一化后可用余弦相似度快速比较
""")


# ============================================================================
# Part 4: Chroma 数据库存储结构
# ============================================================================
print("\n" + "=" * 80)
print("💾 Part 4: Chroma 数据库存储结构")
print("=" * 80)

print("""
Chroma.from_documents() 执行的操作：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Chroma.from_documents(
    documents=doc_splits,          # 20 个 chunks
    collection_name="rag-chroma",  # 集合名称
    embedding=self.embeddings      # Embedding 函数
)

内部流程：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Step 1: 创建/打开集合
┌──────────────────────────────────────┐
│ Collection: "rag-chroma"             │
│ 元数据: embedding_dimension=384     │
└──────────────────────────────────────┘

Step 2: 批量向量化
for chunk in doc_splits:
    vector = embeddings.embed_documents([chunk.page_content])
    ↓

Step 3: 存储数据（每个 chunk）
┌─────────────────────────────────────────────────────────┐
│ ID: "chunk_1"                                           │
│ ├─ Vector: [0.123, -0.456, ..., 0.321]  (384维)        │
│ ├─ Document: "人工智能是计算机科学的一个分支..."       │
│ └─ Metadata: {                                          │
│      "source": "https://...",                           │
│      "chunk_index": 0,                                  │
│      "total_chunks": 20                                 │
│    }                                                    │
├─────────────────────────────────────────────────────────┤
│ ID: "chunk_2"                                           │
│ ├─ Vector: [0.234, 0.567, ..., 0.432]                  │
│ ├─ Document: "机器学习是人工智能的子领域..."           │
│ └─ Metadata: {...}                                      │
├─────────────────────────────────────────────────────────┤
│ ID: "chunk_3"                                           │
│ ├─ Vector: [-0.345, 0.678, ..., -0.543]                │
│ ├─ Document: "深度学习使用多层神经网络..."             │
│ └─ Metadata: {...}                                      │
└─────────────────────────────────────────────────────────┘

Step 4: 构建 HNSW 索引
向量 → HNSW 图结构 → 快速检索
      (层次化导航小世界图)

存储位置：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
默认路径: ./chroma/  (本地目录)
├─ collections/
│  └─ rag-chroma/
│     ├─ data.parquet      # 向量数据
│     ├─ metadata.json     # 元数据
│     └─ index.bin         # HNSW 索引
└─ chroma.sqlite3          # SQLite 数据库
""")


# ============================================================================
# Part 5: HNSW 索引工作原理
# ============================================================================
print("\n" + "=" * 80)
print("🔗 Part 5: HNSW 索引 - 快速检索的秘密")
print("=" * 80)

print("""
HNSW = Hierarchical Navigable Small World
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

为什么需要索引？
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
暴力搜索: O(n) - 计算查询向量与所有向量的距离
  └─ 10000 个向量 → 需要计算 10000 次距离
  └─ 太慢！

HNSW 索引: O(log n) - 层次化图结构导航
  └─ 10000 个向量 → 只需检查约 20-30 个节点
  └─ 快 100+ 倍！

HNSW 结构（简化示例）：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Layer 2 (最稀疏)
  V₁ ←──────→ V₅ ←──────→ V₁₂
   ↓           ↓           ↓

Layer 1
  V₁ ←→ V₃ ←→ V₅ ←→ V₈ ←→ V₁₂
   ↓     ↓     ↓     ↓     ↓

Layer 0 (最密集)
  V₁ ← V₂ ← V₃ ← V₄ ← V₅ ← V₆ ← ... ← V₁₂
  所有向量都在这一层

检索过程：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

查询向量: Q = [0.2, -0.3, 0.5, ...]

Step 1: 从 Layer 2 开始（粗略搜索）
  入口点: V₁
  → 计算 dist(Q, V₁), dist(Q, V₅), dist(Q, V₁₂)
  → V₅ 最近 → 跳到 V₅

Step 2: 下降到 Layer 1（中等精度）
  从 V₅ 开始
  → 检查邻居 V₃, V₈
  → V₈ 最近 → 跳到 V₈

Step 3: 下降到 Layer 0（高精度）
  从 V₈ 开始
  → 检查所有邻居
  → 找到最近的 K 个向量

返回结果: Top K 最相似的 chunks

速度对比：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
暴力搜索: 10000 次距离计算 → 100ms
HNSW 索引: 20-30 次距离计算 → 1ms  ← 快 100 倍！
""")


# ============================================================================
# Part 6: 检索过程详解
# ============================================================================
print("\n" + "=" * 80)
print("🔍 Part 6: 检索过程 - 从查询到结果")
print("=" * 80)

print("""
用户查询: "什么是机器学习？"
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Step 1: 查询向量化
─────────────────────────────────────────────────────────
"什么是机器学习？"
    ↓
embeddings.embed_query("什么是机器学习？")
    ↓
Query Vector: [0.345, -0.678, 0.234, ...]  (384维)


Step 2: HNSW 近似搜索
─────────────────────────────────────────────────────────
vectorstore.similarity_search(
    query="什么是机器学习？",
    k=20  # 返回 Top 20
)
    ↓
Chroma 内部:
1. 查询向量化
2. HNSW 图导航
3. 计算余弦相似度
    ↓
返回 Top 20 chunks:
┌──────────┬─────────┬────────────────────────────┐
│ Chunk ID │  Score  │        Content             │
├──────────┼─────────┼────────────────────────────┤
│ chunk_5  │  0.92   │ "机器学习是人工智能的..." │
│ chunk_2  │  0.88   │ "人工智能包括机器学习..." │
│ chunk_11 │  0.85   │ "监督学习是机器学习..."   │
│ ...      │  ...    │ ...                        │
└──────────┴─────────┴────────────────────────────┘


Step 3: CrossEncoder 重排（你的项目特色）
─────────────────────────────────────────────────────────
reranker.rerank(query, top_20_chunks, top_k=5)
    ↓
每个 chunk 重新打分（深度交互）
    ↓
最终 Top 5:
┌──────────┬─────────┬────────────────────────────┐
│ Chunk ID │  Score  │        Content             │
├──────────┼─────────┼────────────────────────────┤
│ chunk_5  │  8.45   │ "机器学习是人工智能的..." │
│ chunk_11 │  7.89   │ "监督学习是机器学习..."   │
│ chunk_2  │  7.23   │ "人工智能包括机器学习..." │
│ chunk_14 │  6.78   │ "深度学习是机器学习..."   │
│ chunk_8  │  6.12   │ "强化学习允许..."         │
└──────────┴─────────┴────────────────────────────┘


Step 4: 返回给 LLM
─────────────────────────────────────────────────────────
context = "\\n\\n".join([chunk.page_content for chunk in top_5])
    ↓
LLM 生成答案
""")


# ============================================================================
# Part 7: 关键技术细节
# ============================================================================
print("\n" + "=" * 80)
print("⚙️  Part 7: 关键技术细节")
print("=" * 80)

print("""
1. 为什么要归一化向量？
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
encode_kwargs={'normalize_embeddings': True}

原始向量: [1.23, -4.56, 7.89, ...]  # 长度不一
归一化后: [0.12, -0.45, 0.78, ...]  # 长度 = 1

好处:
✅ 余弦相似度 = 点积（计算更快）
✅ 所有向量在同一尺度上
✅ 避免长度影响相似度计算


2. 余弦相似度 vs 欧氏距离
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

余弦相似度（你的项目使用）⭐:
  similarity = v₁ · v₂ / (||v₁|| × ||v₂||)
  范围: [-1, 1]，1 表示完全相同
  特点: 关注方向，忽略长度

欧氏距离:
  distance = √Σ(v₁ᵢ - v₂ᵢ)²
  范围: [0, ∞]，0 表示完全相同
  特点: 关注绝对位置差异

归一化后，两者等价！


3. 批量处理优化
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

不推荐（慢）:
for chunk in chunks:
    vector = embed_documents([chunk])  # 单独处理
    
推荐（快 10 倍）⭐:
vectors = embed_documents(chunks)  # 批量处理
  └─ GPU 并行计算
  └─ 减少模型加载开销


4. 内存优化
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

向量维度选择:
  384 维 (all-MiniLM-L6-v2)   ← 你的项目 ⭐
  └─ 平衡：准确率 vs 存储

  768 维 (BERT-base)
  └─ 更准确但存储翻倍

  1024 维 (large models)
  └─ 最准确但存储 3 倍

存储计算:
20 个 chunks × 384 维 × 4 bytes = 30KB
1000 个 chunks × 384 维 × 4 bytes = 1.5MB
  └─ 非常高效！
""")


# ============================================================================
# Part 8: 完整代码流程
# ============================================================================
print("\n" + "=" * 80)
print("💻 Part 8: 完整代码流程总结")
print("=" * 80)

print("""
你的项目完整流程：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# 1. 初始化 Embedding 模型
embeddings = HuggingFaceEmbeddings(
    model_name="sentence-transformers/all-MiniLM-L6-v2",
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': True}
)

# 2. 文档切割
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=250,
    chunk_overlap=50  # ← 你刚修改的
)
doc_splits = text_splitter.split_documents(docs)

# 3. 向量化 + 存储到 Chroma
vectorstore = Chroma.from_documents(
    documents=doc_splits,      # 输入: 20 个 chunks
    collection_name="rag-chroma",
    embedding=embeddings       # 向量化函数
)
# ↓ 内部自动完成:
#   - 批量向量化: chunks → 384维向量
#   - 存储: 向量 + 原文 + 元数据
#   - 构建 HNSW 索引

# 4. 创建检索器
retriever = vectorstore.as_retriever()

# 5. 检索
docs = retriever.get_relevant_documents("什么是机器学习？")
# ↓ 内部流程:
#   - 查询向量化
#   - HNSW 快速检索
#   - 返回 Top K chunks

# 6. CrossEncoder 重排（可选，你的项目有）
reranked = crossencoder.rerank(query, docs, top_k=5)

# 7. 喂给 LLM 生成答案
answer = llm.generate(context=docs, question=query)

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
""")


# ============================================================================
# Part 9: 性能优化建议
# ============================================================================
print("\n" + "=" * 80)
print("🚀 Part 9: 性能优化建议")
print("=" * 80)

print("""
当前配置评分：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

✅ Embedding 模型: all-MiniLM-L6-v2  (轻量高效) ⭐⭐⭐⭐⭐
✅ 向量归一化: True                 (余弦相似度优化) ⭐⭐⭐⭐⭐
✅ 索引类型: HNSW                   (快速检索) ⭐⭐⭐⭐⭐
✅ Chunk overlap: 50                (保持上下文) ⭐⭐⭐⭐⭐
✅ CrossEncoder 重排                (精准排序) ⭐⭐⭐⭐⭐

总评: 🏆 生产级配置！

可选优化（如需进一步提升）：
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

1. GPU 加速
   model_kwargs={'device': 'cuda'}  # 向量化速度 10x ↑

2. 更大的 Embedding 模型（如需更高准确率）
   "BAAI/bge-large-en-v1.5"  # 1024维，准确率 +5%

3. 批量大小调整
   batch_size=32  # 加快向量化

4. Chroma 持久化配置
   persist_directory="./chroma_db"  # 避免重复向量化
""")


print("\n" + "=" * 80)
print("✅ 解析完成！你现在理解了从切割到向量数据库的完整流程")
print("=" * 80)
print()