Merge remote main

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -0,0 +1,77 @@
+name: "问题反馈：代码或文档问题"
+description: "报告教程中的代码错误、文档错误、链接失效等问题"
+labels: ["bug", "needs triage"]
+title: "[问题反馈] 简要概述问题"
+body:
+  - type: checkboxes
+    id: checklist
+    attributes:
+      label: 提交前检查
+      description: 为了提升处理效率，请先确认以下事项
+      options:
+        - label: 我已在 Issues 中搜索，未找到相同问题
+          required: true
+        - label: 我已使用仓库最新版（或 `main` 分支）进行复现
+          required: true
+  - type: dropdown
+    id: issue_type
+    attributes:
+      label: 问题类型
+      options:
+        - 代码运行错误
+        - 文档描述错误/不清晰
+        - 链接失效/资源缺失
+        - 依赖/环境问题
+        - 其他
+    validations:
+      required: true
+  - type: input
+    id: location
+    attributes:
+      label: 位置（章节/文件路径）
+      description: 如 `docs/chapter4/11_hybrid_search.md` 或 `code/C8/main.py`
+      placeholder: 请输入章节编号、文档路径或代码相对路径
+    validations:
+      required: true
+  - type: textarea
+    id: description
+    attributes:
+      label: 问题描述
+      description: 请清晰说明问题表现与期望行为
+      placeholder: 实际结果 vs 期望结果
+    validations:
+      required: true
+  - type: textarea
+    id: reproduction
+    attributes:
+      label: 复现步骤
+      description: 清晰的复现步骤有助于快速定位
+      placeholder: |
+        1. ...
+        2. ...
+        3. ...
+  - type: textarea
+    id: minimal_code
+    attributes:
+      label: 最小可复现代码/命令
+      description: 请尽量提供最小复现片段，或给出命令行
+      render: shell
+      placeholder: |
+        python code/C8/main.py --config ...
+  - type: textarea
+    id: env
+    attributes:
+      label: 环境信息
+      description: 操作系统、Python/Node 版本、依赖版本等
+      placeholder: |
+        OS: Windows 11 22631
+        Python: 3.10.x
+        依赖: langchain==..., llama-index==...
+  - type: textarea
+    id: logs
+    attributes:
+      label: 日志/报错信息与截图
+      description: 可粘贴报错文本或上传截图
+      render: bash
+
+

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,3 @@
+blank_issues_enabled: false
+
+

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -0,0 +1,54 @@
+name: "需求建议：希望新增教程内容"
+description: "提出希望本仓库新增/扩展的教程主题或方向"
+labels: ["enhancement", "discussion"]
+title: "[需求建议] 简要概述主题/方向"
+body:
+  - type: checkboxes
+    id: checklist
+    attributes:
+      label: 提交前检查
+      options:
+        - label: 我已在 Issues 中搜索，未找到相同或高度相似的建议
+          required: true
+  - type: input
+    id: topic
+    attributes:
+      label: 希望新增的教程主题/模块
+      description: 尽量简洁、明确
+      placeholder: 例如：Agentic RAG、企业级项目
+    validations:
+      required: true
+  - type: textarea
+    id: context
+    attributes:
+      label: 背景与痛点
+      description: 为什么需要这一模块？解决哪些实际问题？
+      placeholder: 场景背景、业务需求、目前的缺口
+    validations:
+      required: true
+  - type: textarea
+    id: scope
+    attributes:
+      label: 建议覆盖的内容范围
+      description: 可按章节/功能点列出
+      placeholder: |
+        1. 主题概述与系统架构
+        2. 环境与数据准备
+        3. 检索/索引/生成一体化示例
+        4. 评估与监控
+  - type: textarea
+    id: audience
+    attributes:
+      label: 目标读者与预期效果
+      description: 面向人群与学完后能做什么
+      placeholder: 面向初学/进阶/生产落地；预期掌握的能力
+  - type: textarea
+    id: references
+    attributes:
+      label: 相关资料/链接（选填）
+      description: 参考项目、论文、博客、Demo 等
+      placeholder: |
+        - 论文/项目链接：...
+        - 竞品或类似教程：...
+
+

.github/ISSUE_TEMPLATE/study_feedback.yml

@@ -0,0 +1,45 @@
+name: "学习反馈：学习体验与改进建议"
+description: "分享你的学习体验、遇到的困惑，以及对教程改进的建议"
+labels: ["feedback"]
+title: "[学习反馈] 简要概述你的体验"
+body:
+  - type: dropdown
+    id: experience
+    attributes:
+      label: 整体学习体验
+      options:
+        - 非常好
+        - 较好
+        - 一般
+        - 需要改进
+        - 较差
+    validations:
+      required: true
+  - type: textarea
+    id: what_helped
+    attributes:
+      label: 觉得教程哪些部分做得比较好？
+      placeholder: 例如：章节结构清晰、代码可运行、图示易懂...
+  - type: textarea
+    id: what_confusing
+    attributes:
+      label: 学习过程中遇到的困惑/不足
+      description: 请指出具体章节或文件，便于改进
+      placeholder: 例如：`docs/chapter3/06_vector_embedding.md` 的某段描述不清晰
+  - type: textarea
+    id: improvement
+    attributes:
+      label: 你希望如何改进？
+      placeholder: 举例说明更好的呈现方式、补充案例、增加背景知识等
+  - type: input
+    id: time_cost
+    attributes:
+      label: 学习用时（选填）
+      placeholder: 例如：总计约 6 小时
+  - type: textarea
+    id: other
+    attributes:
+      label: 其他想说的话（选填）
+      placeholder: 欢迎分享你的收获或对作者的建议
+
+

.gitignore

@@ -0,0 +1,73 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Virtual Environment
+venv/
+ENV/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# Environment variables
+.env
+.env.local
+
+# Logs
+*.log
+logs/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# OS generated files
+.DS_Store
+.DS_Store?
+._*
+.Spotlight-V100
+.Trashes
+ehthumbs.db
+Thumbs.db
+
+# python环境相关目录
+rag/
+
+# Model files
+*.h5
+*.pt
+*.pth
+*.ckpt
+*.pkl
+*.faiss
+
+# Cache
+.cache/
+.cache/modelscope/
+.cargo
+
+# Models
+
+# LLM API keys
+*.key
+*.secret

Extra-chapter/Neo4J-Simple-Application/readme.md

@@ -0,0 +1,486 @@
+# Neo4J 简单应用
+
+## 一、什么是知识图谱
+
+**知识图谱（Knowledge Graph, KG）** 源于自然语言理解，其目标是用一种结构化的方式，来描述现实世界中的实体及其相互关系。它主要由两个核心要素构成：
+
+1.  **节点（Nodes）**：代表现实世界中的“实体”（Entities），例如一个人、一部电影、一家公司或一个具体概念。
+2.  **边（Edges）**：代表实体与实体之间的“关系”（Relations）。
+
+这些元素共同构成了一个庞大的语义网络，其基本结构可以表示为 **（实体）- [关系] -> （实体）** 的三元组（Triples）。例如，“饺子”和“哪吒2”是两个实体，“导演”就是它们之间的关系，构成一个知识三元组：（饺子）- [导演] -> （哪吒2）。
+
+## 二、知识图谱的应用
+
+知识图谱并非一个孤立的学术概念，它在工业界有着广泛且深入的应用，尤其是在需要深度结合领域知识的场景中。
+
+### 2.1 风险识别与网络分析
+
+俗话说“近朱者赤，近墨者黑”，在许多领域，一个实体的风险或属性，往往与其关联的其他实体有很强的相关性。知识图谱正是挖掘这种关联性的利器。
+
+-   **犯罪网络侦查**：公安部门可以利用通话记录、社交关系、转账流水等信息构建犯罪嫌疑人网络。在这个网络中，如果某个节点与多个已知的犯罪分子有直接或间接的联系，那么他参与犯罪的可能性就大大增加。通过分析网络中的核心人物（连接数最多的节点）和资金流向，可以有效地打击整个犯罪团伙。
+-   **信用卡反欺诈**：银行可以将申请人的信息（如电话、地址、公司）构建成一个庞大的关系网络。通过分析这张网络，可以识别出“欺诈团伙”——例如，多个申请人共享同一个联系电话或家庭住址，或者与已知的欺诈分子有紧密的社交关系。
+
+### 2.2 智能诊断与运维
+
+-   **工业设备运维**：将设备的各种“故障现象”、“故障原因”、“解决方案”和“所需零件”构建成知识图谱。当设备出现问题时，系统可以根据上报的现象，在图谱中进行推理，快速定位可能的原因，并给出维修建议，甚至可以提示维修人员需要携带哪些工具和备件，从而提高维修效率。
+-   **医疗辅助诊断**：医疗领域知识繁杂，可以通过构建“病症”、“疾病”、“检查项目”、“治疗方案”、“药品”之间的关系图谱。医生输入患者的症状后，系统可以辅助推荐需要进行的检查，并根据检查结果在图谱中推理，给出可能的诊断建议和治疗方案，帮助实现规范化诊疗。
+
+### 2.3 特定领域聊天机器人
+
+对于通用领域的开放式聊天，大语言模型（LLM）已展现出强大的能力。但在许多垂直领域，基于知识图谱的问答系统（KBQA）因其答案的准确性和可解释性，仍然具有不可替代的价值。其工作流程通常如下：
+
+1.  **意图识别**：首先判断用户提问的意图。例如，“我想买一张明天上午的故宫门票”这个问题的意图是“票务预订”。
+2.  **槽位填充 (实体抽取)**：从问题中抽取出关键信息，即“实体”。例如：`景点: 故宫`, `时间: 明天上午`, `数量: 一张`。
+3.  **知识查询**：利用抽取出的实体，在知识图谱（或数据库）中进行精确查询。
+4.  **回复生成**：将查询到的结果，通过预设的模板生成自然语言回复。
+
+这种方式虽然不如 LLM 灵活，但在机票预订、酒店查询、银行客服等业务逻辑明确的场景中，能够提供更加可靠和可控的服务。
+
+## 三、知识图谱的构建
+
+如何从海量的、非结构化的文本（如新闻、财报、医疗记录）中，自动地构建出结构化的知识图谱，是整个技术流程的核心挑战。
+
+### 3.1 经典构建流程
+
+传统的知识图谱构建过程主要依赖于两项关键的 NLP 技术：
+
+1.  **命名实体识别 (Named Entity Recognition, NER)**：从文本中识别并抽取出特定类别的实体。例如，在“英伟达发布了专为 AI 设计的 Blackwell 芯片”这句话中，识别出“英伟达”（公司）、“Blackwell”（产品）。这些被抽取的实体将成为知识图谱中的 **节点**。
+2.  **关系抽取 (Relation Extraction, RE)**：在识别出实体的基础上，进一步判断实体与实体之间存在何种语义关系。在上面的例子中，模型需要判断“英伟达”和“Blackwell”之间的关系是“发布”。这个关系将成为连接两个节点的 **边**。
+
+通过对大量文本进行这两步处理，我们就能源源不断地抽取出知识三元组，最终汇聚成一个庞大的知识图谱。
+
+### 3.2 大模型带来的革新
+
+随着大语言模型的兴起，传统的 NLP 任务流程正在被重塑。LLM 同样具备强大的实体识别和关系抽取能力，但这并不意味着对传统流程的简单替代，而是呈现出深度融合的趋势。
+
+-   **局限性与挑战**：完全依赖 LLM 会面临成本高昂、数据隐私（使用闭源 API 时）、以及“幻觉”问题，即模型可能会编造事实。
+-   **融合方案**：为了结合知识图谱的准确性和大模型的推理能力，微软提出了 GraphRAG。原理是将知识图谱作为一个可靠、可随时更新的 **外部知识库**，并基于图结构进行“子图检索”（如社区发现、路径搜索等），而非检索孤立事实。当用户提问时：
+    1.  利用模型从问题中识别出核心实体与约束。
+    2.  在图中检索与之高度相关的子图（社区/路径/邻域），获得准确且可解释的事实与关系。
+    3.  将该子图的结构化信息作为上下文，连同原始问题一起输入给大语言模型，生成基于证据的答案。
+
+## 四、图数据库：Neo4j
+
+> [Neo4j 官方文档](https://neo4j.com/docs/)
+
+知识图谱需要专门的数据库进行存储和查询，这类数据库被称为 **图数据库 (Graph Database)**。其中，与传统的关系型数据库（如 MySQL）相比，图数据库的优势在于其对“关系”的查询性能。对于需要进行多层关系遍历的复杂查询（例如，查询“我朋友的朋友”），图数据库的效果远超关系型数据库。而 **Neo4j** 就是目前比较流行的一款开源图数据库。
+
+### 4.1 核心概念
+
+Neo4j 的数据模型主要包含以下几个概念：
+
+-   **节点 (Node)**：节点是图中的基本数据单元，用于表示现实世界中的实体，例如一个人、一家公司、一本书或一个账户。在关系型数据库中，节点可以类比为表中的一行。
+
+-   **标签 (Label)**：用于为节点分类或打上“类型”标记。一个节点可以拥有一个或多个标签。例如，一个节点可以同时拥有 `:Person` 和 `:Author` 两个标签，表示这个人既是一个普通人，也是一位作者。
+
+-   **关系 (Relationship)**：这是图数据库的精髓所在，它以一种富有表现力的方式连接两个节点，并明确地定义了它们之间的联系。每个关系都具有以下特点：
+    -   **有方向**：关系总是从一个“起始节点”指向一个“结束节点”。
+    -   **有类型**：每个关系都必须有一个类型（例如 `:FRIENDS_WITH`, `:PURCHASED`），用来描述连接的性质。
+    -   **可以拥有属性**：和节点一样，关系也可以存储属性，例如，一个 `:PURCHASED` 关系可以有一个 `date` 属性来记录购买日期。
+
+-   **属性 (Property)**：属性是以键值对（Key-Value）形式存储在节点和关系上的详细信息。键是字符串，值可以是各种基本数据类型（如字符串、数字、布尔值）或它们的数组。
+
+这四个概念共同构成了一个灵活而强大的数据模型。
+
+```mermaid
+graph LR
+    A["Alice:Person {name: 'Alice'}"]
+    B["Bob:Person {name: 'Bob'}"]
+    A -- "KNOWS {since: 2020}" --> B
+```
+> 在上图中，`Alice` 和 `Bob` 是 **节点**，`:Person` 是 **标签**，`{name: 'Alice'}` 是 **属性**，`KNOWS` 则是连接它们的 **关系** 类型，而 `{since: 2020}` 是这段关系上的 **属性**。
+
+### 4.2 查询语言：Cypher
+
+Cypher 是 Neo4j 的声明式图形查询语言，它的语法灵感来源于 SQL，但针对图的特性进行了优化。通过 Cypher，我们可以用一种直观且高效的方式来查询和操作图数据。
+
+例如，要查找在电影《黑客帝国》(The Matrix) 中出演过的所有演员，可以使用以下查询：
+
+```cypher
+MATCH (actor:Person)-[:ACTED_IN]->(movie:Movie {title: 'The Matrix'})
+RETURN actor.name
+```
+
+官方的 Cypher 语法速查表（[在线版本](https://neo4j.com/docs/cypher-refcard/4.4/)）汇总了常用的命令、操作符和语法结构，可供读者快速查阅。
+
+<div align="center">
+  <img src="images/1_4_2_1.png" alt="Cypher 语法速查表" width="100%" />
+  <p>图 1.1: Cypher 语法速查表 (Cypher Refcard)</p>
+</div>
+
+### 4.3 安装与使用
+
+对于初学者和开发者，推荐以下两种主流的安装方式。
+
+1.  **Neo4j Desktop (推荐用于本地学习)**
+    -   **安装**:
+        1.  访问 [Neo4j 官网](https://neo4j.com/download/)，在 “Neo4j for Desktop” 板块点击 “Download” 按钮。
+            <div align="center">
+              <img src="images/1_4_3_1.png" alt="Neo4j Desktop 下载页面" width="100%" />
+              <p>图 1.2: 在官网点击下载</p>
+            </div>
+        2.  页面会跳转至一个注册表单。可以填写任意信息，然后点击 “Download Desktop” 按钮，浏览器将自动开始下载安装包。
+            <div align="center">
+              <img src="images/1_4_3_2.png" alt="下载前填写表单" width="100%" />
+              <p>图 1.3: 填写注册表单</p>
+            </div>
+        3.  下载完成后，双击安装文件，程序会自动进行安装。
+        4.  安装完成后首次启动，会看到许可协议界面，点击 “Continue” 即可完成最后的设置。
+            <div align="center">
+              <img src="images/1_4_3_3.png" alt="同意许可协议" width="100%" />
+              <p>图 1.4: 首次启动并同意许可协议</p>
+            </div>
+
+2.  **Docker (推荐用于服务器部署与跨平台开发)**
+    -   **安装**: 只需一行命令即可完成拉取镜像和启动容器。
+        ```bash
+        docker run \
+            --name my-neo4j \
+            -p 7474:7474 -p 7687:7687 \
+            -d \
+            -v $HOME/neo4j/data:/data \
+            -v $HOME/neo4j/logs:/logs \
+            --env NEO4J_AUTH=neo4j/password \
+            neo4j:latest
+        ```
+    -   **参数说明**:
+        -   `-p 7474:7474`: 将容器的 HTTP 端口映射到本机，用于浏览器访问。
+        -   `-p 7687:7687`: 将容器的 Bolt 驱动端口映射到本机，用于代码连接。
+        -   `-v $HOME/neo4j/data:/data`: 将数据目录挂载到本机，确保数据持久化。
+        -   `--env NEO4J_AUTH=neo4j/password`: 设置数据库的初始用户名和密码（此处为 `neo4j/password`）。
+        -   `neo4j:latest`: 使用最新的官方镜像。
+
+安装好 Neo4j 后。我们就可以学习一些 Neo4j 的基本用法了。
+
+## 五、创建并连接数据库
+
+在使用 Neo4j 进行开发时，首先需要在 Neo4j Desktop 中创建一个本地数据库实例（Instance）。这个过程非常直观。
+
+1.  **创建实例**：打开 Neo4j Desktop，在 “Local instances” 页面点击 “Create instance” 按钮。
+
+    <div align="center">
+      <img src="images/2_1_1.png" alt="创建实例" width="100%" />
+      <p>图 2.1: 点击创建实例</p>
+    </div>
+
+2.  **配置实例**：在弹出的窗口中，为实例命名（例如 `base nlp`），选择所需的 Neo4j 版本，并为默认用户 `neo4j` 设置一个能记住的密码。完成后点击 “Create”。
+
+    <div align="center">
+      <img src="images/2_1_2.png" alt="配置实例" width="100%" />
+      <p>图 2.2: 配置实例信息</p>
+    </div>
+
+3.  **启动与连接**：实例创建后会自动启动，状态显示为 “RUNNING”。此时，可以通过浏览器直接访问 `http://127.0.0.1:7474` 来打开 Neo4j Browser。
+
+    <div align="center">
+      <img src="images/2_1_3.png" alt="启动实例" width="100%" />
+      <p>图 2.3: 实例创建成功并运行</p>
+    </div>
+
+    在浏览器打开的连接界面中，使用刚刚设置的密码进行连接。
+
+    <div align="center">
+      <img src="images/2_1_4.png" alt="连接实例" width="100%" />
+      <p>图 2.4: 使用密码连接数据库</p>
+    </div>
+
+## 六、增删查改
+
+数据库操作的核心无外乎增删查改（CRUD），下面来使用 Cypher，围绕一个菜品信息图谱的场景，逐一介绍这些基本操作。
+
+### 6.1 场景设定
+
+为了方便演示，先设定好本次实践所需要用到的实体、属性和关系。
+
+-   **实体/标签 (Labels)**:
+    -   `Ingredient`: 食材，拥有 `name`, `category`（类别）, `origin`（产地）, `tags`（标签，数组）等属性。
+    -   `Dish`: 菜品，拥有 `name`, `cuisine`（菜系）等属性。
+-   **关系 (Relationships)**:
+    -   `(Dish)-[:包含]->(Ingredient)`: 表示某菜品包含某种食材，关系上可以有 `用量` 属性。
+    -   `(Dish)-[:主要食材]->(Ingredient)`: 表示某菜品的主要食材是某种食材。
+    -   `(Dish)-[:调味]->(Ingredient)`: 表示某菜品使用某种食材进行调味。
+
+### 6.2 创建 (CREATE)
+
+`CREATE` 语句用于在图中创建新的节点和关系。
+
+#### 6.2.1 创建节点
+
+创建节点的基本语法是 `CREATE (变量:标签 {属性: 值})`。
+
+-   **变量 (Variable)**: 如 `pork`，是一个临时名称，用于在同一条语句中引用该节点。如果后续不需要引用，可以省略。
+-   **标签 (Label)**: 如 `Ingredient`，用于对节点进行分类。
+-   **属性 (Properties)**: 一个包含键值对的 map/字典，用于描述节点的具体信息。
+
+最基础的创建语句包含一个临时变量（`pork`）、一个标签（`Ingredient`）和一组属性。
+
+```cypher
+CREATE (pork:Ingredient {name:'猪肉', category:'肉类', origin:'杭州'});
+```
+
+如果在创建后不需要立刻使用这个节点（例如，在同一查询中创建关系），可以省略临时变量名，这样语法更简洁。
+
+```cypher
+CREATE (:Ingredient {name:'土豆', category:'蔬菜', origin:'北京'});
+```
+
+还可以在创建节点后，使用 `RETURN` 子句立即将其返回。这对于调试或确认节点是否按预期创建非常有用。`RETURN n` 会在结果面板中直接显示刚刚创建的 `鸡蛋` 节点的信息。
+
+```cypher
+CREATE (n:Ingredient {name:'鸡蛋'}) RETURN n;
+```
+
+执行上述三条命令后，数据库中就创建了三个 `Ingredient` 类型的节点。能够通过 Neo4j Browser 的可视化界面直观地看到这些新创建的数据。
+
+<div align="center">
+  <img src="images/2_2_1_1.png" alt="创建节点后的数据库信息" width="100%" />
+  <p>图 2.5: 执行创建命令后，左侧面板显示已有 3 个 Ingredient 节点</p>
+</div>
+
+点击左侧面板中的 `Ingredient` 标签，Neo4j Browser 会自动执行 `MATCH (n:Ingredient) RETURN n LIMIT 25;` 查询，并在主窗口中展示所有食材节点。如图 2.6 所示，点击其中一个节点（如“土豆”），右侧会显示其详细属性。这里可以观察到：
+
+-   **`<id>` 字段**：这是 Neo4j 为每个节点自动生成的内部唯一标识符。
+-   **Key-Value 结构**：右侧的 “Key” 和 “Value” 两列展示了节点属性是以键值对的形式存储的。
+-   **自定义属性**：`name`、`category`、`origin` 三个字段的值与前面 `CREATE` 语句中设定的值完全一致。
+
+<div align="center">
+  <img src="images/2_2_1_2.png" alt="查询并查看节点详情" width="100%" />
+  <p>图 2.6: 查询并查看新创建的节点及其属性</p>
+</div>
+
+#### 6.2.2 创建关系
+
+关系的创建通常需要先指定关系两端的节点，然后用 `-[变量:类型 {属性}]->` 来定义关系。
+
+-   关系必须有 **方向** 和 **类型 (Type)**。
+-   小括号 `()` 用于表示节点，中括号 `[]` 用于表示关系。
+
+在实际应用中，常常需要一次性创建多个节点以及它们之间的关系。`CREATE` 语句支持通过逗号分隔，在一个查询中完成复杂图谱的构建。下面的例子将创建一个更复杂的菜品关系网络，以体现“多对多”的特性（一道菜包含多种食材，一种食材可用于多道菜）。
+
+```cypher
+CREATE
+	// 创建食材节点
+	(rousi:Ingredient {name:'猪里脊'}),
+	(muer:Ingredient {name:'木耳'}),
+	(huluobo:Ingredient {name:'胡萝卜'}),
+	(qingjiao:Ingredient {name:'青椒'}),
+	// 创建菜品节点
+	(d1:Dish {name:'鱼香肉丝', cuisine:'川菜'}),
+	(d2:Dish {name:'木须肉', cuisine:'鲁菜'}),
+	// 创建关系
+	(d1)-[:包含 {amount:'250g'}]->(rousi), (d1)-[:包含]->(muer), (d1)-[:包含]->(huluobo),
+	(d2)-[:包含 {amount:'150g'}]->(rousi), (d2)-[:包含]->(muer),
+	// 创建双向关系
+	(rousi)-[:被用于]->(d1), (muer)-[:被用于]->(d1), (huluobo)-[:被用于]->(d1),
+	(rousi)-[:被用于]->(d2), (muer)-[:被用于]->(d2);
+```
+这个查询语句做了以下几件事：
+1.  **创建了 4 个 `Ingredient` 节点**：猪里脊、木耳、胡萝卜、青椒。
+2.  **创建了 2 个 `Dish` 节点**：鱼香肉丝、木须肉。
+3.  **创建了 5 条 `包含` 关系**：从菜品指向食材。
+4.  **创建了 5 条 `被用于` 关系**：从食材指向菜品。这样既可以方便地查询“一道菜包含哪些食材”，也可以高效地反向查询“一种食材被用在了哪些菜里”。
+
+执行 `MATCH p=()-[:包含]->() RETURN p LIMIT 25;` 查询可以可视化展示所有“包含”关系。点击关系（箭头），可以在右侧看到其详细信息，例如“鱼香肉丝”到“猪里脊”的关系上，就包含了在 `CREATE` 语句中定义的 `amount: '250g'` 这一属性。
+
+<div align="center">
+  <img src="images/2_2_2_1.png" alt="同时创建节点和关系后的图谱" width="100%" />
+  <p>图 2.7: 创建关系后的图谱结构</p>
+</div>
+
+### 6.3 查询 (MATCH)
+
+`MATCH` 是 Cypher 中用于查询图数据的命令，它允许你描述你想要寻找的节点和关系的模式。
+
+#### 6.3.1 基本查询
+
+最简单的查询是匹配并返回图中的任意节点，可以使用 `LIMIT` 关键字限制返回数量，避免因数据量过大导致浏览器卡顿。
+
+```cypher
+// 匹配并返回图中的任意 25 个节点
+MATCH (n)
+RETURN n
+LIMIT 25;
+```
+
+也可以根据标签和属性进行精确匹配。
+
+```cypher
+// 匹配所有标签为 Ingredient，且名字为'猪里脊'的节点
+MATCH (n:Ingredient {name:'猪里脊'}) RETURN n;
+```
+
+#### 6.3.2 条件查询 (WHERE)
+
+`WHERE` 子句提供了更灵活的过滤能力，可以对节点的属性进行复杂的逻辑判断。
+
+例如，查询名字是'猪里脊'或'鸡蛋'的 `Ingredient` 节点。
+
+```cypher
+MATCH (n:Ingredient)
+WHERE n.name IN ['猪里脊','鸡蛋']
+RETURN n;
+```
+
+也可以使用 `AND`、`OR` 等关键字构建复合查询条件。
+
+```cypher
+// 复合条件：查询指定名称且类别为“肉类”的节点
+MATCH (n:Ingredient)
+WHERE n.name IN ['猪肉', '猪里脊', '鸡蛋'] AND n.category = '肉类'
+RETURN n;
+```
+
+#### 6.3.3 返回指定属性
+
+默认情况下，`RETURN n` 会返回整个节点对象。也可以只返回节点的特定属性，并使用 `AS` 为返回的列起别名，使结果更具可读性。
+
+```cypher
+MATCH (n:Ingredient)
+WHERE n.name IN ['猪里脊','鸡蛋']
+RETURN n.name AS 食材名称, n.category AS 类别;
+```
+
+#### 6.3.4 关联查询
+
+图数据库最强大的地方在于对关系的查询。例如，可以一次性查询“鱼香肉丝”和“木须肉”分别包含了哪些食材。
+
+```cypher
+MATCH (d:Dish)-[:包含]->(i:Ingredient)
+WHERE d.name IN ['鱼香肉丝', '木须肉']
+RETURN d.name AS 菜品, collect(i.name) AS 食材列表;
+```
+> `collect()` 是一个聚合函数，可以将匹配到的多个同类结果（这里是食材名称 `i.name`）收集到一个列表中。
+
+#### 6.3.5 查询并创建 (MATCH + CREATE)
+
+在实际应用中，一个常见的操作是先找到图中已经存在的节点，然后为它们添加新的关系。这可以通过组合使用 `MATCH` 和 `CREATE` 来实现。
+
+例如，我们已经创建了“鱼香肉丝”和“猪里脊”，现在想为它们添加一条“主要食材”的关系。
+
+```cypher
+MATCH
+    (d:Dish {name:'鱼香肉丝'}),
+    (i:Ingredient {name:'猪里脊'})
+MERGE
+    (d)-[r:主要食材]->(i)
+RETURN d, i, r;
+```
+> 这个模式确保了是在已有的、正确的实体之间建立关联，并通过 `MERGE` 避免重复的关系。
+
+#### 6.3.6 排序 (ORDER BY)
+
+可以使用 `ORDER BY` 子句对返回的结果进行排序。默认是升序 (`ASC`)，也可以指定为降序 (`DESC`)。
+
+```cypher
+// 查询所有食材，并按名称升序排序
+MATCH (i:Ingredient)
+RETURN i.name, i.category
+ORDER BY i.name ASC;
+```
+
+### 6.4 更新 (SET & MERGE)
+
+#### 6.4.1 更新属性 (SET)
+
+`SET` 语句用于修改或添加节点/关系的属性。它必须和 `MATCH` 配合使用，先找到要更新的实体，再进行修改。
+
+```cypher
+MATCH (i:Ingredient {name:'猪肉'})
+SET
+    i.is_frozen = true,
+    i.origin = '金华'
+RETURN i;
+```
+
+#### 6.4.2 插入或更新 (MERGE)
+
+在构建知识图谱时，经常遇到这样的场景：如果某个节点已存在，则更新其属性；如果不存在，则创建它。`MERGE` 语句就可以解决这个问题。
+
+`MERGE` 会根据你提供的模式在图中查找，如果找到匹配项，则执行 `ON MATCH` 部分；如果未找到，则执行 `ON CREATE` 部分，从而避免了重复创建实体。
+
+```cypher
+// 查找名为'大蒜'的 Ingredient 节点
+MERGE (n:Ingredient {name: '大蒜'})
+// 如果不存在，则创建该节点，并设置创建时间和初始库存
+ON CREATE SET
+    n.created = timestamp(),
+    n.stock = 100
+// 如果已存在，则更新其库存、访问次数和访问时间
+ON MATCH SET
+  n.stock = coalesce(n.stock, 0) - 1,
+  n.counter = coalesce(n.counter, 0) + 1,
+  n.accessTime = timestamp()
+RETURN n;
+```
+> `coalesce(property, defaultValue)` 是一个非常有用的函数，它会检查属性 `property` 是否存在，如果存在则返回其值，否则返回 `defaultValue`。
+
+### 6.5 删除 (DELETE & REMOVE)
+
+#### 6.5.1 删除属性 (REMOVE)
+
+`REMOVE` 用于移除节点或关系上的某个属性。在下面的例子中，先用 `MATCH` 找到名为“大蒜”的节点，然后移除由 `MERGE` 命令在创建它时添加的 `created` 属性。
+
+```cypher
+MATCH (i:Ingredient {name:'大蒜'})
+REMOVE i.created
+RETURN i;
+```
+
+#### 6.5.2 删除节点和关系 (DELETE)
+
+`DELETE` 用于删除节点和关系。但需要 **特别注意**：Neo4j 不允许直接删除一个还存在关联关系的节点。你必须先删除关系，才能删除节点。
+
+```cypher
+// 错误示范：如果'大蒜'还有关系连着，这条语句会报错
+MATCH (i:Ingredient {name:'大蒜'})
+DELETE i;
+```
+
+正确的做法有两种。第一种是先手动删除与节点相关的所有关系，然后再删除节点本身。
+
+```cypher
+// 正确做法 1：先删除关系，再删除节点
+MATCH (i:Ingredient {name:'大蒜'})-[r]-() // 匹配与'大蒜'相连的任意关系
+DELETE r, i; // 先删除关系 r，再删除节点 i
+```
+
+第二种做法更简洁，也是官方推荐的方式：使用 `DETACH DELETE`。它会自动删除指定节点以及所有与它直接相连的关系。
+
+```cypher
+// 正确做法 2：使用 DETACH DELETE (推荐)
+MATCH (i:Ingredient {name:'大蒜'})
+DETACH DELETE i;
+```
+
+此外，还可以通过节点的内部 ID 进行精确查找和删除。每个节点都有一个由 Neo4j 自动分配的唯一 ID，可以通过 `id()` 函数获取。
+
+```cypher
+// 假设我们通过查询得知“大蒜”的 ID 为 5
+MATCH (i:Ingredient)
+WHERE id(i) = 5
+DETACH DELETE i;
+```
+
+#### 6.5.3 清空数据库
+
+如果想删除数据库中的所有节点和关系，可以使用以下命令：
+
+```cypher
+// 匹配所有节点 n
+MATCH (n)
+// 强制删除节点 n 及其所有关系
+DETACH DELETE n;
+```
+
+#### 6.5.4 软删除
+
+在生产环境中，直接从数据库中物理删除（`DELETE`）数据是一种高风险操作。一种更安全、更常见的做法是“软删除”。软删除并非真的将数据移除，而是通过 `SET` 命令为其添加一个状态属性，将其标记为“已删除”或“不活跃”。
+
+```cypher
+// 将“木耳”标记为不活跃
+MATCH (i:Ingredient {name:'木耳'})
+SET i.is_active = false;
+```
+这样，在后续的查询中，只需要增加一个 `WHERE i.is_active = true` 的过滤条件，就能只使用那些“活跃”的数据，而被软删除的数据依然保留在数据库中，以备审计或恢复。
+
+> 删除操作是高风险行为，尤其是在生产环境中。执行前请务必确认操作对象和范围，并做好数据备份。

Extra-chapter/PowerRAG-SDK-Text-QA/code/.env.example

@@ -0,0 +1,26 @@
+# PowerRAG (RAGFlow) SDK demo config
+
+# SDK endpoint (from your docker-compose env: SVR_HTTP_PORT=9380)
+RAGFLOW_BASE_URL=http://127.0.0.1:9380
+
+# SDK API key (format: ragflow-...; created via /v1/api/new_token)
+RAGFLOW_API_KEY=ragflow-REPLACE_ME
+
+# Optional: override dataset name created by the demo
+RAGFLOW_DATASET_NAME=powerrag_text_qa_demo
+
+# Optional: override embedding model for dataset creation (recommended to leave empty and use tenant default)
+# Format: <model>@<factory>
+# Example:
+# RAGFLOW_EMBEDDING_MODEL=text-embedding-3-small@OpenAI
+RAGFLOW_EMBEDDING_MODEL=
+
+# -----------------------------
+# Optional: embedding provider config (used by the README “API 配置 embedding” steps)
+# -----------------------------
+
+# Use the factory/model name shown by your PowerRAG UI/API.
+EMB_FACTORY=REPLACE_ME
+EMB_MODEL=REPLACE_ME
+EMB_API_BASE=REPLACE_ME
+EMB_API_KEY=REPLACE_ME

Extra-chapter/PowerRAG-SDK-Text-QA/code/config.py

@@ -0,0 +1,46 @@
+"""
+PowerRAG (RAGFlow) SDK Demo configuration.
+
+This module follows the `code/` directory convention:
+- Provide a small config object
+- Load `.env` automatically (if present)
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def _bool_env(name: str, default: bool = False) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    raw = raw.strip().lower()
+    if raw in {"1", "true", "yes", "y", "on"}:
+        return True
+    if raw in {"0", "false", "no", "n", "off"}:
+        return False
+    return default
+
+
+@dataclass(frozen=True)
+class PowerRAGDemoConfig:
+    base_url: str = os.getenv("RAGFLOW_BASE_URL", "http://127.0.0.1:9380").strip()
+    api_key: str = os.getenv("RAGFLOW_API_KEY", "").strip()
+    dataset_name: str = os.getenv("RAGFLOW_DATASET_NAME", "powerrag_text_qa_demo").strip()
+    embedding_model: str = os.getenv("RAGFLOW_EMBEDDING_MODEL", "").strip()
+
+    top_k: int = int(os.getenv("RAGFLOW_TOP_K", "5"))
+    candidate_k: int = int(os.getenv("RAGFLOW_CANDIDATE_K", "1024"))
+    similarity_threshold: float = float(os.getenv("RAGFLOW_SIMILARITY_THRESHOLD", "0.2"))
+    vector_similarity_weight: float = float(os.getenv("RAGFLOW_VECTOR_SIMILARITY_WEIGHT", "0.3"))
+    keyword: bool = _bool_env("RAGFLOW_KEYWORD", False)
+
+
+DEFAULT_CONFIG = PowerRAGDemoConfig()
+

Extra-chapter/PowerRAG-SDK-Text-QA/code/main.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from typing import Any
+
+from config import DEFAULT_CONFIG
+
+
+def _env(name: str, default: str | None = None) -> str | None:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return default
+    return value.strip()
+
+
+def _require(value: str | None, hint: str) -> str:
+    if value is None or value.strip() == "":
+        raise SystemExit(hint)
+    return value.strip()
+
+
+def _read_bytes(path: Path) -> bytes:
+    try:
+        return path.read_bytes()
+    except FileNotFoundError:
+        raise SystemExit(f"File not found: {path}")
+
+
+def _safe_get(obj: Any, attr: str, default: Any = None) -> Any:
+    try:
+        return getattr(obj, attr)
+    except Exception:
+        return default
+
+
+def main(argv: list[str]) -> int:
+    parser = argparse.ArgumentParser(
+        description="PowerRAG (RAGFlow) SDK demo: upload Markdown, parse, retrieve top-k chunks.",
+    )
+    parser.add_argument("--file", type=Path, required=True, help="Markdown file path, e.g. ./data/sample.md")
+    parser.add_argument("--question", type=str, required=True, help="User question for retrieval")
+    parser.add_argument("--top-k", type=int, default=DEFAULT_CONFIG.top_k, help="How many chunks to return (mapped to page_size)")
+    parser.add_argument(
+        "--embedding-model",
+        type=str,
+        default=DEFAULT_CONFIG.embedding_model or _env("RAGFLOW_EMBEDDING_MODEL"),
+        help=(
+            "Embedding model string in '<model>@<factory>' format. "
+            "If omitted, server tenant default is used."
+        ),
+    )
+    parser.add_argument("--candidate-k", type=int, default=DEFAULT_CONFIG.candidate_k, help="RAGFlow.retrieve(top_k=...) candidate pool size")
+    parser.add_argument("--similarity-threshold", type=float, default=DEFAULT_CONFIG.similarity_threshold, help="Filter chunks below this similarity")
+    parser.add_argument("--vector-similarity-weight", type=float, default=DEFAULT_CONFIG.vector_similarity_weight, help="Weight of vector similarity in hybrid score")
+    parser.add_argument("--keyword", action="store_true", default=DEFAULT_CONFIG.keyword, help="Enable keyword matching (hybrid retrieval)")
+    parser.add_argument("--dataset-name", type=str, default=DEFAULT_CONFIG.dataset_name, help="Dataset name to create")
+    parser.add_argument(
+        "--base-url",
+        type=str,
+        default=DEFAULT_CONFIG.base_url or _env("RAGFLOW_BASE_URL") or _env("POWERRAG_BASE_URL") or _env("BASE_URL"),
+        help="RAGFlow/PowerRAG base_url (or env RAGFLOW_BASE_URL / POWERRAG_BASE_URL / BASE_URL)",
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=DEFAULT_CONFIG.api_key or _env("RAGFLOW_API_KEY") or _env("POWERRAG_API_KEY") or _env("API_KEY"),
+        help="RAGFlow/PowerRAG api_key (or env RAGFLOW_API_KEY / POWERRAG_API_KEY / API_KEY)",
+    )
+    parser.add_argument("--cleanup", action="store_true", help="Delete created dataset after finishing")
+
+    args = parser.parse_args(argv)
+
+    base_url = _require(args.base_url, "Missing base_url. Use --base-url or set env RAGFLOW_BASE_URL.")
+    api_key = _require(args.api_key, "Missing api_key. Use --api-key or set env RAGFLOW_API_KEY.")
+
+    if args.top_k <= 0:
+        raise SystemExit("--top-k must be > 0")
+    if args.candidate_k <= 0:
+        raise SystemExit("--candidate-k must be > 0")
+
+    blob = _read_bytes(args.file)
+    display_name = args.file.name
+    if not display_name.lower().endswith(".md"):
+        display_name = f"{display_name}.md"
+
+    try:
+        from ragflow_sdk import RAGFlow  # type: ignore
+    except Exception as e:
+        raise SystemExit(
+            "Failed to import ragflow_sdk. Install dependencies first:\n"
+            "  pip install -r requirements.txt\n"
+            f"Original error: {e}"
+        )
+
+    rag = RAGFlow(api_key=api_key, base_url=base_url)
+
+    dataset_kwargs: dict[str, Any] = {"name": args.dataset_name}
+    if args.embedding_model:
+        dataset_kwargs["embedding_model"] = args.embedding_model
+    dataset = rag.create_dataset(**dataset_kwargs)
+    try:
+        docs = dataset.upload_documents([{"display_name": display_name, "blob": blob}])
+        if not docs:
+            raise SystemExit("Upload succeeded but no document returned by SDK.")
+        doc = docs[0]
+
+        parse_results = dataset.parse_documents([doc.id])
+        # parse_results: list[tuple[doc_id, status, success_count, failure_count]] (per API ref)
+        print("Parse results:")
+        print(parse_results)
+        if parse_results and isinstance(parse_results, list):
+            statuses = {r[1] for r in parse_results if isinstance(r, (list, tuple)) and len(r) >= 2}
+            if statuses and statuses != {"DONE"}:
+                raise SystemExit(
+                    "Document parsing failed (status not DONE). "
+                    "Most common cause is missing/unauthorized embedding model.\n"
+                    "Try:\n"
+                    "  - set tenant default embedding model in UI or via /v1/user/set_tenant_info, OR\n"
+                    "  - rerun with --embedding-model '<model>@<factory>' (must be supported & configured for the tenant)\n"
+                    "If it still fails, check PowerRAG logs inside the container (task executor) for the detailed error.\n"
+                )
+
+        chunks = rag.retrieve(
+            question=args.question,
+            dataset_ids=[dataset.id],
+            document_ids=[doc.id],
+            page=1,
+            page_size=args.top_k,
+            similarity_threshold=args.similarity_threshold,
+            vector_similarity_weight=args.vector_similarity_weight,
+            top_k=args.candidate_k,
+            keyword=args.keyword,
+        )
+
+        print("\nRetrieved chunks:")
+        if not chunks:
+            print("(empty)")
+            return 0
+
+        for i, c in enumerate(chunks, start=1):
+            similarity = _safe_get(c, "similarity")
+            vector_similarity = _safe_get(c, "vector_similarity")
+            term_similarity = _safe_get(c, "term_similarity")
+            content = _safe_get(c, "content", "")
+            content_preview = (content or "").strip().replace("\n", " ")
+            if len(content_preview) > 260:
+                content_preview = content_preview[:260] + "…"
+            print(f"{i:02d}. similarity={similarity} vector={vector_similarity} term={term_similarity}")
+            print(f"    {content_preview}")
+
+        return 0
+    finally:
+        if args.cleanup:
+            try:
+                rag.delete_datasets(ids=[dataset.id])
+            except Exception as e:
+                print(f"Warning: failed to cleanup dataset {dataset.id}: {e}", file=sys.stderr)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))

Extra-chapter/PowerRAG-SDK-Text-QA/code/requirements.txt

@@ -0,0 +1,3 @@
+ragflow-sdk
+python-dotenv
+

Extra-chapter/PowerRAG-SDK-Text-QA/data/questions.txt

@@ -0,0 +1,5 @@
+1) 这个 demo 的验收标准是什么？
+2) 餐厅排队系统里，如果顾客过号，通常怎么处理？
+3) 已发货未签收的退款规则是什么？
+4) 如何估算排队等待时间？
+

Extra-chapter/PowerRAG-SDK-Text-QA/data/sample.md

@@ -0,0 +1,40 @@
+# PowerRAG 文本问答 Demo · 示例文档
+
+## 1. 项目背景
+
+本示例用于演示：上传一份 Markdown 文档 → 服务端自动解析与分块 → 基于问题检索相关 chunks。
+
+## 2. 关键概念
+
+- **分块（Chunk）**：把长文切成多个小段，便于向量化与检索。
+- **向量检索（Vector Search）**：把文本映射到向量空间，通过相似度找到相关片段。
+- **Top-k**：返回最相关的 k 个片段。
+
+## 3. 规则与约束
+
+1) 只有当“检索到的 chunks 与问题语义相关”时，才算成功。  
+2) 本 demo 不要求大模型生成最终回答（可选）。  
+
+## 4. 示例内容：餐厅排队系统
+
+我们要做一个餐厅排队系统，核心流程如下：
+
+1. 顾客在前台取号，系统生成排队号（例如 A001）。
+2. 服务员在就餐区空位出现时叫号，顾客到号后入座。
+3. 如果顾客过号，可选择重新排队或延后若干位。
+4. 系统需要支持查询当前排队情况，以及某个号码前面还有多少人。
+
+### 4.1 常见问题
+
+- “过号后怎么处理？”：可以延后或重新取号，策略由门店决定。
+- “如何估算等待时间？”：可以用平均翻台时间 × 前方人数估算。
+- “如何处理多人同时取号？”：需要对取号操作加锁或用原子自增保证顺序。
+
+## 5. 示例内容：退款规则
+
+退款规则如下：
+
+- 未发货：可全额退款。
+- 已发货未签收：可申请退款，但需要承担退货运费。
+- 已签收：7 天内可退货退款；超过 7 天视情况处理。
+

Extra-chapter/PowerRAG-SDK-Text-QA/readme.md

@@ -0,0 +1,404 @@
+# PowerRAG SDK 文本问答检索 Demo
+
+## 一、这篇专题要解决什么问题？
+
+很多同学做 RAG 时会先把注意力放在“怎么让大模型回答得更像人”。但只要检索没找对上下文，生成再花哨也只是“把错讲得更顺”。
+
+这个专题做一件更朴素、也更值得先掌握的事：
+
+> **只做检索，不做生成。**
+
+你会把一份 Markdown 文档交给服务端，让服务端完成解析、切分、向量化，然后用问题去做 Top‑K 检索，拿回最相关的原文片段（chunks）。
+
+**验收标准也很直接**：Top‑K chunks 是否与问题语义相关（不要求最终答案）。
+
+本专题目录结构：
+
+- `readme.md`：本文（教学文档）
+- `images/`：配图
+- `code/`：可运行脚本与配置（`main.py`、`config.py`、`.env.example`、`requirements.txt`）
+- `data/`：可复现样例数据（`sample.md` + `questions.txt`）
+
+---
+
+## 二、技术方案：从 Markdown 到 Top‑K chunks（图文讲清楚）
+
+下面这张图展示了端到端链路，也基本对应 `code/main.py` 的执行顺序。
+
+<div align="center">
+  <img src="images/10_1_1.webp" alt="端到端流程图：上传→解析/切分→向量化→Top-K 检索" width="100%" />
+  <p>图 10.1: 端到端流程（本 demo 只验收检索结果，不要求生成最终回答）</p>
+</div>
+
+为了避免“看完图还是不知道自己要做什么”，这里把图 10.1 的关键节点按顺序讲清楚（你可以边对照图边往下读）：
+
+**（1）本地输入：Markdown 文档**
+
+你可以直接用本专题提供的 `data/sample.md`。这份文件故意写得短：包含“排队规则”和“退款规则”，方便你用不同问题去验证检索是否命中。
+
+**（2）Upload：上传到 dataset**
+
+上传不是“把文本发过去就结束”，它的意义在于：服务端要把这份文档纳入某个 **dataset**（容器）里，后续切分出来的 chunks、embedding、索引都挂在这个容器下面。
+
+**（3）Parse/Chunk：解析 + 切分**
+
+这一步会把 Markdown 解析成可检索的文本结构，并按服务端策略切成多个 chunk。
+
+> ⚠️ 图里标了一个常见失败点：如果你的 tenant 没有配置默认 embedding（`embd_id` 为空或未授权），解析任务可能直接 FAIL。
+
+**（4）Embedding：向量化**
+
+每个 chunk 会被映射成向量（embedding）。这一步是向量检索的前提——没有向量，后面就谈不上“语义相似”。
+
+**（5）写入向量库/索引**
+
+chunk + embedding 会写入向量索引（图里叫 Vector Store / Index）。
+
+**（6）Retrieve Top‑K：检索并返回 chunks**
+
+输入一个问题（question），服务端从索引里找出最相关的 K 个 chunk，并把这些原文片段返回给你。本 demo 的验收就看这里：**返回的 chunks 是否包含你期望的规则段落**。
+
+---
+
+到这里，你应该已经能把这条链路从头到尾“顺着说一遍”了：
+
+> 文档上传 → 服务端解析/切分/向量化 → 写入索引 → 问题检索 → 返回 Top‑K chunks。
+
+但很多初学者还有一个常见困惑：**这些名词到底对应什么对象？我拿到的结果到底是谁？**
+
+所以下面我们换一个视角：不再看“流程”，而是看“对象之间的关系”。
+
+---
+
+再看图 10.2（对象关系）。这张图的目的只有一个：把“你上传的文件”和“检索返回的结果”彻底区分开。
+
+很多同学第一次用 RAG 平台 SDK，会把这些概念混在一起。你只要记住：
+
+- **dataset**：容器（装很多文档）
+- **document**：你上传的那份文件
+- **chunk**：文档切分出来的文本片段（检索返回的就是它）
+
+<div align="center">
+  <img src="images/10_1_2.webp" alt="对象关系图：dataset-document-chunk-embedding 与 Top-K 返回" width="100%" />
+  <p>图 10.2: 对象关系与返回结构（检索返回的核心对象是 chunk）</p>
+</div>
+
+图 10.2 里最容易忽略、但最关键的一点是：**检索返回的是 chunk，不是 document。**
+
+- document 是“你上传的整份文件”
+- chunk 是“切分后的片段”，它才是检索、重排、压缩、最终拼上下文的基本单位
+
+所以你在终端里看到的 Top‑K 结果，应该是一段段原文片段，而不是整篇 Markdown。
+
+> 💡 小白自检：我怎么判断“这段 chunk 就是我想要的那段”？
+>
+> 很简单：用你自己的语言把问题再复述一遍，然后在返回的 chunk 里找“能直接支撑答案的原文句子”。
+> 例如你问“已发货未签收能不能退款”，chunk 里应当出现“已发货未签收：可申请退款，但需要承担退货运费”这一类关键句。
+
+---
+## 三、实现思路：从零写一版“最小检索脚本”（带代码块）
+
+先给一个“最小骨架”（你可以把它当作伪代码，但它基本就是 `code/main.py` 的主干）：
+
+```python
+rag = RAGFlow(api_key=..., base_url=...)
+
+# 1) 创建 dataset（容器）
+dataset = rag.create_dataset(name=...)
+
+# 2) 上传文档（拿到 doc.id）
+doc = dataset.upload_documents([{...}])[0]
+
+# 3) 解析/切分/向量化（失败大多发生在这里）
+parse_results = dataset.parse_documents([doc.id])
+
+# 4) 检索 Top-K chunks（验收点）
+chunks = rag.retrieve(question=..., dataset_ids=[dataset.id], document_ids=[doc.id], page_size=top_k)
+```
+
+下面把每一步展开讲清楚（并配上代码片段）。
+
+### 3.1 参数与配置：先让脚本可复现
+
+先从命令行参数入手，理解脚本“能调什么”。`code/main.py` 里最常用的是这几个：
+
+```python
+parser.add_argument("--file", type=Path, required=True)
+parser.add_argument("--question", type=str, required=True)
+parser.add_argument("--top-k", type=int, default=DEFAULT_CONFIG.top_k)
+parser.add_argument("--dataset-name", type=str, default=DEFAULT_CONFIG.dataset_name)
+parser.add_argument("--base-url", type=str, default=DEFAULT_CONFIG.base_url)
+parser.add_argument("--api-key", type=str, default=DEFAULT_CONFIG.api_key)
+```
+
+- `--file`：你要上传哪份 Markdown
+- `--question`：你想验证的提问
+- `--top-k`：返回多少个 chunk
+- `--dataset-name`：本次创建/使用的数据集名字
+- `--base-url/--api-key`：PowerRAG 服务端地址与 SDK token
+
+这几个参数足够让你完成“换文档、换问题、调 Top‑K、连不同服务端”这四类最常见实验。
+
+> 💡 小白自检：为什么这里既支持命令行参数，又支持 `.env`？
+>
+> 因为这两种场景都很常见：
+>
+> - 你本地调试时，喜欢用 `.env` 固定住 base_url/api_key
+> - 你改参数做实验时，喜欢命令行直接覆盖（不用反复改文件）
+
+### 3.2 初始化 SDK：先连上再说
+
+```python
+from ragflow_sdk import RAGFlow
+
+rag = RAGFlow(api_key=api_key, base_url=base_url)
+```
+
+这里没有花活：就是把请求的 base_url 和 token 配好。
+
+### 3.3 创建 dataset：把文档放进“一个篮子里”
+
+```python
+dataset_kwargs = {"name": args.dataset_name}
+if args.embedding_model:
+    dataset_kwargs["embedding_model"] = args.embedding_model
+dataset = rag.create_dataset(**dataset_kwargs)
+```
+
+为什么要先有 dataset？因为“上传/解析/检索”都需要一个边界。
+你不希望每次检索都在整个租户的所有文档里搜；你希望“只在这次实验的文档集合里搜”。
+
+> 💡 小白自检：能不能不建 dataset，直接上传然后检索？
+>
+> 取决于平台能力。但在 PowerRAG/RAGFlow 这类系统里，dataset 是“组织边界”。
+> 没有边界，检索要么全库搜（不可控），要么压根没有地方挂索引。
+
+### 3.4 上传 document：得到 doc.id，后面都靠它
+
+```python
+docs = dataset.upload_documents([
+    {"display_name": display_name, "blob": blob}
+])
+doc = docs[0]
+```
+
+上传成功后，SDK 会返回一个 document 对象（至少包含 `doc.id`）。
+后续的 parse 和 retrieve 都要用它来限定范围。
+
+> 💡 小白自检：为什么要限定 `document_ids=[doc.id]`？
+>
+> 因为你这次实验只关心“这份文档”的检索效果。
+> 如果不限定，dataset 里有多份文档时，你可能会检索到别的文档的 chunk，导致结果看起来“跑偏”。
+
+### 3.5 解析 / 切分 / 向量化：最容易踩坑的一步
+
+```python
+parse_results = dataset.parse_documents([doc.id])
+print("Parse results:")
+print(parse_results)
+```
+
+脚本会把 parse 的状态打印出来，并且做了一个很直接的判断：
+
+```python
+statuses = {r[1] for r in parse_results if isinstance(r, (list, tuple)) and len(r) >= 2}
+if statuses and statuses != {"DONE"}:
+    raise SystemExit("Document parsing failed (status not DONE)...")
+```
+
+你可以把它理解为“验收关卡”：
+
+- **DONE**：说明服务端已经把文档切成 chunk，并完成（或至少开始完成）向量化与索引写入
+- **FAIL/其他状态**：先别着急改代码，优先排查 tenant 默认 embedding
+
+> 经验：`Model(@None) not authorized` 基本就是在提示“默认 embedding 没配/没权限”。
+
+> 💡 小白自检：为什么 embedding 配置会影响“解析（parse）”？
+>
+> 因为这里的 parse 往往不是“纯语法解析 Markdown”，而是一条“解析 → 切分 → 向量化 → 写索引”的流水线任务。
+> embedding 不可用时，流水线中途失败，平台就会把整个任务标为 FAIL。
+
+### 3.6 检索 Top‑K：你真正要验收的结果
+
+```python
+chunks = rag.retrieve(
+    question=args.question,
+    dataset_ids=[dataset.id],
+    document_ids=[doc.id],
+    page=1,
+    page_size=args.top_k,
+    similarity_threshold=args.similarity_threshold,
+    vector_similarity_weight=args.vector_similarity_weight,
+    top_k=args.candidate_k,
+    keyword=args.keyword,
+)
+```
+
+这里有两个点值得你留意（也是很多人调参的入口）：
+
+- `page_size=args.top_k`：你最终想看多少条 chunk
+- `similarity_threshold`：太高会过滤掉结果导致空，太低会混进无关段落
+
+最后脚本会把每条 chunk 的内容预览打印出来：
+
+```python
+for i, c in enumerate(chunks, start=1):
+    content = _safe_get(c, "content", "")
+    print(f"{i:02d}. {content[:260]}")
+```
+
+你要做的“人工验收”也很简单：看看这几段文字是不是回答问题所需的那几段原文。
+
+> 💡 小白自检：Top‑K 是不是越大越好？
+>
+> 不是。Top‑K 太大容易把无关 chunk 混进来；太小又可能漏掉关键段落。
+> 教学 demo 里一般用 3~8 都够用。
+
+---
+
+### 3.7 先跑通一次（最短路径）
+
+> ⚠️ 注意：解析/向量化依赖 embedding。如果你的 tenant 没有配置默认 embedding（`embd_id` 为空或未授权），解析阶段会 FAIL。不要先怀疑 Python。
+
+```bash
+# 1) 安装依赖
+cd Extra-chapter/PowerRAG-SDK-Text-QA/code
+python -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+
+# 2) 配置 .env（在 code/ 目录下）
+cp .env.example .env
+
+# 3) 回到专题根目录运行（data/ 路径更直观）
+cd ..
+python code/main.py \
+  --file data/sample.md \
+  --question "已发货未签收的退款规则是什么？" \
+  --top-k 5 \
+  --cleanup
+```
+
+你会看到两段关键输出：
+
+1. `Parse results`：解析/分块状态（期望 `DONE`）
+2. `Retrieved chunks`：Top‑K chunks 的内容预览
+
+---
+
+---
+
+## 四、经验总结与坑点（把时间花在对的地方）
+
+很多时候问题不在“你写的 Python”，而在“服务端是不是已经把 embedding 产出来了”。
+
+<div align="center">
+  <img src="images/10_1_3.webp" alt="简化时序图：上传→解析→写入索引→Top-K 检索→返回 chunks" width="100%" />
+  <p>图 10.3: code/main.py 与服务端 API 的交互顺序（简化版）</p>
+</div>
+
+如果你只记一条顺序，就记这句：
+
+> **先上传 → 再解析（产出 chunk+embedding）→ 最后检索（返回 chunk）**
+
+很多“为什么检索不到”的问题，本质是解析还没成功，索引里根本没有向量。
+
+---
+
+### 4.1 Parse results 是 FAIL
+
+优先检查 tenant 的默认 embedding（`embd_id`）是否已配置且可用。典型错误：
+
+- `Model(@None) not authorized`
+- `Parse results: ... FAIL ...`
+
+如果已经配置仍失败，直接看 task executor 日志最省时间：
+
+```bash
+docker exec powerrag-powerrag-1 sh -lc 'tail -n 200 /ragflow/logs/task_executor_* | tail -n 200'
+```
+
+### 4.2 401/403：token 类型搞混
+
+PowerRAG 常见会同时出现两类 token：
+
+- Web 层 `AUTH`（用于 `/v1/*`）
+- SDK 的 `ragflow-...` token（用于 `/api/v1/*`，通常写在 `Authorization: Bearer <ragflow-...>`）
+
+如果你看到 401/403，先确认 token 类型和接口前缀是否匹配。
+
+---
+
+---
+
+## 附录：用 API 配默认 embedding + 生成 ragflow token（重操作区）
+
+> 这部分是“环境/账号/服务端配置”，放到附录，避免主线被淹没。
+
+### A1. 用 API 配好 embedding（通用）
+
+这一步需要一个 Web 层的 `AUTH`（`/v1/*` 使用），它和 SDK 的 `ragflow-...` key 不是一回事。
+
+你可以把 embedding 配置写进 `.env`（见 `.env.example` 的 `EMB_*`），下面命令会读取 `EMB_FACTORY/EMB_MODEL/EMB_API_BASE/EMB_API_KEY`。
+
+#### A1.1 获取 `AUTH`（注册并从响应头拿 Authorization）
+
+PowerRAG 的 `/v1/user/register` 要求 password 先用服务端的 RSA public key 加密。最省事的方式是在容器内调用它自带的加密函数：
+
+```bash
+BASE_URL="http://127.0.0.1:9380"
+
+ENC_PW="$(docker exec powerrag-powerrag-1 sh -lc 'python - <<"PY"\nfrom api.utils.crypt import crypt\nprint(crypt("powerrag"))\nPY')"
+
+EMAIL="powerrag.demo.$(date +%s)@example.com"
+AUTH="$(curl -sS -D - -o /dev/null -X POST "$BASE_URL/v1/user/register" \
+  -H 'Content-Type: application/json' \
+  -d "{\"nickname\":\"demo\",\"email\":\"$EMAIL\",\"password\":\"$ENC_PW\"}" \
+| awk 'BEGIN{IGNORECASE=1} /^authorization:/{print $2}' | tr -d '\r')"
+```
+
+#### A1.2 绑定 embedding 的外部 API
+
+> 注意：`max_tokens` 需要显式传，否则可能报数据库字段错误。
+
+```bash
+curl -sS -X POST "$BASE_URL/v1/llm/add_llm" \
+  -H "Authorization: $AUTH" \
+  -H 'Content-Type: application/json' \
+  -d '{
+    "llm_factory": "'"${EMB_FACTORY}"'",
+    "model_type": "embedding",
+    "llm_name": "'"${EMB_MODEL}"'",
+    "api_base": "'"${EMB_API_BASE}"'",
+    "api_key": "'"${EMB_API_KEY}"'",
+    "max_tokens": 8192
+  }'
+```
+
+#### A1.3 设置 tenant 默认 `embd_id`
+
+```bash
+TENANT_ID="$(curl -sS -H "Authorization: $AUTH" "$BASE_URL/v1/user/tenant_info" | python -c 'import sys,json; print(json.load(sys.stdin)["data"]["tenant_id"])')"
+
+curl -sS -X POST "$BASE_URL/v1/user/set_tenant_info" \
+  -H "Authorization: $AUTH" \
+  -H 'Content-Type: application/json' \
+  -d "{\"tenant_id\":\"$TENANT_ID\",\"llm_id\":\"\",\"embd_id\":\"${EMB_MODEL}@${EMB_FACTORY}\",\"asr_id\":\"\",\"img2txt_id\":\"\"}"
+```
+
+### A2. 生成 SDK 的 `ragflow-...` api_key
+
+SDK 接口在 `/api/v1/*`，它不认 `AUTH`，需要 `ragflow-...` 这种 token（放在 header：`Authorization: Bearer <ragflow-...>`）。
+
+用 `AUTH` 创建一个 SDK key：
+
+```bash
+DIALOG_ID="$(python -c 'import uuid; print(uuid.uuid4().hex)')"
+API_KEY="$(curl -sS -X POST "$BASE_URL/v1/api/new_token" \
+  -H "Authorization: $AUTH" \
+  -H 'Content-Type: application/json' \
+  -d "{\"dialog_id\":\"$DIALOG_ID\"}" \
+| python -c 'import sys,json; print(json.load(sys.stdin)["data"]["token"])')"
+
+echo "$API_KEY"
+```

Extra-chapter/README.md

@@ -0,0 +1,149 @@
+<div align="center">
+
+  <h2>🚀 All-in-RAG · Extra Chapter</h2>
+
+  <p><em>主教程之外的「知识拓展与社区实践」专区</em></p>
+
+</div>
+
+---
+
+## 📖 Extra Chapter 是什么？
+
+`Extra-chapter` 目录用于存放 **不直接属于主线章节，但对 RAG / LLM 应用非常有价值的补充内容**。这些内容可以是你对某个相关技术（如 Neo4j、Milvus、GraphRAG 等）的系统整理，也可以是与你的研究/工作紧密相关的专题总结。
+
+在主仓库的「知识拓展」部分，我们会以清单形式挂出这里的优秀专题，例如：
+
+- `Neo4J 简单应用`（本目录下的第一个示例专题）
+- `PowerRAG SDK 文本检索 Demo`（上传 Markdown → 解析/切分/向量化 → Top‑K 检索）
+
+我们希望通过 Extra Chapter：
+
+- **补充主教程**：覆盖 RAG 周边的生态组件和工程化话题  
+- **沉淀实践经验**：记录真实项目中的设计思路与解决方案  
+- **鼓励多元视角**：允许不同风格的探索与实验性内容  
+- **形成可引用的知识单元**：每个专题都可以被单独阅读和复用  
+
+---
+
+## 📂 推荐目录结构
+
+每一个新专题建议在 `Extra-chapter/` 下新建一个独立子目录，并至少包含一个 `readme.md` 作为主文档。参考结构如下（可按需要增减）：
+
+```text
+Extra-chapter/
+├── your-topic-name/                # 你的专题目录（必需）
+│   ├── readme.md                   # 专题主文档（必需）
+│   ├── images/                     # 图片资源（建议有图片时使用）
+│   │   ├── figure1.png
+│   │   └── figure2.jpg
+│   ├── code/                       # 代码示例（如有代码建议单独放）
+│   │   ├── demo.py
+│   │   └── requirements.txt
+│   └── data/                       # 示例数据（如有）
+│       └── sample_data.json
+└── README.md                       # 本说明文件
+```
+
+> 当前仓库中的 `Neo4J 简单应用` 专题位于 `Extra-chapter/Neo4J/readme.md`，你可以参考其组织方式，但不必完全照搬。
+
+---
+
+## 🧱 命名与排版建议
+
+### 1. 目录与文件命名
+
+- **专题目录名 (`your-topic-name/`)**
+  - 使用有语义的英文或中英文混合
+  - 避免过长，尽量能一眼看出主题，例如：
+    - `Neo4J-Simple-Application`
+    - `graph-rag-practice`
+
+- **主文档**
+  - 固定为：`readme.md`（小写），便于 GitHub 直接展示
+
+- **图片 / 代码 / 数据**
+  - 统一放在 `images/`、`code/`、`data/` 等子目录中
+  - 文件名建议包含用途或内容关键信息，例如：`pipeline-overview.png`、`retriever_benchmark.ipynb`
+
+### 2. 内容组织与标题层级
+
+推荐的章节骨架（可根据需要微调）：
+
+- `# 专题标题`
+- `## 背景与动机`
+- `## 场景或问题描述`
+- `## 技术方案 / 实现思路`
+- `## 实践步骤 / 代码示例`
+- `## 经验总结与坑点`
+
+标题层级请保持清晰，避免过多嵌套（一般不超过三级：`##` / `###` / `####`）。
+
+---
+
+## ✅ 内容质量与范围
+
+为了让内容对读者真正有帮助，提交 PR 前请尽量满足以下要求：
+
+- **与 RAG / LLM 应用相关**
+  - 例如：图数据库在 RAG 中的用法、检索评估方案、多模态扩展、部署/监控经验等
+
+- **技术与事实尽量准确**
+  - 如有推测或实验性的结论，请明确标注
+  - 关键代码建议给出可运行环境说明（Python 版本、依赖库等）
+
+- **结构清晰、可独立阅读**
+  - 读者不看主教程，也能从该专题收获完整的一段知识
+
+- **尊重版权与引用规范**
+  - 如引用论文、博客或开源项目，请在文末列出参考资料或链接
+
+---
+
+## 🔀 提交 PR 的具体步骤（建议）
+
+1. **Fork 仓库并创建分支**
+   - 从 `main` 或当前默认分支拉出个人分支，例如：`feat/extra-chapter-neo4j-practice`
+
+2. **在 `Extra-chapter/` 下创建你的专题目录**
+   - 遵循上文的目录组织方案
+   - 如有需要，可在根目录 `README.md` 的「第六部分：知识拓展」列表中，添加指向你专题的链接（格式参考已有的 `Neo4J 简单应用`）
+
+3. **本地检查排版与图片路径**
+   - 确认 Markdown 渲染正常、图片路径相对位置正确
+
+4. **编写有信息量的 Commit Message**
+
+推荐包含以下信息（可按实际场景调整）：
+
+```text
+Extra-chapter: <你的专题标题>
+
+- 新增专题目录：Extra-chapter/<your-topic-name>
+- 主要内容：简要说明该专题解决了什么问题 / 分享了哪些经验
+- 代码与数据：如有，说明放在哪些子目录，如何运行
+- 个人信息（可选但推荐）：你的 GitHub 链接、研究 / 工作方向等
+```
+
+5. **发起 PR**
+   - PR 标题中建议包含：`[Extra-chapter]` + 你的专题名
+   - 在 PR 描述里可以：
+     - 说明该专题与 All-in-RAG 主线章节的关系
+     - 标注阅读顺序建议（例如：适合在读完第 3 章之后再看）
+
+---
+
+## 💬 沟通与反馈
+
+如果你在撰写 Extra Chapter 的过程中：
+
+- 不确定选题是否合适
+- 想讨论目录设计或技术路线
+- 希望对草稿先做一次技术性 Review
+
+可以在仓库的 `Discussions` 或 `Issues` 中发起话题，标题中带上 `Extra-chapter` 关键字，方便维护者和其他贡献者一起参与讨论。
+
+---
+
+欢迎你把自己的实践经验沉淀在这里，让更多学习 RAG / LLM 应用的同学受益。 🎉
+

README.md

@@ -1,3 +1,235 @@
-# all-in-rag
+# All-in-RAG | 大模型应用开发实战一：RAG技术全栈指南
 
-Fork https://github.com/datawhalechina/all-in-rag#
+<div align='center'>
+  <img src="./docs/logo.svg" alt="All-in-RAG Logo" width="70%">
+</div>
+
+<div align="center">
+  <h2>🔍 检索增强生成 (RAG) 技术全栈指南</h2>
+  <p><em>从理论到实践，从基础到进阶，构建你的RAG技术体系</em></p>
+</div>
+
+<div align="center">
+  <img src="https://img.shields.io/github/stars/datawhalechina/all-in-rag?style=for-the-badge&logo=github&color=ff6b6b" alt="GitHub stars"/>
+  <img src="https://img.shields.io/github/forks/datawhalechina/all-in-rag?style=for-the-badge&logo=github&color=4ecdc4" alt="GitHub forks"/>
+  <img src="https://img.shields.io/badge/Python-3.12.7-blue?style=for-the-badge&logo=python&logoColor=white" alt="Python"/>
+  <a href="https://zread.ai/datawhalechina/all-in-rag">
+    <img src="https://img.shields.io/badge/Ask_Zread-_.svg?style=for-the-badge&color=00b0aa&labelColor=000000&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZD0iTTQuOTYxNTYgMS42MDAxSDIuMjQxNTZDMS44ODgxIDEuNjAwMSAxLjYwMTU2IDEuODg2NjQgMS42MDE1NiAyLjI0MDFWNC45NjAxQzEuNjAxNTYgNS4zMTM1NiAxLjg4ODEgNS42MDAxIDIuMjQxNTYgNS42MDAxSDQuOTYxNTZDNS4zMTUwMiA1LjYwMDEgNS42MDE1NiA1LjMxMzU2IDUuNjAxNTYgNC45NjAxVjIuMjQwMUM1LjYwMTU2IDEuODg2NjQgNS4zMTUwMiAxLjYwMDEgNC45NjE1NiAxLjYwMDFaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00Ljk2MTU2IDEwLjM5OTlIMi4yNDE1NkMxLjg4ODEgMTAuMzk5OSAxLjYwMTU2IDEwLjY4NjQgMS42MDE1NiAxMS4wMzk5VjEzLjc1OTlDMS42MDE1NiAxNC4xMTM0IDEuODg4MSAxNC4zOTk5IDIuMjQxNTYgMTQuMzk5OUg0Ljk2MTU2QzUuMzE1MDIgMTQuMzk5OSA1LjYwMTU2IDE0LjExMzQgNS42MDE1NiAxMy43NTk5VjExLjAzOTlDNS42MDE1NiAxMC42ODY0IDUuMzE1MDIgMTAuMzk5OSA0Ljk2MTU2IDEwLjM5OTlaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik0xMy43NTg0IDEuNjAwMUgxMS4wMzg0QzEwLjY4NSAxLjYwMDEgMTAuMzk4NCAxLjg4NjY0IDEwLjM5ODQgMi4yNDAxVjQuOTYwMUMxMC4zOTg0IDUuMzEzNTYgMTAuNjg1IDUuNjAwMSAxMS4wMzg0IDUuNjAwMUgxMy43NTg0QzE0LjExMTkgNS42MDAxIDE0LjM5ODQgNS4zMTM1NiAxNC4zOTg0IDQuOTYwMVYyLjI0MDFDMTQuMzk4NCAxLjg4NjY0IDE0LjExMTkgMS42MDAxIDEzLjc1ODQgMS42MDAxWiIgZmlsbD0iI2ZmZiIvPgo8cGF0aCBkPSJNNCAxMkwxMiA0TDQgMTJaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00IDEyTDEyIDQiIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSIxLjUiIHN0cm9rZS1saW5lY2FwPSJyb3VuZCIvPgo8L3N2Zz4K&logoColor=ffffff" alt="zread"/>
+  </a>
+</div>
+
+<div align="center">
+  <a href="https://datawhalechina.github.io/all-in-rag/">
+    <img src="https://img.shields.io/badge/📖_在线阅读-立即开始-success?style=for-the-badge&logoColor=white" alt="在线阅读"/>
+  </a>
+  <a href="README_en.md">
+    <img src="https://img.shields.io/badge/🌍_English-Version-blue?style=for-the-badge&logoColor=white" alt="English Version"/>
+  </a>
+  <a href="https://github.com/datawhalechina">
+    <img src="https://img.shields.io/badge/💬_讨论交流-加入我们-purple?style=for-the-badge&logoColor=white" alt="讨论交流"/>
+  </a>
+</div>
+
+<div align="center">
+  <br>
+  <table>
+    <tr>
+      <td align="center">🎯 <strong>系统化学习</strong><br>完整的RAG技术体系</td>
+      <td align="center">🛠️ <strong>动手实践</strong><br>丰富的项目案例</td>
+      <td align="center">🚀 <strong>生产就绪</strong><br>工程化最佳实践</td>
+      <td align="center">📊 <strong>多模态支持</strong><br>文本+图像检索</td>
+    </tr>
+  </table>
+</div>
+
+## 项目简介（中文 | [English](README_en.md)）
+
+本项目是一个面向大模型应用开发者的RAG（检索增强生成）技术全栈教程，旨在通过体系化的学习路径和动手实践项目，帮助开发者掌握基于大语言模型的RAG应用开发技能，构建生产级的智能问答和知识检索系统。
+
+**主要内容包括：**
+
+1. **RAG技术基础**：深入浅出地介绍RAG的核心概念、技术原理和应用场景
+2. **数据处理全流程**：从数据加载、清洗到文本分块的完整数据准备流程
+3. **索引构建与优化**：向量嵌入、多模态嵌入、向量数据库构建及索引优化技术
+4. **检索技术进阶**：混合检索、查询构建、Text2SQL等高级检索技术
+5. **生成集成与评估**：格式化生成、系统评估与优化方法
+6. **项目实战**：从基础到进阶的完整RAG应用开发实践
+
+## 项目意义
+
+随着大语言模型的快速发展，RAG技术已成为构建智能问答系统、知识检索应用的核心技术。然而，现有的RAG教程往往零散且缺乏系统性，初学者难以形成完整的技术体系认知。
+
+本项目从实践出发，结合最新的RAG技术发展趋势，构建了一套完整的RAG学习体系，帮助开发者：
+- 系统掌握RAG技术的理论基础和实践技能
+- 理解RAG系统的完整架构和各组件的作用
+- 具备独立开发RAG应用的能力
+- 掌握RAG系统的评估和优化方法
+
+## 项目受众
+
+**本项目适合以下人群学习：**
+- 具备Python编程基础，对RAG技术感兴趣的开发者
+- 希望系统学习RAG技术的AI工程师
+- 想要构建智能问答系统的产品开发者
+- 对检索增强生成技术有学习需求的研究人员
+
+**前置要求：**
+- 掌握Python基础语法和常用库的使用
+- 能够简单使用docker
+- 了解基本的LLM概念（推荐但非必需）
+- 具备基础的Linux命令行操作能力
+
+## 项目亮点
+
+1. **体系化学习路径**：从基础概念到高级应用，构建完整的RAG技术学习体系
+2. **理论与实践并重**：每个章节都包含理论讲解和代码实践，确保学以致用
+3. **多模态支持**：不仅涵盖文本RAG，还包括多模态嵌入和检索技术
+4. **工程化导向**：注重实际应用中的工程化问题，包括性能优化、系统评估等
+5. **丰富的实战项目**：提供从基础到进阶的多个实战项目，帮助巩固学习成果
+
+## 内容大纲
+
+### 第一部分：RAG基础入门
+
+**第一章 解锁RAG** [📖 查看章节](./docs/chapter1)
+- [x] [RAG简介](./docs/chapter1/01_RAG_intro.md) - RAG技术概述与应用场景
+- [x] [准备工作](./docs/chapter1/02_preparation.md) - 环境配置与准备
+- [x] [四步构建RAG](./docs/chapter1/03_get_start_rag.md) - 快速上手RAG开发
+- [x] [附：环境部署](./docs/chapter1/virtualenv.md) - Python虚拟环境部署方案补充 (贡献者: [@anarchysaiko](https://github.com/anarchysaiko))
+
+**第二章 数据准备** [📖 查看章节](./docs/chapter2)
+- [x] [数据加载](./docs/chapter2/04_data_load.md) - 多格式文档处理与加载
+- [x] [文本分块](./docs/chapter2/05_text_chunking.md) - 文本切分策略与优化
+
+### 第二部分：索引构建与优化
+
+**第三章 索引构建** [📖 查看章节](./docs/chapter3)
+- [x] [向量嵌入](./docs/chapter3/06_vector_embedding.md) - 文本向量化技术详解
+- [x] [多模态嵌入](./docs/chapter3/07_multimodal_embedding.md) - 图文多模态向量化
+- [x] [向量数据库](./docs/chapter3/08_vector_db.md) - 向量存储与检索系统
+- [x] [Milvus实践](./docs/chapter3/09_milvus.md) - Milvus多模态检索实战
+- [x] [索引优化](./docs/chapter3/10_index_optimization.md) - 索引性能调优技巧
+
+### 第三部分：检索技术进阶
+
+**第四章 检索优化** [📖 查看章节](./docs/chapter4)
+- [x] [混合检索](./docs/chapter4/11_hybrid_search.md) - 稠密+稀疏检索融合
+- [x] [查询构建](./docs/chapter4/12_query_construction.md) - 智能查询理解与构建
+- [x] [Text2SQL](./docs/chapter4/13_text2sql.md) - 自然语言转SQL查询
+- [x] [查询重构与分发](./docs/chapter4/14_query_rewriting.md) - 查询优化策略
+- [x] [检索进阶技术](./docs/chapter4/15_advanced_retrieval_techniques.md) - 高级检索算法
+
+### 第四部分：生成与评估
+
+**第五章 生成集成** [📖 查看章节](./docs/chapter5)
+- [x] [格式化生成](./docs/chapter5/16_formatted_generation.md) - 结构化输出与格式控制
+
+**第六章 RAG系统评估** [📖 查看章节](./docs/chapter6)
+- [x] [评估介绍](./docs/chapter6/18_system_evaluation.md) - RAG系统评估方法论
+- [x] [评估工具](./docs/chapter6/19_common_tools.md) - 常用评估工具与指标
+
+### 第五部分：高级应用与实战
+
+**第七章 高级RAG架构（拓展部分）** [📖 查看章节](./docs/chapter7)
+
+- [x] [基于知识图谱的RAG](./docs/chapter7/20_kg_rag.md)
+
+**第八章 项目实战一** [📖 查看章节](./docs/chapter8)
+- [x] [环境配置与项目架构](./docs/chapter8/01_env_architecture.md)
+- [x] [数据准备模块实现](./docs/chapter8/02_data_preparation.md)
+- [x] [索引构建与检索优化](./docs/chapter8/03_index_retrieval.md)
+- [x] [生成集成与系统整合](./docs/chapter8/04_generation_sys.md)
+
+**第九章 项目实战一优化（选修篇）** [📖 查看章节](./docs/chapter9)
+
+[🍽️ 项目展示](https://github.com/FutureUnreal/What-to-eat-today)
+- [x] [图RAG架构设计](./docs/chapter9/01_graph_rag_architecture.md)
+- [x] [图数据建模与准备](./docs/chapter9/02_graph_data_modeling.md)
+- [x] [Milvus索引构建](./docs/chapter9/03_index_construction.md)
+- [x] [智能查询路由与检索策略](./docs/chapter9/04_intelligent_query_routing.md)
+
+**第十章 项目实战二（选修篇）** [📖 查看章节](./docs/chapter10) *规划中*
+
+### Extra-chapter
+
+- [Neo4J 简单应用](./Extra-chapter/Neo4J-Simple-Application/readme.md) （贡献者: [dalvqw](https://github.com/FutureUnreal)）
+
+> 如果你在使用 RAG / 向量数据库 / Agentic RAG 等相关技术时，也有值得分享的经验与专题内容，非常欢迎以独立章节的形式投稿到 [Extra Chapter](./Extra-chapter/) 中。提交前请先阅读 Extra Chapter 的[贡献与 PR 指南](./Extra-chapter/README.md)，我们会根据内容的完整度、实践深度与参考价值综合评估是否合并，并视情况在主教程中进行引用或扩展说明。
+
+## 目录结构说明
+
+```
+all-in-rag/
+├── docs/           # 教程文档
+├── code/           # 代码示例
+├── data/           # 示例数据
+├── models/         # 预训练模型
+├── Extra-chapter/  # 扩展章节与社区实践内容
+└── README.md       # 项目说明
+```
+
+## 实战项目展示
+
+### 第八章 项目一：
+
+![项目一](./project01.png)
+
+### 第九章 项目一（Graph RAG优化）：
+
+![项目一（Graph RAG优化）](./project01_graph.png)
+
+### 第十章 项目二：
+
+## 致谢
+
+**核心贡献者**
+- [dalvqw-项目负责人](https://github.com/FutureUnreal)（项目发起人与主要贡献者）
+
+**额外章节贡献者**
+- [孙超-内容创作者](https://github.com/anarchysaiko)（Datawhale成员-上海工程技术大学）
+
+### 特别感谢
+- 感谢 [@Sm1les](https://github.com/Sm1les) 对本项目的帮助与支持
+- 感谢所有为本项目做出贡献的开发者们
+- 感谢开源社区提供的优秀工具和框架支持
+- 特别感谢以下为教程做出贡献的开发者！
+
+[![Contributors](https://contrib.rocks/image?repo=datawhalechina/all-in-rag)](https://github.com/datawhalechina/all-in-rag/graphs/contributors)
+
+*Made with [contrib.rocks](https://contrib.rocks).*
+
+## 参与贡献
+
+我们欢迎所有形式的贡献，包括但不限于：
+
+- 🚨 **Bug报告**：发现问题请提交 [Issue](https://github.com/datawhalechina/all-in-rag/issues)
+- 💭 **教程建议**：有好的想法欢迎在 [Discussions](https://github.com/datawhalechina/all-in-rag/discussions) 中讨论
+- 📚 **文档改进**：帮助完善文档内容和示例代码（当前仅支持 Extra-chapter 优质内容pr）
+
+## Star History
+
+[![all-in-rag stats](https://datawhalechina.github.io/members-visualization/badges/all-in-rag.png)](https://datawhalechina.github.io/members-visualization/repo-badge?repo=all-in-rag)
+
+<div align="center">
+  <p>如果这个项目对你有帮助，请给我们一个 ⭐️</p>
+  <p>让更多人发现这个项目（护食？发来！）</p>
+</div>
+
+![star](./emoji.png)
+
+## 关于 Datawhale
+
+<div align='center'>
+    <img src="https://raw.githubusercontent.com/datawhalechina/pumpkin-book/master/res/qrcode.jpeg" alt="Datawhale" width="30%">
+    <p>扫描二维码关注 Datawhale 公众号，获取更多优质开源内容</p>
+</div>
+
+---
+
+## 许可证
+
+<a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/"><img alt="知识共享许可协议" style="border-width:0" src="https://img.shields.io/badge/license-CC%20BY--NC--SA%204.0-lightgrey" /></a>
+
+本作品采用 [知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议](http://creativecommons.org/licenses/by-nc-sa/4.0/) 进行许可。
+
+---

README_en.md

@@ -0,0 +1,222 @@
+# All-in-RAG | Large Model Application Development Practice: RAG Technology Full-Stack Guide
+
+<div align='center'>
+  <img src="./docs/logo.svg" alt="All-in-RAG Logo" width="70%">
+</div>
+
+<div align="center">
+  <h2>🔍 Retrieval-Augmented Generation (RAG) Technology Full-Stack Guide</h2>
+  <p><em>From theory to practice, from basics to advanced, build your RAG technology system</em></p>
+</div>
+
+<div align="center">
+  <img src="https://img.shields.io/github/stars/datawhalechina/all-in-rag?style=for-the-badge&logo=github&color=ff6b6b" alt="GitHub stars"/>
+  <img src="https://img.shields.io/github/forks/datawhalechina/all-in-rag?style=for-the-badge&logo=github&color=4ecdc4" alt="GitHub forks"/>
+  <img src="https://img.shields.io/badge/Python-3.12.7-blue?style=for-the-badge&logo=python&logoColor=white" alt="Python"/>
+</div>
+
+<div align="center">
+  <a href="https://datawhalechina.github.io/all-in-rag/#/en/">
+    <img src="https://img.shields.io/badge/📖_Online_Reading-Get_Started-success?style=for-the-badge&logoColor=white" alt="Online Reading"/>
+  </a>
+  <a href="README.md">
+    <img src="https://img.shields.io/badge/🌍_中文-Version-blue?style=for-the-badge&logoColor=white" alt="Chinese Version"/>
+  </a>
+  <a href="https://github.com/datawhalechina">
+    <img src="https://img.shields.io/badge/💬_Discussion-Join_Us-purple?style=for-the-badge&logoColor=white" alt="Discussion"/>
+  </a>
+</div>
+
+<div align="center">
+  <br>
+  <table>
+    <tr>
+      <td align="center">🎯 <strong>Systematic Learning</strong><br>Complete RAG technology system</td>
+      <td align="center">🛠️ <strong>Hands-on Practice</strong><br>Rich project examples</td>
+      <td align="center">🚀 <strong>Production Ready</strong><br>Engineering best practices</td>
+      <td align="center">📊 <strong>Multimodal Support</strong><br>Text + Image retrieval</td>
+    </tr>
+  </table>
+</div>
+
+## Project Introduction（[中文](README.md) | English）
+
+This project is a comprehensive RAG (Retrieval-Augmented Generation) technology full-stack tutorial for large model application developers. It aims to help developers master RAG application development skills based on large language models through systematic learning paths and hands-on practice projects, building production-grade intelligent Q&A and knowledge retrieval systems.
+
+**Main content includes:**
+
+1. **RAG Technology Fundamentals**: In-depth introduction to RAG core concepts, technical principles, and application scenarios
+2. **Complete Data Processing Pipeline**: From data loading, cleaning to text chunking - the complete data preparation process
+3. **Index Construction and Optimization**: Vector embedding, multimodal embedding, vector database construction and index optimization techniques
+4. **Advanced Retrieval Techniques**: Hybrid retrieval, query construction, Text2SQL and other advanced retrieval technologies
+5. **Generation Integration and Evaluation**: Formatted generation, system evaluation and optimization methods
+6. **Project Practice**: Complete RAG application development practice from basic to advanced
+
+## Project Significance
+
+With the rapid development of large language models, RAG technology has become the core technology for building intelligent Q&A systems and knowledge retrieval applications. However, existing RAG tutorials are often scattered and lack systematicity, making it difficult for beginners to form a complete understanding of the technical system.
+
+Starting from practice and combining the latest RAG technology development trends, this project builds a complete RAG learning system to help developers:
+- Systematically master the theoretical foundation and practical skills of RAG technology
+- Understand the complete architecture of RAG systems and the role of each component
+- Develop the ability to independently develop RAG applications
+- Master evaluation and optimization methods for RAG systems
+
+## Target Audience
+
+**This project is suitable for the following groups:**
+- Developers with Python programming foundation who are interested in RAG technology
+- AI engineers who want to systematically learn RAG technology
+- Product developers who want to build intelligent Q&A systems
+- Researchers with learning needs for retrieval-augmented generation technology
+
+**Prerequisites:**
+- Master Python basic syntax and usage of common libraries
+- Ability to use Docker simply
+- Understanding of basic LLM concepts (recommended but not required)
+- Basic Linux command line operation skills
+
+## Project Highlights
+
+1. **Systematic Learning Path**: From basic concepts to advanced applications, building a complete RAG technology learning system
+2. **Theory and Practice Combined**: Each chapter includes theoretical explanation and code practice to ensure learning and application
+3. **Multimodal Support**: Covers not only text RAG, but also multimodal embedding and retrieval technologies
+4. **Engineering-Oriented**: Focus on engineering problems in practical applications, including performance optimization, system evaluation, etc.
+5. **Rich Practical Projects**: Provides multiple practical projects from basic to advanced to help consolidate learning outcomes
+
+## Content Outline
+
+### Part I: RAG Fundamentals
+
+**Chapter 1 Unlocking RAG** [📖 View Chapter](./docs/en/chapter1)
+1. [x] [RAG Introduction](./docs/en/chapter1/01_RAG_intro.md) - RAG technology overview and application scenarios
+2. [x] [Preparation](./docs/en/chapter1/02_preparation.md) - Environment configuration and preparation
+3. [x] [Four Steps to Build RAG](./docs/en/chapter1/03_get_start_rag.md) - Quick start with RAG development
+
+**Chapter 2 Data Preparation** [📖 View Chapter](./docs/en/chapter2)
+1. [x] [Data Loading](./docs/en/chapter2/04_data_load.md) - Multi-format document processing and loading
+2. [x] [Text Chunking](./docs/en/chapter2/05_text_chunking.md) - Text segmentation strategies and optimization
+
+### Part II: Index Construction and Optimization
+
+**Chapter 3 Index Construction** [📖 View Chapter](./docs/en/chapter3)
+1. [x] [Vector Embedding](./docs/en/chapter3/06_vector_embedding.md) - Detailed explanation of text vectorization technology
+2. [x] [Multimodal Embedding](./docs/en/chapter3/07_multimodal_embedding.md) - Image-text multimodal vectorization
+3. [x] [Vector Database](./docs/en/chapter3/08_vector_db.md) - Vector storage and retrieval systems
+4. [x] [Milvus Practice](./docs/en/chapter3/09_milvus.md) - Milvus multimodal retrieval practice
+5. [x] [Index Optimization](./docs/en/chapter3/10_index_optimization.md) - Index performance tuning techniques
+
+### Part III: Advanced Retrieval Techniques
+
+**Chapter 4 Retrieval Optimization** [📖 View Chapter](./docs/en/chapter4)
+1. [x] [Hybrid Search](./docs/en/chapter4/11_hybrid_search.md) - Dense + sparse retrieval fusion
+2. [x] [Query Construction](./docs/en/chapter4/12_query_construction.md) - Intelligent query understanding and construction
+3. [x] [Text2SQL](./docs/en/chapter4/13_text2sql.md) - Natural language to SQL query
+4. [x] [Query Rewriting and Routing](./docs/en/chapter4/14_query_rewriting.md) - Query optimization strategies
+5. [x] [Advanced Retrieval Techniques](./docs/en/chapter4/15_advanced_retrieval_techniques.md) - Advanced retrieval algorithms
+
+### Part IV: Generation and Evaluation
+
+**Chapter 5 Generation Integration** [📖 View Chapter](./docs/en/chapter5)
+1. [x] [Formatted Generation](./docs/en/chapter5/16_formatted_generation.md) - Structured output and format control
+
+**Chapter 6 RAG System Evaluation** [📖 View Chapter](./docs/en/chapter6)
+1. [x] [Evaluation Introduction](./docs/en/chapter6/18_system_evaluation.md) - RAG system evaluation methodology
+2. [x] [Evaluation Tools](./docs/en/chapter6/19_common_tools.md) - Common evaluation tools and metrics
+
+### Part V: Advanced Applications and Practice
+
+**Chapter 7 Advanced RAG Architecture (Extended Elective)** [📖 View Chapter](./docs/en/chapter7)
+
+1. [x] [Knowledge Graph-based RAG](./docs/en/chapter7/20_kg_rag.md)
+
+**Chapter 8 Project Practice I (Basic)** [📖 View Chapter](./docs/en/chapter8)
+1. [x] [Environment Configuration and Project Architecture](./docs/en/chapter8/01_env_architecture.md)
+2. [x] [Data Preparation Module Implementation](./docs/en/chapter8/02_data_preparation.md)
+3. [x] [Index Construction and Retrieval Optimization](./docs/en/chapter8/03_index_retrieval.md)
+4. [x] [Generation Integration and System Integration](./docs/en/chapter8/04_generation_sys.md)
+
+**Chapter 9 Project Practice I Optimization (Elective)** [📖 View Chapter](./docs/en/chapter9)
+
+[🍽️ Project Demo](https://github.com/FutureUnreal/What-to-eat-today)
+1. [x] [Graph RAG Architecture Design](./docs/en/chapter9/01_graph_rag_architecture.md)
+2. [x] [Graph Data Modeling and Preparation](./docs/en/chapter9/02_graph_data_modeling.md)
+3. [x] [Milvus Index Construction](./docs/en/chapter9/03_index_construction.md)
+4. [x] [Intelligent Query Routing and Retrieval Strategy](./docs/en/chapter9/04_intelligent_query_routing.md)
+
+**Chapter 10 Project Practice II (Elective)** [📖 View Chapter](./docs/en/chapter10) *In Planning*
+
+## Directory Structure
+
+```
+all-in-rag/
+├── docs/           # Tutorial documentation
+├── code/           # Code examples
+├── data/           # Sample data
+├── models/         # Pre-trained models
+└── README.md       # Project description
+```
+
+## Practical Project Showcase
+
+### Chapter 8 Project I:
+
+![Project I](./project01.png)
+
+### Chapter 9 Project I (Graph RAG Optimization):
+
+![Project I (Graph RAG Optimization)](./project01_graph.png)
+
+### Chapter 10 Project II:
+
+## Acknowledgments
+
+**Core Contributors**
+- [Yin Dalv - Project Lead](https://github.com/FutureUnreal) (Project initiator and main contributor)
+
+### Special Thanks
+- Thanks to [@Sm1les](https://github.com/Sm1les) for help and support on this project
+- Thanks to all developers who contributed to this project
+- Thanks to the open source community for providing excellent tools and framework support
+- Special thanks to the following developers who contributed to the tutorial!
+
+[![Contributors](https://contrib.rocks/image?repo=datawhalechina/all-in-rag)](https://github.com/datawhalechina/all-in-rag/graphs/contributors)
+
+*Made with [contrib.rocks](https://contrib.rocks).*
+
+## Contributing
+
+We welcome all forms of contributions, including but not limited to:
+
+- 🚨 **Bug Reports**: Please submit [Issues](https://github.com/datawhalechina/all-in-rag/issues) if you find problems
+- 💭 **Feature Suggestions**: Welcome to discuss good ideas in [Discussions](https://github.com/datawhalechina/all-in-rag/discussions)
+- 📚 **Documentation Improvement**: Help improve documentation content and example code
+- ⚡ **Code Contributions**: Submit [Pull Requests](https://github.com/datawhalechina/all-in-rag/pulls) to improve the project
+
+## Star History
+
+[![all-in-rag stats](https://datawhalechina.github.io/members-visualization/badges/all-in-rag.png)](https://datawhalechina.github.io/members-visualization/repo-badge?repo=all-in-rag)
+
+<div align="center">
+  <p>If this project helps you, please give us a ⭐️</p>
+  <p>Let more people discover this project (Food protection? Bring it on!)</p>
+</div>
+
+![star](./emoji.png)
+
+## About Datawhale
+
+<div align='center'>
+    <img src="https://raw.githubusercontent.com/datawhalechina/pumpkin-book/master/res/qrcode.jpeg" alt="Datawhale" width="30%">
+    <p>Scan the QR code to follow Datawhale WeChat Official Account for more quality open source content</p>
+</div>
+
+---
+
+## License
+
+<a rel="license" href="http://creativecommons.org/licenses/by-nc-sa/4.0/"><img alt="Creative Commons License" style="border-width:0" src="https://img.shields.io/badge/license-CC%20BY--NC--SA%204.0-lightgrey" /></a>
+
+This work is licensed under a [Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License](http://creativecommons.org/licenses/by-nc-sa/4.0/).
+
+---

code/C1/01_langchain_example.py

@@ -0,0 +1,75 @@
+import os
+# hugging face镜像设置，如果国内环境无法使用启用该设置
+# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
+from dotenv import load_dotenv
+from langchain_community.document_loaders import UnstructuredMarkdownLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.vectorstores import InMemoryVectorStore
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+
+load_dotenv()
+
+markdown_path = "../../data/C1/markdown/easy-rl-chapter1.md"
+
+# 加载本地markdown文件
+loader = UnstructuredMarkdownLoader(markdown_path)
+docs = loader.load()
+
+# 文本分块
+text_splitter = RecursiveCharacterTextSplitter()
+chunks = text_splitter.split_documents(docs)
+
+# 中文嵌入模型
+embeddings = HuggingFaceEmbeddings(
+    model_name="BAAI/bge-small-zh-v1.5",
+    model_kwargs={'device': 'cpu'},
+    encode_kwargs={'normalize_embeddings': True}
+)
+  
+# 构建向量存储
+vectorstore = InMemoryVectorStore(embeddings)
+vectorstore.add_documents(chunks)
+
+# 提示词模板
+prompt = ChatPromptTemplate.from_template("""请根据下面提供的上下文信息来回答问题。
+请确保你的回答完全基于这些上下文。
+如果上下文中没有足够的信息来回答问题，请直接告知：“抱歉，我无法根据提供的上下文找到相关信息来回答此问题。”
+
+上下文:
+{context}
+
+问题: {question}
+
+回答:"""
+                                          )
+
+# 配置大语言模型
+
+# 使用 AIHubmix
+llm = ChatOpenAI(
+    model="glm-4.7-flash-free",
+    temperature=0.7,
+    max_tokens=4096,
+    api_key=os.getenv("DEEPSEEK_API_KEY"),
+    base_url="https://aihubmix.com/v1"
+)
+
+# llm = ChatOpenAI(
+#     model="deepseek-chat",
+#     temperature=0.7,
+#     max_tokens=4096,
+#     api_key=os.getenv("DEEPSEEK_API_KEY"),
+#     base_url="https://api.deepseek.com"
+# )
+
+# 用户查询
+question = "文中举了哪些例子？"
+
+# 在向量存储中查询相关文档
+retrieved_docs = vectorstore.similarity_search(question, k=3)
+docs_content = "\n\n".join(doc.page_content for doc in retrieved_docs)
+
+answer = llm.invoke(prompt.format(question=question, context=docs_content))
+print(answer)

code/C1/02_llamaIndex_example.py

@@ -0,0 +1,33 @@
+import os
+# os.environ['HF_ENDPOINT']='https://hf-mirror.com'
+from dotenv import load_dotenv
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings 
+from llama_index.llms.openai_like import OpenAILike
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+load_dotenv()
+
+# 使用 AIHubmix
+Settings.llm = OpenAILike(
+    model="glm-4.7-flash-free",
+    api_key=os.getenv("DEEPSEEK_API_KEY"),
+    api_base="https://aihubmix.com/v1",
+    is_chat_model=True
+)
+
+# Settings.llm = OpenAI(
+#     model="deepseek-chat",
+#     api_key=os.getenv("DEEPSEEK_API_KEY"),
+#     api_base="https://api.deepseek.com"
+# )
+Settings.embed_model = HuggingFaceEmbedding("BAAI/bge-small-zh-v1.5")
+
+docs = SimpleDirectoryReader(input_files=["../../data/C1/markdown/easy-rl-chapter1.md"]).load_data()
+
+index = VectorStoreIndex.from_documents(docs)
+
+query_engine = index.as_query_engine()
+
+print(query_engine.get_prompts())
+
+print(query_engine.query("文中举了哪些例子?"))

code/C1/fix_nltk.py

@@ -0,0 +1,4 @@
+import nltk
+
+nltk.download('punkt', force=True)
+nltk.download('averaged_perceptron_tagger', force=True)

code/C2/01_unstructured_example.py

@@ -0,0 +1,25 @@
+from unstructured.partition.auto import partition
+
+# PDF文件路径
+pdf_path = "../../data/C2/pdf/rag.pdf"
+
+# 使用Unstructured加载并解析PDF文档
+elements = partition(
+    filename=pdf_path,
+    content_type="application/pdf"
+)
+
+# 打印解析结果
+print(f"解析完成: {len(elements)} 个元素, {sum(len(str(e)) for e in elements)} 字符")
+
+# 统计元素类型
+from collections import Counter
+types = Counter(e.category for e in elements)
+print(f"元素类型: {dict(types)}")
+
+# 显示所有元素
+print("\n所有元素:")
+for i, element in enumerate(elements, 1):
+    print(f"Element {i} ({element.category}):")
+    print(element)
+    print("=" * 60)

code/C2/02_character_splitter.py

@@ -0,0 +1,23 @@
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+
+# 1. 文档加载
+loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
+docs = loader.load()
+
+# 2. 初始化固定大小分块器
+text_splitter = CharacterTextSplitter(
+    chunk_size=200,    # 每个块的大小
+    chunk_overlap=10   # 块之间的重叠大小
+)
+
+# 3. 执行分块
+chunks = text_splitter.split_documents(docs)
+
+# 4. 打印结果
+print(f"文本被切分为 {len(chunks)} 个块。\n")
+print("--- 前5个块内容示例 ---")
+for i, chunk in enumerate(chunks[:5]):
+    print("=" * 60)
+    # chunk 是一个 Document 对象，需要访问它的 .page_content 属性来获取文本
+    print(f'块 {i+1} (长度: {len(chunk.page_content)}): "{chunk.page_content}"')

code/C2/03_recursive_character_splitter.py

@@ -0,0 +1,20 @@
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+
+loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
+docs = loader.load()
+
+text_splitter = RecursiveCharacterTextSplitter(
+    # 针对中英文混合文本，定义一个更全面的分隔符列表
+    separators=["\n\n", "\n", "。", "，", " ", ""], # 按顺序尝试分割
+    chunk_size=200,
+    chunk_overlap=10
+)
+
+chunks = text_splitter.split_documents(docs)
+
+print(f"文本被切分为 {len(chunks)} 个块。\n")
+print("--- 前5个块内容示例 ---")
+for i, chunk in enumerate(chunks[:5]):
+    print("=" * 60)
+    print(f'块 {i+1} (长度: {len(chunk.page_content)}): "{chunk.page_content}"')

code/C2/04_semantic_chunker.py

@@ -0,0 +1,26 @@
+from langchain_experimental.text_splitter import SemanticChunker
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.document_loaders import TextLoader
+
+embeddings = HuggingFaceEmbeddings(
+    model_name="BAAI/bge-small-zh-v1.5",
+    model_kwargs={'device': 'cpu'},
+    encode_kwargs={'normalize_embeddings': True}
+)
+
+# 初始化 SemanticChunker
+text_splitter = SemanticChunker(
+    embeddings,
+    breakpoint_threshold_type="percentile" # 也可以是 "standard_deviation", "interquartile", "gradient"
+)
+
+loader = TextLoader("../../data/C2/txt/蜂医.txt", encoding="utf-8")
+documents = loader.load()
+
+docs = text_splitter.split_documents(documents)
+
+print(f"文本被切分为 {len(docs)} 个块。\n")
+print("--- 前2个块内容示例 ---")
+for i, chunk in enumerate(docs[:2]):
+    print("=" * 60)
+    print(f'块 {i+1} (长度: {len(chunk.page_content)}):\n"{chunk.page_content}"')

code/C3/01_bge_visualized.py

@@ -0,0 +1,32 @@
+import torch
+from visual_bge.visual_bge.modeling import Visualized_BGE
+
+model = Visualized_BGE(model_name_bge="BAAI/bge-base-en-v1.5",
+                      model_weight="../../models/bge/Visualized_base_en_v1.5.pth")
+model.eval()
+
+with torch.no_grad():
+    text_emb = model.encode(text="datawhale开源组织的logo")
+    img_emb_1 = model.encode(image="../../data/C3/imgs/datawhale01.png")
+    multi_emb_1 = model.encode(image="../../data/C3/imgs/datawhale01.png", text="datawhale开源组织的logo")
+    img_emb_2 = model.encode(image="../../data/C3/imgs/datawhale02.png")
+    multi_emb_2 = model.encode(image="../../data/C3/imgs/datawhale02.png", text="datawhale开源组织的logo")
+
+# 计算相似度
+sim_1 = img_emb_1 @ img_emb_2.T
+sim_2 = img_emb_1 @ multi_emb_1.T
+sim_3 = text_emb @ multi_emb_1.T
+sim_4 = multi_emb_1 @ multi_emb_2.T
+
+print("=== 相似度计算结果 ===")
+print(f"纯图像 vs 纯图像: {sim_1}")
+print(f"图文结合1 vs 纯图像: {sim_2}")
+print(f"图文结合1 vs 纯文本: {sim_3}")
+print(f"图文结合1 vs 图文结合2: {sim_4}")
+
+# 向量信息分析
+print("\n=== 嵌入向量信息 ===")
+print(f"多模态向量维度: {multi_emb_1.shape}")
+print(f"图像向量维度: {img_emb_1.shape}")
+print(f"多模态向量示例 (前10个元素): {multi_emb_1[0][:10]}")
+print(f"图像向量示例 (前10个元素):   {img_emb_1[0][:10]}")

code/C3/02_langchain_faiss.py

@@ -0,0 +1,37 @@
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_core.documents import Document
+
+# 1. 示例文本和嵌入模型
+texts = [
+    "张三是法外狂徒",
+    "FAISS是一个用于高效相似性搜索和密集向量聚类的库。",
+    "LangChain是一个用于开发由语言模型驱动的应用程序的框架。"
+]
+docs = [Document(page_content=t) for t in texts]
+embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+
+# 2. 创建向量存储并保存到本地
+vectorstore = FAISS.from_documents(docs, embeddings)
+
+local_faiss_path = "./faiss_index_store"
+vectorstore.save_local(local_faiss_path)
+
+print(f"FAISS index has been saved to {local_faiss_path}")
+
+# 3. 加载索引并执行查询
+# 加载时需指定相同的嵌入模型，并允许反序列化
+loaded_vectorstore = FAISS.load_local(
+    local_faiss_path,
+    embeddings,
+    allow_dangerous_deserialization=True
+)
+
+# 执行相似性搜索
+query = "FAISS是做什么的？"
+results = loaded_vectorstore.similarity_search(query, k=1)
+
+print(f"\n查询: '{query}'")
+print("相似度最高的文档:")
+for doc in results:
+    print(f"- {doc.page_content}")

code/C3/03_llamaindex_vector.py

@@ -0,0 +1,19 @@
+from llama_index.core import VectorStoreIndex, Document, Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+# 1. 配置全局嵌入模型
+Settings.embed_model = HuggingFaceEmbedding("BAAI/bge-small-zh-v1.5")
+
+# 2. 创建示例文档
+texts = [
+    "张三是法外狂徒",
+    "LlamaIndex是一个用于构建和查询私有或领域特定数据的框架。",
+    "它提供了数据连接、索引和查询接口等工具。"
+]
+docs = [Document(text=t) for t in texts]
+
+# 3. 创建索引并持久化到本地
+index = VectorStoreIndex.from_documents(docs)
+persist_path = "./llamaindex_index_store"
+index.storage_context.persist(persist_dir=persist_path)
+print(f"LlamaIndex 索引已保存至: {persist_path}")

code/C3/04_multi_milvus.py

@@ -0,0 +1,160 @@
+import os
+from tqdm import tqdm
+from glob import glob
+import torch
+from visual_bge.visual_bge.modeling import Visualized_BGE
+from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
+import numpy as np
+import cv2
+from PIL import Image
+
+# 1. 初始化设置
+MODEL_NAME = "BAAI/bge-base-en-v1.5"
+MODEL_PATH = "../../models/bge/Visualized_base_en_v1.5.pth"
+DATA_DIR = "../../data/C3"
+COLLECTION_NAME = "multimodal_demo"
+MILVUS_URI = "http://localhost:19530"
+
+# 2. 定义工具 (编码器和可视化函数)
+class Encoder:
+    """编码器类，用于将图像和文本编码为向量。"""
+    def __init__(self, model_name: str, model_path: str):
+        self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)
+        self.model.eval()
+
+    def encode_query(self, image_path: str, text: str) -> list[float]:
+        with torch.no_grad():
+            query_emb = self.model.encode(image=image_path, text=text)
+        return query_emb.tolist()[0]
+
+    def encode_image(self, image_path: str) -> list[float]:
+        with torch.no_grad():
+            query_emb = self.model.encode(image=image_path)
+        return query_emb.tolist()[0]
+
+def visualize_results(query_image_path: str, retrieved_images: list, img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
+    """从检索到的图像列表创建一个全景图用于可视化。"""
+    panoramic_width = img_width * row_count
+    panoramic_height = img_height * row_count
+    panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
+    query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
+
+    # 处理查询图像
+    query_pil = Image.open(query_image_path).convert("RGB")
+    query_cv = np.array(query_pil)[:, :, ::-1]
+    resized_query = cv2.resize(query_cv, (img_width, img_height))
+    bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
+    query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
+    cv2.putText(query_display_area, "Query", (10, panoramic_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+
+    # 处理检索到的图像
+    for i, img_path in enumerate(retrieved_images):
+        row, col = i // row_count, i % row_count
+        start_row, start_col = row * img_height, col * img_width
+        
+        retrieved_pil = Image.open(img_path).convert("RGB")
+        retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
+        resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
+        bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
+        panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
+        
+        # 添加索引号
+        cv2.putText(panoramic_image, str(i), (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+    return np.hstack([query_display_area, panoramic_image])
+
+# 3. 初始化客户端
+print("--> 正在初始化编码器和Milvus客户端...")
+encoder = Encoder(MODEL_NAME, MODEL_PATH)
+milvus_client = MilvusClient(uri=MILVUS_URI)
+
+# 4. 创建 Milvus Collection
+print(f"\n--> 正在创建 Collection '{COLLECTION_NAME}'")
+if milvus_client.has_collection(COLLECTION_NAME):
+    milvus_client.drop_collection(COLLECTION_NAME)
+    print(f"已删除已存在的 Collection: '{COLLECTION_NAME}'")
+
+image_list = glob(os.path.join(DATA_DIR, "dragon", "*.png"))
+if not image_list:
+    raise FileNotFoundError(f"在 {DATA_DIR}/dragon/ 中未找到任何 .png 图像。")
+dim = len(encoder.encode_image(image_list[0]))
+
+fields = [
+    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
+    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
+    FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
+]
+
+# 创建集合 Schema
+schema = CollectionSchema(fields, description="多模态图文检索")
+print("Schema 结构:")
+print(schema)
+
+# 创建集合
+milvus_client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
+print(f"成功创建 Collection: '{COLLECTION_NAME}'")
+print("Collection 结构:")
+print(milvus_client.describe_collection(collection_name=COLLECTION_NAME))
+
+# 5. 准备并插入数据
+print(f"\n--> 正在向 '{COLLECTION_NAME}' 插入数据")
+data_to_insert = []
+for image_path in tqdm(image_list, desc="生成图像嵌入"):
+    vector = encoder.encode_image(image_path)
+    data_to_insert.append({"vector": vector, "image_path": image_path})
+
+if data_to_insert:
+    result = milvus_client.insert(collection_name=COLLECTION_NAME, data=data_to_insert)
+    print(f"成功插入 {result['insert_count']} 条数据。")
+
+# 6. 创建索引
+print(f"\n--> 正在为 '{COLLECTION_NAME}' 创建索引")
+index_params = milvus_client.prepare_index_params()
+index_params.add_index(
+    field_name="vector",
+    index_type="HNSW",
+    metric_type="COSINE",
+    params={"M": 16, "efConstruction": 256}
+)
+milvus_client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
+print("成功为向量字段创建 HNSW 索引。")
+print("索引详情:")
+print(milvus_client.describe_index(collection_name=COLLECTION_NAME, index_name="vector"))
+milvus_client.load_collection(collection_name=COLLECTION_NAME)
+print("已加载 Collection 到内存中。")
+
+# 7. 执行多模态检索
+print(f"\n--> 正在 '{COLLECTION_NAME}' 中执行检索")
+query_image_path = os.path.join(DATA_DIR, "dragon", "query.png")
+query_text = "一条龙"
+query_vector = encoder.encode_query(image_path=query_image_path, text=query_text)
+
+search_results = milvus_client.search(
+    collection_name=COLLECTION_NAME,
+    data=[query_vector],
+    output_fields=["image_path"],
+    limit=5,
+    search_params={"metric_type": "COSINE", "params": {"ef": 128}}
+)[0]
+
+retrieved_images = []
+print("检索结果:")
+for i, hit in enumerate(search_results):
+    print(f"  Top {i+1}: ID={hit['id']}, 距离={hit['distance']:.4f}, 路径='{hit['entity']['image_path']}'")
+    retrieved_images.append(hit['entity']['image_path'])
+
+# 8. 可视化与清理
+print(f"\n--> 正在可视化结果并清理资源")
+if not retrieved_images:
+    print("没有检索到任何图像。")
+else:
+    panoramic_image = visualize_results(query_image_path, retrieved_images)
+    combined_image_path = os.path.join(DATA_DIR, "search_result.png")
+    cv2.imwrite(combined_image_path, panoramic_image)
+    print(f"结果图像已保存到: {combined_image_path}")
+    Image.open(combined_image_path).show()
+
+milvus_client.release_collection(collection_name=COLLECTION_NAME)
+print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
+milvus_client.drop_collection(COLLECTION_NAME)
+print(f"已删除 Collection: '{COLLECTION_NAME}'")

code/C3/05_sentence_window_retrieval.py

@@ -0,0 +1,51 @@
+import os
+from llama_index.core.node_parser import SentenceWindowNodeParser, SentenceSplitter
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+from llama_index.llms.deepseek import DeepSeek
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.postprocessor import MetadataReplacementPostProcessor
+
+# 1. 配置模型
+Settings.llm = DeepSeek(model="deepseek-chat", temperature=0.1, api_key=os.getenv("DEEPSEEK_API_KEY"))
+Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en")
+
+# 2. 加载文档
+documents = SimpleDirectoryReader(
+    input_files=["../../data/C3/pdf/IPCC_AR6_WGII_Chapter03.pdf"]
+).load_data()
+
+# 3. 创建节点与构建索引
+# 3.1 句子窗口索引
+node_parser = SentenceWindowNodeParser.from_defaults(
+    window_size=3,
+    window_metadata_key="window",
+    original_text_metadata_key="original_text",
+)
+sentence_nodes = node_parser.get_nodes_from_documents(documents)
+sentence_index = VectorStoreIndex(sentence_nodes)
+
+# 3.2 常规分块索引 (基准)
+base_parser = SentenceSplitter(chunk_size=512)
+base_nodes = base_parser.get_nodes_from_documents(documents)
+base_index = VectorStoreIndex(base_nodes)
+
+# 4. 构建查询引擎
+sentence_query_engine = sentence_index.as_query_engine(
+    similarity_top_k=2,
+    node_postprocessors=[
+        MetadataReplacementPostProcessor(target_metadata_key="window")
+    ],
+)
+base_query_engine = base_index.as_query_engine(similarity_top_k=2)
+
+# 5. 执行查询并对比结果
+query = "What are the concerns surrounding the AMOC?"
+print(f"查询: {query}\n")
+
+print("--- 句子窗口检索结果 ---")
+window_response = sentence_query_engine.query(query)
+print(f"回答: {window_response}\n")
+
+print("--- 常规检索结果 ---")
+base_response = base_query_engine.query(query)
+print(f"回答: {base_response}\n")

code/C3/06_recursive_retrieval.py

@@ -0,0 +1,63 @@
+import os
+import pandas as pd
+from dotenv import load_dotenv
+from llama_index.core import VectorStoreIndex
+from llama_index.core.schema import IndexNode
+from llama_index.experimental.query_engine import PandasQueryEngine
+from llama_index.core.retrievers import RecursiveRetriever
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.llms.deepseek import DeepSeek
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core import Settings
+
+load_dotenv()
+
+# 配置模型
+Settings.llm = DeepSeek(model="deepseek-chat", api_key=os.getenv("DEEPSEEK_API_KEY"))
+Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")
+
+# 1.加载数据并为每个工作表创建查询引擎和摘要节点
+excel_file = '../../data/C3/excel/movie.xlsx'
+xls = pd.ExcelFile(excel_file)
+
+df_query_engines = {}
+all_nodes = []
+
+for sheet_name in xls.sheet_names:
+    df = pd.read_excel(xls, sheet_name=sheet_name)
+    
+    # 为当前工作表（DataFrame）创建一个 PandasQueryEngine
+    query_engine = PandasQueryEngine(df=df, llm=Settings.llm, verbose=True)
+    
+    # 为当前工作表创建一个摘要节点（IndexNode）
+    year = sheet_name.replace('年份_', '')
+    summary = f"这个表格包含了年份为 {year} 的电影信息，可以用来回答关于这一年电影的具体问题。"
+    node = IndexNode(text=summary, index_id=sheet_name)
+    all_nodes.append(node)
+    
+    # 存储工作表名称到其查询引擎的映射
+    df_query_engines[sheet_name] = query_engine
+
+# 2. 创建顶层索引（只包含摘要节点）
+vector_index = VectorStoreIndex(all_nodes)
+
+# 3. 创建递归检索器
+# 3.1 创建顶层检索器，用于在摘要节点中检索
+vector_retriever = vector_index.as_retriever(similarity_top_k=1)
+
+# 3.2 创建递归检索器
+recursive_retriever = RecursiveRetriever(
+    "vector",
+    retriever_dict={"vector": vector_retriever},
+    query_engine_dict=df_query_engines,
+    verbose=True,
+)
+
+# 4. 创建查询引擎
+query_engine = RetrieverQueryEngine.from_args(recursive_retriever)
+
+# 5. 执行查询
+query = "1994年评分人数最少的电影是哪一部？"
+print(f"查询: {query}")
+response = query_engine.query(query)
+print(f"回答: {response}")

code/C3/07_recursive_retrieval_v2.py

@@ -0,0 +1,101 @@
+import os
+import pandas as pd
+from dotenv import load_dotenv
+from llama_index.core import VectorStoreIndex, Document, Settings
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter
+from llama_index.llms.deepseek import DeepSeek
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+load_dotenv()
+
+# 配置模型
+Settings.llm = DeepSeek(model="deepseek-chat", api_key=os.getenv("DEEPSEEK_API_KEY"))
+Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")
+
+# 1. 加载和预处理数据
+excel_file = '../../data/C3/excel/movie.xlsx'
+xls = pd.ExcelFile(excel_file)
+
+summary_docs = []
+content_docs = []
+
+print("开始加载和处理Excel文件...")
+for sheet_name in xls.sheet_names:
+    df = pd.read_excel(xls, sheet_name=sheet_name)
+    
+    # 数据清洗
+    if '评分人数' in df.columns:
+        df['评分人数'] = df['评分人数'].astype(str).str.replace('人评价', '').str.strip()
+        df['评分人数'] = pd.to_numeric(df['评分人数'], errors='coerce').fillna(0).astype(int)
+
+    # 创建摘要文档 (用于路由)
+    year = sheet_name.replace('年份_', '')
+    summary_text = f"这个表格包含了年份为 {year} 的电影信息，包括电影名称、导演、评分、评分人数等。"
+    summary_doc = Document(
+        text=summary_text,
+        metadata={"sheet_name": sheet_name}
+    )
+    summary_docs.append(summary_doc)
+    
+    # 创建内容文档 (用于最终问答)
+    content_text = df.to_string(index=False)
+    content_doc = Document(
+        text=content_text,
+        metadata={"sheet_name": sheet_name}
+    )
+    content_docs.append(content_doc)
+
+print("数据加载和处理完成。\n")
+
+# 2. 构建向量索引
+# 使用默认的内存SimpleVectorStore，它支持元数据过滤
+
+# 2.1 为摘要创建索引
+summary_index = VectorStoreIndex(summary_docs)
+
+# 2.2 为内容创建索引
+content_index = VectorStoreIndex(content_docs)
+
+print("摘要索引和内容索引构建完成。\n")
+
+# 3. 定义两步式查询逻辑
+def query_safe_recursive(query_str):
+    print(f"--- 开始执行查询 ---")
+    print(f"查询: {query_str}")
+    
+    # 第一步：路由 - 在摘要索引中找到最相关的表格
+    print("\n第一步：在摘要索引中进行路由...")
+    summary_retriever = VectorIndexRetriever(index=summary_index, similarity_top_k=1)
+    retrieved_nodes = summary_retriever.retrieve(query_str)
+    
+    if not retrieved_nodes:
+        return "抱歉，未能找到相关的电影年份信息。"
+    
+    # 获取匹配到的工作表名称
+    matched_sheet_name = retrieved_nodes[0].node.metadata['sheet_name']
+    print(f"路由结果：匹配到工作表 -> {matched_sheet_name}")
+    
+    # 第二步：检索 - 在内容索引中根据工作表名称过滤并检索具体内容
+    print("\n第二步：在内容索引中检索具体信息...")
+    content_retriever = VectorIndexRetriever(
+        index=content_index,
+        similarity_top_k=1, # 通常只返回最匹配的整个表格即可
+        filters=MetadataFilters(
+            filters=[ExactMatchFilter(key="sheet_name", value=matched_sheet_name)]
+        )
+    )
+    
+    # 创建查询引擎并执行查询
+    query_engine = RetrieverQueryEngine.from_args(content_retriever)
+    response = query_engine.query(query_str)
+    
+    print("--- 查询执行结束 ---\n")
+    return response
+
+# 4. 执行查询
+query = "1994年评分人数最少的电影是哪一部？"
+response = query_safe_recursive(query)
+
+print(f"最终回答: {response}")

code/C3/download_model.py

@@ -0,0 +1,68 @@
+import requests
+from pathlib import Path
+
+def download_visualized_bge_model():
+    """
+    下载 Visual BGE 模型权重文件
+    如果模型文件不存在，则从 Hugging Face 下载
+    """
+    # 定义模型路径和下载URL
+    model_dir = Path("../../models/bge")
+    model_file = model_dir / "Visualized_base_en_v1.5.pth"
+    download_url = "https://huggingface.co/BAAI/bge-visualized/resolve/main/Visualized_base_en_v1.5.pth?download=true"
+    
+    # 检查模型文件是否已存在
+    if model_file.exists():
+        print(f"模型文件已存在: {model_file}")
+        print(f"文件大小: {model_file.stat().st_size / (1024*1024):.1f} MB")
+        return str(model_file)
+    
+    # 创建目录
+    model_dir.mkdir(parents=True, exist_ok=True)
+    print(f"创建模型目录: {model_dir}")
+    
+    # 下载模型
+    print(f"开始下载模型...")
+    print(f"下载地址: {download_url}")
+    
+    try:
+        response = requests.get(download_url, stream=True)
+        response.raise_for_status()
+        
+        # 获取文件大小
+        total_size = int(response.headers.get('content-length', 0))
+        downloaded_size = 0
+        
+        with open(model_file, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+                    downloaded_size += len(chunk)
+                    
+                    # 显示下载进度
+                    if total_size > 0:
+                        progress = (downloaded_size / total_size) * 100
+                        print(f"\r下载进度: {progress:.1f}% ({downloaded_size/(1024*1024):.1f}/{total_size/(1024*1024):.1f} MB)", end='')
+        
+        print(f"\n模型下载完成: {model_file}")
+        print(f"文件大小: {model_file.stat().st_size / (1024*1024):.1f} MB")
+        return str(model_file)
+        
+    except requests.exceptions.RequestException as e:
+        print(f"下载失败: {e}")
+        # 如果下载失败，删除不完整的文件
+        if model_file.exists():
+            model_file.unlink()
+        return None
+    except Exception as e:
+        print(f"发生错误: {e}")
+        if model_file.exists():
+            model_file.unlink()
+        return None
+
+if __name__ == "__main__":
+    model_path = download_visualized_bge_model()
+    if model_path:
+        print(f"✅ 模型准备就绪: {model_path}")
+    else:
+        print("❌ 模型下载失败")

code/C3/visual_bge/__init__.py

@@ -0,0 +1 @@
+from .visual_bge.modeling import Visualized_BGE

code/C3/visual_bge/setup.py

@@ -0,0 +1,18 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="visual_bge",
+    version="0.1.0",
+    description='visual_bge',
+    long_description="./README.md",
+    long_description_content_type="text/markdown",
+    url='https://github.com/FlagOpen/FlagEmbedding/tree/master/research/visual_bge',
+    packages=find_packages(),
+    install_requires=[
+        'torchvision',
+        'timm',
+        'einops',
+        'ftfy'
+    ],
+    python_requires='>=3.6',
+)

code/C3/visual_bge/visual_bge/eva_clip/__init__.py

@@ -0,0 +1,11 @@
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_eva_vision_and_transforms
+from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+from .loss import ClipLoss
+from .model import CLIP, CustomCLIP, CLIPTextCfg, CLIPVisionCfg,\
+    convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
+from .openai import load_openai_model, list_openai_models
+from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model,\
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from .tokenizer import SimpleTokenizer, tokenize
+from .transform import image_transform

code/C3/visual_bge/visual_bge/eva_clip/constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

code/C3/visual_bge/visual_bge/eva_clip/eva_vit_model.py

@@ -0,0 +1,532 @@
+# --------------------------------------------------------
+# Adapted from  https://github.com/microsoft/unilm/tree/master/beit
+# --------------------------------------------------------
+import math
+import os
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+except:
+    from timm.layers import drop_path, to_2tuple, trunc_normal_
+    
+from .transformer import PatchDropout
+from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
+
+if os.getenv('ENV_TYPE') == 'deepspeed':
+    try:
+        from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
+    except:
+        from torch.utils.checkpoint import checkpoint
+else:
+    from torch.utils.checkpoint import checkpoint
+
+try:
+    import xformers.ops as xops
+except ImportError:
+    xops = None
+    # print("Please 'pip install xformers'")
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self, 
+        in_features, 
+        hidden_features=None, 
+        out_features=None, 
+        act_layer=nn.GELU, 
+        norm_layer=nn.LayerNorm, 
+        drop=0.,
+        subln=False,
+
+        ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        # x = self.drop(x)
+        # commit this for the orignal BERT implement 
+        x = self.ffn_ln(x)
+
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+class SwiGLU(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.SiLU, drop=0., 
+                norm_layer=nn.LayerNorm, subln=False):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.w1 = nn.Linear(in_features, hidden_features)
+        self.w2 = nn.Linear(in_features, hidden_features)
+
+        self.act = act_layer()
+        self.ffn_ln = norm_layer(hidden_features) if subln else nn.Identity()
+        self.w3 = nn.Linear(hidden_features, out_features)
+        
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x1 = self.w1(x)
+        x2 = self.w2(x)
+        hidden = self.act(x1) * x2
+        x = self.ffn_ln(hidden)
+        x = self.w3(x)
+        x = self.drop(x)
+        return x
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., window_size=None, attn_head_dim=None, xattn=False, rope=None, subln=False, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.subln = subln
+        if self.subln:
+            self.q_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.k_proj = nn.Linear(dim, all_head_dim, bias=False)
+            self.v_proj = nn.Linear(dim, all_head_dim, bias=False)
+        else:
+            self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        if window_size:
+            self.window_size = window_size
+            self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+            self.relative_position_bias_table = nn.Parameter(
+                torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+            # cls to token & token 2 cls & cls to cls
+
+            # get pair-wise relative position index for each token inside the window
+            coords_h = torch.arange(window_size[0])
+            coords_w = torch.arange(window_size[1])
+            coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+            coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+            relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+            relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+            relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+            relative_coords[:, :, 1] += window_size[1] - 1
+            relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+            relative_position_index = \
+                torch.zeros(size=(window_size[0] * window_size[1] + 1, ) * 2, dtype=relative_coords.dtype)
+            relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+            relative_position_index[0, 0:] = self.num_relative_distance - 3
+            relative_position_index[0:, 0] = self.num_relative_distance - 2
+            relative_position_index[0, 0] = self.num_relative_distance - 1
+
+            self.register_buffer("relative_position_index", relative_position_index)
+        else:
+            self.window_size = None
+            self.relative_position_bias_table = None
+            self.relative_position_index = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.inner_attn_ln = norm_layer(all_head_dim) if subln else nn.Identity()
+        # self.proj = nn.Linear(all_head_dim, all_head_dim)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn
+        self.xattn_drop = attn_drop
+
+        self.rope = rope
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        B, N, C = x.shape
+        if self.subln: 
+            q = F.linear(input=x, weight=self.q_proj.weight, bias=self.q_bias)
+            k = F.linear(input=x, weight=self.k_proj.weight, bias=None)
+            v = F.linear(input=x, weight=self.v_proj.weight, bias=self.v_bias)
+
+            q = q.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)     # B, num_heads, N, C
+            k = k.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3)  
+            v = v.reshape(B, N, self.num_heads, -1).permute(0, 2, 1, 3) 
+        else: 
+
+            qkv_bias = None
+            if self.q_bias is not None:
+                qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+            
+            qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+            qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)   # 3, B, num_heads, N, C
+            q, k, v = qkv[0], qkv[1], qkv[2]
+
+        if self.rope:
+            # slightly fast impl
+            q_t = q[:, :, 1:, :]
+            ro_q_t = self.rope(q_t)
+            q = torch.cat((q[:, :, :1, :], ro_q_t), -2).type_as(v)
+
+            k_t = k[:, :, 1:, :]
+            ro_k_t = self.rope(k_t)
+            k = torch.cat((k[:, :, :1, :], ro_k_t), -2).type_as(v)
+
+        if xops is not None:
+            q = q.permute(0, 2, 1, 3)   # B, num_heads, N, C -> B, N, num_heads, C
+            k = k.permute(0, 2, 1, 3)
+            v = v.permute(0, 2, 1, 3)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale,
+                )
+            x = x.reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        else:
+            q = q * self.scale
+            attn = (q @ k.transpose(-2, -1))
+
+            if self.relative_position_bias_table is not None:
+                relative_position_bias = \
+                    self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                        self.window_size[0] * self.window_size[1] + 1,
+                        self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+                relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+                attn = attn + relative_position_bias.unsqueeze(0).type_as(attn)
+
+            if rel_pos_bias is not None:
+                attn = attn + rel_pos_bias.type_as(attn)
+
+            if attn_mask is not None:
+                attn_mask = attn_mask.bool()
+                attn = attn.masked_fill(~attn_mask[:, None, None, :], float("-inf"))
+            
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+            x = self.inner_attn_ln(x)
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 window_size=None, attn_head_dim=None, xattn=False, rope=None, postnorm=False,
+                 subln=False, naiveswiglu=False):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, window_size=window_size, attn_head_dim=attn_head_dim,
+            xattn=xattn, rope=rope, subln=subln, norm_layer=norm_layer)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+
+        if naiveswiglu:
+            self.mlp = SwiGLU(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                subln=subln,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.mlp = Mlp(
+                in_features=dim, 
+                hidden_features=mlp_hidden_dim, 
+                act_layer=act_layer,
+                subln=subln,
+                drop=drop
+            )
+
+        if init_values is not None and init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+        self.postnorm = postnorm
+
+    def forward(self, x, rel_pos_bias=None, attn_mask=None):
+        if self.gamma_1 is None:
+            if self.postnorm:
+                x = x + self.drop_path(self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            if self.postnorm:
+                x = x + self.drop_path(self.gamma_1 * self.norm1(self.attn(x, rel_pos_bias=rel_pos_bias, attn_mask=attn_mask)))
+                x = x + self.drop_path(self.gamma_2 * self.norm2(self.mlp(x)))
+            else:
+                x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias, attn_mask=attn_mask))
+                x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x, **kwargs):
+        B, C, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2) # [10, 3, 224, 224] -> [10, 196, 768]
+        return x
+
+
+class RelativePositionBias(nn.Module):
+
+    def __init__(self, window_size, num_heads):
+        super().__init__()
+        self.window_size = window_size
+        self.num_relative_distance = (2 * window_size[0] - 1) * (2 * window_size[1] - 1) + 3
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros(self.num_relative_distance, num_heads))  # 2*Wh-1 * 2*Ww-1, nH
+        # cls to token & token 2 cls & cls to cls
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(window_size[0])
+        coords_w = torch.arange(window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2, Wh*Ww
+        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * window_size[1] - 1
+        relative_position_index = \
+            torch.zeros(size=(window_size[0] * window_size[1] + 1,) * 2, dtype=relative_coords.dtype)
+        relative_position_index[1:, 1:] = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        relative_position_index[0, 0:] = self.num_relative_distance - 3
+        relative_position_index[0:, 0] = self.num_relative_distance - 2
+        relative_position_index[0, 0] = self.num_relative_distance - 1
+
+        self.register_buffer("relative_position_index", relative_position_index)
+
+    def forward(self):
+        relative_position_bias = \
+            self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
+                self.window_size[0] * self.window_size[1] + 1,
+                self.window_size[0] * self.window_size[1] + 1, -1)  # Wh*Ww,Wh*Ww,nH
+        return relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
+
+
+class EVAVisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, patch_dropout=0.,
+                 use_abs_pos_emb=True, use_rel_pos_bias=False, use_shared_rel_pos_bias=False, rope=False,
+                 use_mean_pooling=True, init_scale=0.001, grad_checkpointing=False, xattn=False, postnorm=False,
+                 pt_hw_seq_len=16, intp_freq=False, naiveswiglu=False, subln=False):
+        super().__init__()
+        self.image_size = img_size
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        # self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        if use_abs_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        else:
+            self.pos_embed = None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        if use_shared_rel_pos_bias:
+            self.rel_pos_bias = RelativePositionBias(window_size=self.patch_embed.patch_shape, num_heads=num_heads)
+        else:
+            self.rel_pos_bias = None
+
+        if rope:
+            half_head_dim = embed_dim // num_heads // 2
+            hw_seq_len = img_size // patch_size
+            self.rope = VisionRotaryEmbeddingFast(
+                dim=half_head_dim,
+                pt_seq_len=pt_hw_seq_len,
+                ft_seq_len=hw_seq_len if intp_freq else None,
+                # patch_dropout=patch_dropout
+            )
+        else: 
+            self.rope = None
+
+        self.naiveswiglu = naiveswiglu
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.use_rel_pos_bias = use_rel_pos_bias
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values, window_size=self.patch_embed.patch_shape if use_rel_pos_bias else None,
+                xattn=xattn, rope=self.rope, postnorm=postnorm, subln=subln, naiveswiglu=naiveswiglu)
+            for i in range(depth)])
+        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=.02)
+
+        trunc_normal_(self.cls_token, std=.02)
+        # trunc_normal_(self.mask_token, std=.02)
+
+        self.apply(self._init_weights)
+        self.fix_init_weight()
+
+        if isinstance(self.head, nn.Linear):
+            trunc_normal_(self.head.weight, std=.02)
+            self.head.weight.data.mul_(init_scale)
+            self.head.bias.data.mul_(init_scale)
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+
+        self.grad_checkpointing = grad_checkpointing
+
+    def fix_init_weight(self):
+        def rescale(param, layer_id):
+            param.div_(math.sqrt(2.0 * layer_id))
+
+        for layer_id, layer in enumerate(self.blocks):
+            rescale(layer.attn.proj.weight.data, layer_id + 1)
+            if self.naiveswiglu:
+                rescale(layer.mlp.w3.weight.data, layer_id + 1)
+            else:
+                rescale(layer.mlp.fc2.weight.data, layer_id + 1)
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.blocks[0].mlp.fc2.weight.dtype
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_num_layers(self):
+        return len(self.blocks)
+    
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        for param in self.parameters():
+            param.requires_grad = False
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def reset_classifier(self, num_classes, global_pool=''):
+        self.num_classes = num_classes
+        self.head = nn.Linear(self.embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+    def forward_features(self, x, return_all_features=False):
+        
+        x = self.patch_embed(x)
+        batch_size, seq_len, _ = x.size()
+
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+        x = torch.cat((cls_tokens, x), dim=1)
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        if os.getenv('RoPE') == '1':
+            if self.training and not isinstance(self.patch_dropout, nn.Identity):
+                x, patch_indices_keep = self.patch_dropout(x)
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=patch_indices_keep)
+            else:
+                self.rope.forward = partial(self.rope.forward, patch_indices_keep=None)
+                x = self.patch_dropout(x)
+        else:
+            x = self.patch_dropout(x)
+
+        rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None
+        for blk in self.blocks:
+            if self.grad_checkpointing:
+                # x = checkpoint(blk, x, (rel_pos_bias,))
+                x = checkpoint(blk, x, rel_pos_bias)
+            else:
+                x = blk(x, rel_pos_bias=rel_pos_bias)
+
+        if not return_all_features:
+            x = self.norm(x)
+            if self.fc_norm is not None:
+                return self.fc_norm(x.mean(1))
+            else:
+                return x[:, 0]
+        return x
+
+    def forward(self, x, return_all_features=True):
+        if return_all_features:
+            return self.forward_features(x, return_all_features)
+        x = self.forward_features(x)
+        x = self.head(x)
+        return x

code/C3/visual_bge/visual_bge/eva_clip/factory.py

@@ -0,0 +1,519 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Optional, Tuple, Union, Dict, Any
+import torch
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from .model import CLIP, CustomCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+    get_cast_dtype
+from .openai import load_openai_model
+from .pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained, list_pretrained_tags_by_model
+from .transform import image_transform
+from .tokenizer import HFTokenizer, tokenize
+from .utils import resize_clip_pos_embed, resize_evaclip_pos_embed, resize_visual_pos_embed, resize_eva_pos_embed
+
+
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+
+    for cf in config_files:
+        with open(cf, "r", encoding="utf8") as f:
+            model_cfg = json.load(f)
+            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+                _MODEL_CONFIGS[cf.stem] = model_cfg
+
+    _MODEL_CONFIGS = dict(sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0])))
+
+
+_rescan_model_configs()  # initial populate of model config registry
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+def get_tokenizer(model_name):
+    config = get_model_config(model_name)
+    tokenizer = HFTokenizer(config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+    return tokenizer
+
+
+# loading openai CLIP weights when is_openai=True for training
+def load_state_dict(checkpoint_path: str, map_location: str='cpu', model_key: str='model|module|state_dict', is_openai: bool=False, skip_list: list=[]):
+    if is_openai:
+        model = torch.jit.load(checkpoint_path, map_location="cpu").eval()
+        state_dict = model.state_dict()
+        for key in ["input_resolution", "context_length", "vocab_size"]:
+            state_dict.pop(key, None)
+    else:
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        for mk in model_key.split('|'):
+            if isinstance(checkpoint, dict) and mk in checkpoint:
+                state_dict = checkpoint[mk]
+                break
+            else:
+                state_dict = checkpoint
+        if next(iter(state_dict.items()))[0].startswith('module'):
+            state_dict = {k[7:]: v for k, v in state_dict.items()}
+    
+    for k in skip_list:
+        if k in list(state_dict.keys()):
+            logging.info(f"Removing key {k} from pretrained checkpoint")
+            del state_dict[k]
+
+    if os.getenv('RoPE') == '1':
+        for k in list(state_dict.keys()):
+            if 'freqs_cos' in k or 'freqs_sin' in k:
+                del state_dict[k]
+    return state_dict
+
+
+
+def load_checkpoint(model, checkpoint_path, model_key="model|module|state_dict", strict=True):
+    state_dict = load_state_dict(checkpoint_path, model_key=model_key, is_openai=False)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    if 'text.logit_scale' in state_dict and hasattr(model, 'logit_scale'):
+        state_dict['logit_scale'] = state_dict['text.logit_scale']
+        del state_dict['text.logit_scale']
+
+    # resize_clip_pos_embed for CLIP and open CLIP
+    if 'visual.positional_embedding' in state_dict:
+        resize_clip_pos_embed(state_dict, model)
+    # specified to eva_vit_model
+    elif 'visual.pos_embed' in state_dict:
+        resize_evaclip_pos_embed(state_dict, model)
+
+    # resize_clip_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+    logging.info(f"incompatible_keys.missing_keys: {incompatible_keys.missing_keys}")
+    return incompatible_keys
+
+def load_clip_visual_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
+    state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+
+    for k in list(state_dict.keys()):
+        if not k.startswith('visual.'):
+            del state_dict[k]
+    for k in list(state_dict.keys()):
+        if k.startswith('visual.'):
+            new_k = k[7:]
+            state_dict[new_k] = state_dict[k]
+            del state_dict[k]
+    return state_dict
+
+def load_clip_text_state_dict(checkpoint_path: str, map_location: str='cpu', is_openai: bool=False, skip_list:list=[]):
+    state_dict = load_state_dict(checkpoint_path, map_location=map_location, is_openai=is_openai, skip_list=skip_list)
+
+    for k in list(state_dict.keys()):
+        if k.startswith('visual.'):
+            del state_dict[k]
+    return state_dict
+
+def get_pretrained_tag(pretrained_model):
+    pretrained_model = pretrained_model.lower()
+    if "laion" in pretrained_model or "open_clip" in pretrained_model:
+        return "open_clip"
+    elif "openai" in pretrained_model:
+        return "clip"
+    elif "eva" in pretrained_model and "clip" in pretrained_model:
+        return "eva_clip"
+    else:
+        return "other"
+
+def load_pretrained_checkpoint(
+        model,
+        visual_checkpoint_path,
+        text_checkpoint_path,
+        strict=True,
+        visual_model=None,
+        text_model=None,
+        model_key="model|module|state_dict",
+        skip_list=[]):
+    visual_tag = get_pretrained_tag(visual_model)
+    text_tag = get_pretrained_tag(text_model)
+
+    logging.info(f"num of model state_dict keys: {len(model.state_dict().keys())}")
+    visual_incompatible_keys, text_incompatible_keys = None, None
+    if visual_checkpoint_path:
+        if visual_tag == "eva_clip" or visual_tag == "open_clip":
+            visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=False, skip_list=skip_list)
+        elif visual_tag == "clip":
+            visual_state_dict = load_clip_visual_state_dict(visual_checkpoint_path, is_openai=True, skip_list=skip_list)
+        else:
+            visual_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
+    
+        # resize_clip_pos_embed for CLIP and open CLIP
+        if 'positional_embedding' in visual_state_dict:
+            resize_visual_pos_embed(visual_state_dict, model)
+        # specified to EVA model
+        elif 'pos_embed' in visual_state_dict:
+            resize_eva_pos_embed(visual_state_dict, model)
+
+        visual_incompatible_keys = model.visual.load_state_dict(visual_state_dict, strict=strict)
+        logging.info(f"num of loaded visual_state_dict keys: {len(visual_state_dict.keys())}")
+        logging.info(f"visual_incompatible_keys.missing_keys: {visual_incompatible_keys.missing_keys}")
+
+    if text_checkpoint_path:
+        if text_tag == "eva_clip" or text_tag == "open_clip":
+            text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=False, skip_list=skip_list)
+        elif text_tag == "clip":
+            text_state_dict = load_clip_text_state_dict(text_checkpoint_path, is_openai=True, skip_list=skip_list)
+        else:
+            text_state_dict = load_state_dict(visual_checkpoint_path, model_key=model_key, is_openai=False, skip_list=skip_list)
+
+        text_incompatible_keys = model.text.load_state_dict(text_state_dict, strict=strict)
+        
+        logging.info(f"num of loaded text_state_dict keys: {len(text_state_dict.keys())}")
+        logging.info(f"text_incompatible_keys.missing_keys: {text_incompatible_keys.missing_keys}")
+
+    return visual_incompatible_keys, text_incompatible_keys
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_clip: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        pretrained_image: str = '',
+        pretrained_text: str = '',
+        pretrained_hf: bool = True,
+        pretrained_visual_model: str = None,
+        pretrained_text_model: str = None,
+        cache_dir: Optional[str] = None,
+        skip_list: list  = [],
+        is_only_visual: bool = False,
+        is_only_text: bool = False,
+):
+    model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    if pretrained and pretrained.lower() == 'openai':
+        logging.info(f'Loading pretrained {model_name} from OpenAI.')
+        model = load_openai_model(
+            model_name,
+            precision=precision,
+            device=device,
+            jit=jit,
+            cache_dir=cache_dir,
+        )
+    else:
+        model_cfg = get_model_config(model_name)
+        if model_cfg is not None:
+            logging.info(f'Loaded {model_name} model config.')
+        else:
+            logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
+            raise RuntimeError(f'Model config for {model_name} not found.')
+
+        if 'rope' in model_cfg.get('vision_cfg', {}):
+            if model_cfg['vision_cfg']['rope']:
+                os.environ['RoPE'] = "1"
+        else:
+            os.environ['RoPE'] = "0"
+
+        if force_quick_gelu:
+            # override for use of QuickGELU on non-OpenAI transformer models
+            model_cfg["quick_gelu"] = True
+        
+        if force_patch_dropout is not None:
+            # override the default patch dropout value
+            model_cfg['vision_cfg']["patch_dropout"] = force_patch_dropout
+
+        cast_dtype = get_cast_dtype(precision)
+        custom_clip = model_cfg.pop('custom_text', False) or force_custom_clip or ('hf_model_name' in model_cfg['text_cfg'])
+
+
+        if custom_clip:
+            if 'hf_model_name' in model_cfg.get('text_cfg', {}):
+                model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+            model = CustomCLIP(**model_cfg, cast_dtype=cast_dtype, is_only_visual=is_only_visual, is_only_text=is_only_text)
+        else:
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+            print("Not CustomCLIP: If you have set building only visual or text tower, you may still get a complete CLIP model.")
+
+        pretrained_cfg = {}
+        if pretrained:
+            checkpoint_path = ''
+            pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+            if pretrained_cfg:
+                checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            elif os.path.exists(pretrained):
+                checkpoint_path = pretrained
+
+            if checkpoint_path:
+                logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+                load_checkpoint(model,
+                               checkpoint_path,
+                               model_key="model|module|state_dict",
+                               strict=False
+                               ) 
+            else:
+                error_str = (
+                    f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                    f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+                logging.warning(error_str)
+                raise RuntimeError(error_str)
+        else:
+            visual_checkpoint_path = ''
+            text_checkpoint_path = ''
+            
+            if pretrained_image:
+                pretrained_visual_model = pretrained_visual_model.replace('/', '-')  # for callers using old naming with / in ViT names
+                pretrained_image_cfg = get_pretrained_cfg(pretrained_visual_model, pretrained_image)
+                if 'timm_model_name' in model_cfg.get('vision_cfg', {}):
+                    # pretrained weight loading for timm models set via vision_cfg
+                    model_cfg['vision_cfg']['timm_model_pretrained'] = True
+                elif pretrained_image_cfg:
+                    visual_checkpoint_path = download_pretrained(pretrained_image_cfg, cache_dir=cache_dir)
+                elif os.path.exists(pretrained_image):
+                    visual_checkpoint_path = pretrained_image
+                else:
+                    logging.warning(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
+                    raise RuntimeError(f'Pretrained weights ({visual_checkpoint_path}) not found for model {model_name}.visual.')
+
+            if pretrained_text:
+                pretrained_text_model = pretrained_text_model.replace('/', '-')  # for callers using old naming with / in ViT names
+                pretrained_text_cfg = get_pretrained_cfg(pretrained_text_model, pretrained_text)
+                if pretrained_image_cfg:
+                    text_checkpoint_path = download_pretrained(pretrained_text_cfg, cache_dir=cache_dir)
+                elif os.path.exists(pretrained_text):
+                    text_checkpoint_path = pretrained_text
+                else:
+                    logging.warning(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
+                    raise RuntimeError(f'Pretrained weights ({text_checkpoint_path}) not found for model {model_name}.text.')
+            
+            if visual_checkpoint_path:
+                logging.info(f'Loading pretrained {model_name}.visual weights ({visual_checkpoint_path}).')
+            if text_checkpoint_path:
+                logging.info(f'Loading pretrained {model_name}.text weights ({text_checkpoint_path}).')
+
+            if visual_checkpoint_path or text_checkpoint_path:
+                load_pretrained_checkpoint(
+                    model,
+                    visual_checkpoint_path,
+                    text_checkpoint_path,
+                    strict=False,
+                    visual_model=pretrained_visual_model,
+                    text_model=pretrained_text_model,
+                    model_key="model|module|state_dict",
+                    skip_list=skip_list
+                )
+        
+        if "fp16" in precision or "bf16" in precision:
+            logging.info(f'convert precision to {precision}')
+            model = model.to(torch.bfloat16) if 'bf16' in precision else model.to(torch.float16)
+
+        model.to(device=device)
+
+        # set image / mean metadata from pretrained_cfg if available, or use default
+        if not is_only_text:
+            model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+            model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+
+        if jit:
+            model = torch.jit.script(model)
+
+    return model
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_clip: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        pretrained_image: str = '',
+        pretrained_text: str = '',
+        pretrained_hf: bool = True,
+        pretrained_visual_model: str = None,
+        pretrained_text_model: str = None,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+        skip_list: list = [],
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_clip=force_custom_clip,
+        force_patch_dropout=force_patch_dropout,
+        pretrained_image=pretrained_image,
+        pretrained_text=pretrained_text,
+        pretrained_hf=pretrained_hf,
+        pretrained_visual_model=pretrained_visual_model,
+        pretrained_text_model=pretrained_text_model,
+        cache_dir=cache_dir,
+        skip_list=skip_list,
+    )
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess_train = image_transform(
+        model.visual.image_size,
+        is_train=True,
+        mean=image_mean,
+        std=image_std
+    )
+    preprocess_val = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std
+    )
+
+    return model, preprocess_train, preprocess_val
+
+def create_eva_vision_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_clip: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        pretrained_image: str = '',
+        pretrained_text: str = '',
+        pretrained_hf: bool = True,
+        pretrained_visual_model: str = None,
+        pretrained_text_model: str = None,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+        skip_list: list = [],
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_clip=force_custom_clip,
+        force_patch_dropout=force_patch_dropout,
+        pretrained_image=pretrained_image,
+        pretrained_text=pretrained_text,
+        pretrained_hf=pretrained_hf,
+        pretrained_visual_model=pretrained_visual_model,
+        pretrained_text_model=pretrained_text_model,
+        cache_dir=cache_dir,
+        skip_list=skip_list,
+        is_only_visual=True, # only use visual tower
+    )
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess_train = image_transform(
+        model.visual.image_size,
+        is_train=True,
+        mean=image_mean,
+        std=image_std
+    )
+    preprocess_val = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std
+    )
+
+    return model, preprocess_train, preprocess_val
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: str,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_clip: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+        is_frozen: bool = False,
+):
+    if not is_pretrained_cfg(model_name, pretrained) and not os.path.exists(pretrained):
+        raise RuntimeError(
+            f'{pretrained} is not a valid pretrained cfg or checkpoint for {model_name}.'
+            f' Use open_clip.list_pretrained() to find one.')
+
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_clip=force_custom_clip,
+        force_patch_dropout=force_patch_dropout,
+        cache_dir=cache_dir,
+    )
+
+    if is_frozen:
+        for param in model.parameters():
+            param.requires_grad = False
+
+    if not return_transform:
+        return model
+
+    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+    image_std = image_std or getattr(model.visual, 'image_std', None)
+    preprocess = image_transform(
+        model.visual.image_size,
+        is_train=False,
+        mean=image_mean,
+        std=image_std
+    )
+
+    return model, preprocess

code/C3/visual_bge/visual_bge/eva_clip/hf_configs.py

@@ -0,0 +1,57 @@
+# HF architecture dict:
+arch_dict = {
+  # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
+  "roberta": {
+      "config_names": {
+          "context_length": "max_position_embeddings",
+          "vocab_size": "vocab_size",
+          "width": "hidden_size",
+          "heads": "num_attention_heads",
+          "layers": "num_hidden_layers",
+          "layer_attr": "layer",
+          "token_embeddings_attr": "embeddings"
+      },
+      "pooler": "mean_pooler",
+  },
+  # https://huggingface.co/docs/transformers/model_doc/xlm-roberta#transformers.XLMRobertaConfig
+  "xlm-roberta": {
+      "config_names": {
+          "context_length": "max_position_embeddings",
+          "vocab_size": "vocab_size",
+          "width": "hidden_size",
+          "heads": "num_attention_heads",
+          "layers": "num_hidden_layers",
+          "layer_attr": "layer",
+          "token_embeddings_attr": "embeddings"
+      },
+      "pooler": "mean_pooler",
+  },
+  # https://huggingface.co/docs/transformers/model_doc/mt5#mt5
+  "mt5": {
+      "config_names": {
+          # unlimited seqlen
+          # https://github.com/google-research/text-to-text-transfer-transformer/issues/273
+          # https://github.com/huggingface/transformers/blob/v4.24.0/src/transformers/models/t5/modeling_t5.py#L374
+          "context_length": "",
+          "vocab_size": "vocab_size",
+          "width": "d_model",
+          "heads": "num_heads",
+          "layers": "num_layers",
+          "layer_attr": "block",
+          "token_embeddings_attr": "embed_tokens"
+      },
+      "pooler": "mean_pooler",
+  },
+  "bert": {
+    "config_names": {
+      "context_length": "max_position_embeddings",
+      "vocab_size": "vocab_size",
+      "width": "hidden_size",
+      "heads": "num_attention_heads",
+      "layers": "num_hidden_layers",
+      "layer_attr": "layer",
+      "token_embeddings_attr": "embeddings"
+    },
+    "pooler": "mean_pooler",
+  }
+}

code/C3/visual_bge/visual_bge/eva_clip/hf_model.py

@@ -0,0 +1,248 @@
+""" huggingface model adapter
+
+Wraps HuggingFace transformers (https://github.com/huggingface/transformers) models for use as a text tower in CLIP model.
+"""
+
+import re
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch import TensorType
+try:
+    import transformers
+    from transformers import AutoModel, AutoModelForMaskedLM, AutoTokenizer, AutoConfig, PretrainedConfig
+    from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, \
+        BaseModelOutputWithPoolingAndCrossAttentions
+except ImportError as e:
+    transformers = None
+
+
+    class BaseModelOutput:
+        pass
+
+
+    class PretrainedConfig:
+        pass
+
+from .hf_configs import arch_dict
+
+# utils
+def _camel2snake(s):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', s).lower()
+
+# TODO: ?last - for gpt-like models
+_POOLERS = {}
+
+def register_pooler(cls):
+    """Decorator registering pooler class"""
+    _POOLERS[_camel2snake(cls.__name__)] = cls
+    return cls
+
+
+@register_pooler
+class MeanPooler(nn.Module):
+    """Mean pooling"""
+    def forward(self, x:BaseModelOutput, attention_mask:TensorType):
+        masked_output = x.last_hidden_state * attention_mask.unsqueeze(-1)
+        return masked_output.sum(dim=1) / attention_mask.sum(-1, keepdim=True)
+
+@register_pooler
+class MaxPooler(nn.Module):
+    """Max pooling"""
+    def forward(self, x:BaseModelOutput, attention_mask:TensorType):
+        masked_output = x.last_hidden_state.masked_fill(attention_mask.unsqueeze(-1), -torch.inf)
+        return masked_output.max(1).values
+
+@register_pooler
+class ClsPooler(nn.Module):
+    """CLS token pooling"""
+    def __init__(self, use_pooler_output=True):
+        super().__init__()
+        self.cls_token_position = 0
+        self.use_pooler_output = use_pooler_output
+
+    def forward(self, x:BaseModelOutput, attention_mask:TensorType):
+        
+        if (self.use_pooler_output and 
+            isinstance(x, (BaseModelOutputWithPooling, BaseModelOutputWithPoolingAndCrossAttentions)) and
+            (x.pooler_output is not None)
+            ):
+            return x.pooler_output
+        
+        return x.last_hidden_state[:, self.cls_token_position, :]
+
+class HFTextEncoder(nn.Module):
+    """HuggingFace model adapter"""
+    def __init__(
+            self, 
+            model_name_or_path: str,
+            output_dim: int,
+            tokenizer_name: str = None,
+            config: PretrainedConfig = None,
+            pooler_type: str = None,
+            proj: str = None,
+            pretrained: bool = True,
+            masked_language_modeling: bool = False):
+        super().__init__()
+
+        self.output_dim = output_dim
+
+        # TODO: find better way to get this information
+        uses_transformer_pooler = (pooler_type == "cls_pooler")
+
+        if transformers is None:
+            raise RuntimeError("Please `pip install transformers` to use pre-trained HuggingFace models")
+        if config is None:
+            self.config = AutoConfig.from_pretrained(model_name_or_path)
+            if masked_language_modeling:
+                create_func, model_args = (AutoModelForMaskedLM.from_pretrained, model_name_or_path) if pretrained else (
+                    AutoModelForMaskedLM.from_config, self.config)
+            else:
+                create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else (
+                    AutoModel.from_config, self.config)
+            # TODO: do all model configs have this attribute? PretrainedConfig does so yes??
+            if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder:
+                self.transformer = create_func(model_args)
+                self.transformer = self.transformer.encoder
+            else:
+                self.transformer = create_func(model_args, add_pooling_layer=uses_transformer_pooler)
+        else:
+            self.config = config
+            if masked_language_modeling:
+                self.transformer = AutoModelForMaskedLM.from_config(config)
+            else:
+                self.transformer = AutoModel.from_config(config)
+
+        if pooler_type is None: # get default arch pooler
+            self.pooler = _POOLERS[(arch_dict[self.config.model_type]["pooler"])]()
+        else:
+            self.pooler = _POOLERS[pooler_type]()
+
+        d_model = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["width"])
+        if (d_model == output_dim) and (proj is None): # do we always need a proj?
+            self.proj = nn.Identity()
+        elif proj == 'linear':
+            self.proj = nn.Linear(d_model, output_dim, bias=False)
+        elif proj == 'mlp':
+            hidden_size = (d_model + output_dim) // 2
+            self.proj = nn.Sequential(
+                nn.Linear(d_model, hidden_size, bias=False),
+                nn.GELU(),
+                nn.Linear(hidden_size, output_dim, bias=False),
+            )
+
+        # self.itm_proj = nn.Linear(d_model, 2, bias=False)
+        # self.mlm_proj = nn.Linear(d_model, self.config.vocab_size), bias=False)
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+    # def forward_itm(self, x:TensorType, image_embeds:TensorType) -> TensorType:
+    #     image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(x.device)  
+    #     attn_mask = (x != self.config.pad_token_id).long()
+    #     out = self.transformer(
+    #         input_ids=x, 
+    #         attention_mask=attn_mask,
+    #         encoder_hidden_states = image_embeds,
+    #         encoder_attention_mask = image_atts,
+    #         )
+    #     pooled_out = self.pooler(out, attn_mask)
+
+    #     return self.itm_proj(pooled_out)
+
+    def mask(self, input_ids, vocab_size, device, targets=None, masked_indices=None, probability_matrix=None):
+        if masked_indices is None:                                       
+            masked_indices = torch.bernoulli(probability_matrix).bool()
+                                               
+        masked_indices[input_ids == self.tokenizer.pad_token_id] = False
+        masked_indices[input_ids == self.tokenizer.cls_token_id] = False
+        
+        if targets is not None:
+            targets[~masked_indices] = -100 # We only compute loss on masked tokens            
+
+        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+        indices_replaced = torch.bernoulli(torch.full(input_ids.shape, 0.8)).bool() & masked_indices
+        input_ids[indices_replaced] = self.tokenizer.mask_token_id
+
+        # 10% of the time, we replace masked input tokens with random word
+        indices_random = torch.bernoulli(torch.full(input_ids.shape, 0.5)).bool() & masked_indices & ~indices_replaced
+        random_words = torch.randint(vocab_size, input_ids.shape, dtype=torch.long).to(device)
+        input_ids[indices_random] = random_words[indices_random]                     
+        # The rest of the time (10% of the time) we keep the masked input tokens unchanged   
+        
+        if targets is not None:
+            return input_ids, targets
+        else:
+            return input_ids
+
+    def forward_mlm(self, input_ids, image_embeds, mlm_probability=0.25):
+        labels = input_ids.clone()
+        attn_mask = (input_ids != self.config.pad_token_id).long()
+        image_atts = torch.ones(image_embeds.size()[:-1],dtype=torch.long).to(input_ids.device) 
+        vocab_size = getattr(self.config, arch_dict[self.config.model_type]["config_names"]["vocab_size"])
+        probability_matrix = torch.full(labels.shape, mlm_probability)
+        input_ids, labels = self.mask(input_ids, vocab_size, input_ids.device, targets=labels,
+                                      probability_matrix = probability_matrix)
+        mlm_output = self.transformer(input_ids,
+                        attention_mask = attn_mask,
+                        encoder_hidden_states = image_embeds,
+                        encoder_attention_mask = image_atts,
+                        return_dict = True,
+                        labels = labels,
+                    )
+        return mlm_output.loss
+        # mlm_output = self.transformer(input_ids,
+        #                 attention_mask = attn_mask,
+        #                 encoder_hidden_states = image_embeds,
+        #                 encoder_attention_mask = image_atts,
+        #                 return_dict = True,
+        #             ).last_hidden_state
+        # logits = self.mlm_proj(mlm_output)
+
+        # # logits = logits[:, :-1, :].contiguous().view(-1, vocab_size)
+        # logits = logits[:, 1:, :].contiguous().view(-1, vocab_size)
+        # labels = labels[:, 1:].contiguous().view(-1)
+
+        # mlm_loss = F.cross_entropy(
+        #     logits,
+        #     labels,
+        #     # label_smoothing=0.1,
+        # )
+        # return mlm_loss
+
+
+    def forward(self, x:TensorType) -> TensorType:
+        attn_mask = (x != self.config.pad_token_id).long()
+        out = self.transformer(input_ids=x, attention_mask=attn_mask)
+        pooled_out = self.pooler(out, attn_mask)
+
+        return self.proj(pooled_out)
+
+    def lock(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
+        if not unlocked_layers: # full freezing
+             for n, p in self.transformer.named_parameters():
+                 p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+             return
+
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        print(f"Unlocking {unlocked_layers}/{len(layer_list) + 1} layers of hf model")
+        embeddings = getattr(
+            self.transformer, arch_dict[self.config.model_type]["config_names"]["token_embeddings_attr"])
+        modules = [embeddings, *layer_list][:-unlocked_layers]
+        # freeze layers
+        for module in modules:
+            for n, p in module.named_parameters():
+                p.requires_grad = (not freeze_layer_norm) if "LayerNorm" in n.split(".") else False
+
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.gradient_checkpointing_enable()
+
+    def get_num_layers(self):
+        encoder = self.transformer.encoder if hasattr(self.transformer, 'encoder') else self.transformer
+        layer_list = getattr(encoder, arch_dict[self.config.model_type]["config_names"]["layer_attr"])
+        return len(layer_list)
+
+    def init_parameters(self):
+        pass

code/C3/visual_bge/visual_bge/eva_clip/loss.py

@@ -0,0 +1,138 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+try:
+    import torch.distributed.nn
+    from torch import distributed as dist
+    has_distributed = True
+except ImportError:
+    has_distributed = False
+
+try:
+    import horovod.torch as hvd
+except ImportError:
+    hvd = None
+
+from timm.loss import LabelSmoothingCrossEntropy
+
+
+def gather_features(
+        image_features,
+        text_features,
+        local_loss=False,
+        gather_with_grad=False,
+        rank=0,
+        world_size=1,
+        use_horovod=False
+):
+    assert has_distributed, 'torch.distributed did not import correctly, please use a PyTorch version with support.'
+    if use_horovod:
+        assert hvd is not None, 'Please install horovod'
+        if gather_with_grad:
+            all_image_features = hvd.allgather(image_features)
+            all_text_features = hvd.allgather(text_features)
+        else:
+            with torch.no_grad():
+                all_image_features = hvd.allgather(image_features)
+                all_text_features = hvd.allgather(text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features = list(all_image_features.chunk(world_size, dim=0))
+                gathered_text_features = list(all_text_features.chunk(world_size, dim=0))
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+                all_image_features = torch.cat(gathered_image_features, dim=0)
+                all_text_features = torch.cat(gathered_text_features, dim=0)
+    else:
+        # We gather tensors from all gpus
+        if gather_with_grad:
+            all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features), dim=0)
+            all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features), dim=0)
+            # all_image_features = torch.cat(torch.distributed.nn.all_gather(image_features, async_op=True), dim=0)
+            # all_text_features = torch.cat(torch.distributed.nn.all_gather(text_features, async_op=True), dim=0)
+        else:
+            gathered_image_features = [torch.zeros_like(image_features) for _ in range(world_size)]
+            gathered_text_features = [torch.zeros_like(text_features) for _ in range(world_size)]
+            dist.all_gather(gathered_image_features, image_features)
+            dist.all_gather(gathered_text_features, text_features)
+            if not local_loss:
+                # ensure grads for local rank when all_* features don't have a gradient
+                gathered_image_features[rank] = image_features
+                gathered_text_features[rank] = text_features
+            all_image_features = torch.cat(gathered_image_features, dim=0)
+            all_text_features = torch.cat(gathered_text_features, dim=0)
+
+    return all_image_features, all_text_features
+
+
+class ClipLoss(nn.Module):
+
+    def __init__(
+            self,
+            local_loss=False,
+            gather_with_grad=False,
+            cache_labels=False,
+            rank=0,
+            world_size=1,
+            use_horovod=False,
+            smoothing=0.,
+    ):
+        super().__init__()
+        self.local_loss = local_loss
+        self.gather_with_grad = gather_with_grad
+        self.cache_labels = cache_labels
+        self.rank = rank
+        self.world_size = world_size
+        self.use_horovod = use_horovod
+        self.label_smoothing_cross_entropy = LabelSmoothingCrossEntropy(smoothing=smoothing) if smoothing > 0 else None
+
+        # cache state
+        self.prev_num_logits = 0
+        self.labels = {}
+
+    def forward(self, image_features, text_features, logit_scale=1.):
+        device = image_features.device
+        if self.world_size > 1:
+            all_image_features, all_text_features = gather_features(
+                image_features, text_features,
+                self.local_loss, self.gather_with_grad, self.rank, self.world_size, self.use_horovod)
+
+            if self.local_loss:
+                logits_per_image = logit_scale * image_features @ all_text_features.T
+                logits_per_text = logit_scale * text_features @ all_image_features.T
+            else:
+                logits_per_image = logit_scale * all_image_features @ all_text_features.T
+                logits_per_text = logits_per_image.T
+        else:
+            logits_per_image = logit_scale * image_features @ text_features.T
+            logits_per_text = logit_scale * text_features @ image_features.T
+        # calculated ground-truth and cache if enabled
+        num_logits = logits_per_image.shape[0]
+        if self.prev_num_logits != num_logits or device not in self.labels:
+            labels = torch.arange(num_logits, device=device, dtype=torch.long)
+            if self.world_size > 1 and self.local_loss:
+                labels = labels + num_logits * self.rank
+            if self.cache_labels:
+                self.labels[device] = labels
+                self.prev_num_logits = num_logits
+        else:
+            labels = self.labels[device]
+        
+        if self.label_smoothing_cross_entropy:
+            total_loss = (
+                self.label_smoothing_cross_entropy(logits_per_image, labels) +
+                self.label_smoothing_cross_entropy(logits_per_text, labels)
+                ) / 2
+        else:
+            total_loss = (
+                F.cross_entropy(logits_per_image, labels) +
+                F.cross_entropy(logits_per_text, labels)
+                ) / 2
+            
+        acc = None
+        i2t_acc = (logits_per_image.argmax(-1) == labels).sum() / len(logits_per_image)
+        t2i_acc = (logits_per_text.argmax(-1) == labels).sum() / len(logits_per_text)
+        acc = {"i2t": i2t_acc, "t2i": t2i_acc}
+        return total_loss, acc

code/C3/visual_bge/visual_bge/eva_clip/model.py

@@ -0,0 +1,452 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import os
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+from functools import partial
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+try:
+    from .hf_model import HFTextEncoder
+except:
+    HFTextEncoder = None
+from .modified_resnet import ModifiedResNet
+from .timm_model import TimmModel
+from .eva_vit_model import EVAVisionTransformer
+from .transformer import LayerNorm, QuickGELU, Attention, VisionTransformer, TextTransformer
+
+# try:
+#     from apex.normalization import FusedLayerNorm
+# except:
+FusedLayerNorm = LayerNorm
+    # print("Please 'pip install apex'")
+
+try:
+    import xformers.ops as xops
+except ImportError:
+    xops = None
+    # print("Please 'pip install xformers'")
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    patch_dropout: float = 0. # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    global_average_pool: bool = False # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    drop_path_rate: Optional[float] = None  # drop path rate
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    eva_model_name: str = None # a valid eva model name overrides layers, width, patch_size
+    qkv_bias: bool = True
+    fusedLN: bool = False
+    xattn: bool = False
+    postnorm: bool = False
+    rope: bool = False
+    pt_hw_seq_len: int = 16   # 224/14
+    intp_freq: bool = False
+    naiveswiglu: bool = False
+    subln: bool = False
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    masked_language_modeling: bool = False
+    fusedLN: bool = False
+    xattn: bool = False
+    attn_mask: bool = True
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.eva_model_name:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNorm
+        
+        visual = EVAVisionTransformer(
+            img_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            num_classes=embed_dim,
+            use_mean_pooling=vision_cfg.global_average_pool, #False
+            init_values=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            embed_dim=vision_cfg.width,
+            depth=vision_cfg.layers,
+            num_heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            qkv_bias=vision_cfg.qkv_bias,
+            drop_path_rate=vision_cfg.drop_path_rate,
+            norm_layer= partial(FusedLayerNorm, eps=1e-6) if vision_cfg.fusedLN else partial(norm_layer, eps=1e-6),
+            xattn=vision_cfg.xattn,
+            rope=vision_cfg.rope,
+            postnorm=vision_cfg.postnorm,
+            pt_hw_seq_len= vision_cfg.pt_hw_seq_len,   # 224/14
+            intp_freq= vision_cfg.intp_freq,
+            naiveswiglu= vision_cfg.naiveswiglu,
+            subln= vision_cfg.subln
+        )
+    elif vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size
+        )
+        act_layer = nn.GELU  # so that text transformer doesn't use QuickGELU w/ timm models
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNormFp32 if cast_dtype in (torch.float16, torch.bfloat16) else LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            global_average_pool=vision_cfg.global_average_pool,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            tokenizer_name=text_cfg.hf_tokenizer_name,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            masked_language_modeling=text_cfg.masked_language_modeling
+       )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer= FusedLayerNorm if text_cfg.fusedLN else norm_layer,
+            xattn=text_cfg.xattn,
+            attn_mask=text_cfg.attn_mask,
+        )
+    return text
+
+class CLIP(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+    ):
+        super().__init__()
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+
+        text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.transformer = text.transformer
+        self.vocab_size = text.vocab_size
+        self.token_embedding = text.token_embedding
+        self.positional_embedding = text.positional_embedding
+        self.ln_final = text.ln_final
+        self.text_projection = text.text_projection
+        self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+    
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'logit_scale'}
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image, normalize=True)
+        text_features = self.encode_text(text, normalize=True)
+        return image_features, text_features, self.logit_scale.exp()
+
+
+class CustomCLIP(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            itm_task: bool = False,
+            is_only_visual: bool = False,
+            is_only_text: bool = False,
+    ):
+        super().__init__()
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        self.text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07)) #可学习参数
+        if is_only_visual:
+            self.text = None
+        if is_only_text:
+            self.visual = None
+
+    def lock_image_tower(self, unlocked_groups=0, freeze_bn_stats=False):
+        # lock image tower as per LiT - https://arxiv.org/abs/2111.07991
+        self.visual.lock(unlocked_groups=unlocked_groups, freeze_bn_stats=freeze_bn_stats)
+
+    def lock_text_tower(self, unlocked_layers:int=0, freeze_layer_norm:bool=True):
+        self.text.lock(unlocked_layers, freeze_layer_norm)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        if self.text is not None:
+            self.text.set_grad_checkpointing(enable)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'logit_scale'}
+
+    def encode_image(self, image, normalize: bool = False):
+        features = self.visual(image)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        features = self.text(text)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def forward(self, image, text):
+        if self.visual is not None:
+            image_features = self.encode_image(image, normalize=True)
+        else:
+            image_features = None
+        if self.text is not None:
+            text_features = self.encode_text(text, normalize=True)
+        else:
+            text_features = None
+        return image_features, text_features, self.logit_scale.exp()
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr, None)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, nn.Parameter):
+            l.data = l.data.to(dtype)
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name) and isinstance(l, nn.Parameter):
+                attr = getattr(l, name, None)
+                if attr is not None:
+                    attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+                'logit_scale'
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def trace_model(model, batch_size=256, device=torch.device('cpu')):
+    model.eval()
+    image_size = model.visual.image_size
+    example_images = torch.ones((batch_size, 3, image_size, image_size), device=device)
+    example_text = torch.zeros((batch_size, model.context_length), dtype=torch.int, device=device)
+    model = torch.jit.trace_module(
+        model,
+        inputs=dict(
+            forward=(example_images, example_text),
+            encode_text=(example_text,),
+            encode_image=(example_images,)
+        ))
+    model.visual.image_size = image_size
+    return model

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA01-CLIP-B-16.json

@@ -0,0 +1,19 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "patch_size": 16,
+        "eva_model_name": "eva-clip-b-16",
+        "ls_init_value": 0.1,
+        "drop_path_rate": 0.0
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA01-CLIP-g-14-plus.json

@@ -0,0 +1,24 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 40,
+        "width": 1408,
+        "head_width": 88,
+        "mlp_ratio": 4.3637,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-g-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA01-CLIP-g-14.json

@@ -0,0 +1,24 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 40,
+        "width": 1408,
+        "head_width": 88,
+        "mlp_ratio": 4.3637,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-g-14-x",
+        "drop_path_rate": 0.4,
+        "xattn": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA02-CLIP-B-16.json

@@ -0,0 +1,30 @@
+{
+    "embed_dim": 512,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 12,
+        "width": 768,
+        "head_width": 64,
+        "patch_size": 16,
+        "mlp_ratio": 2.6667,
+        "eva_model_name": "eva-clip-b-16-X",
+        "drop_path_rate": 0.0,
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true,
+        "patch_dropout": 0.5
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 512,
+        "heads": 8,
+        "layers": 12,
+        "xattn": true,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA02-CLIP-L-14-336.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 336,
+        "layers": 24,
+        "width": 1024,
+        "drop_path_rate": 0,
+        "head_width": 64,
+        "mlp_ratio": 2.6667,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-l-14-336",
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA02-CLIP-L-14.json

@@ -0,0 +1,29 @@
+{
+    "embed_dim": 768,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 24,
+        "width": 1024,
+        "drop_path_rate": 0,
+        "head_width": 64,
+        "mlp_ratio": 2.6667,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-l-14",
+        "xattn": true,
+        "fusedLN": true,
+        "rope": true,
+        "pt_hw_seq_len": 16,
+        "intp_freq": true,
+        "naiveswiglu": true,
+        "subln": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 768,
+        "heads": 12,
+        "layers": 12,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA02-CLIP-bigE-14-plus.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 64,
+        "width": 1792,
+        "head_width": 112,
+        "mlp_ratio": 8.571428571428571,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-4b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/model_configs/EVA02-CLIP-bigE-14.json

@@ -0,0 +1,25 @@
+{
+    "embed_dim": 1024,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 64,
+        "width": 1792,
+        "head_width": 112,
+        "mlp_ratio": 8.571428571428571,
+        "patch_size": 14,
+        "eva_model_name": "eva-clip-4b-14-x",
+        "drop_path_rate": 0,
+        "xattn": true,
+        "postnorm": true,
+        "fusedLN": true
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1024,
+        "heads": 16,
+        "layers": 24,
+        "xattn": false,
+        "fusedLN": true
+    }
+}

code/C3/visual_bge/visual_bge/eva_clip/modified_resnet.py

@@ -0,0 +1,181 @@
+from collections import OrderedDict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .utils import freeze_batch_norm_2d
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.act1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.act2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.act3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.act1(self.bn1(self.conv1(x)))
+        out = self.act2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.act3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0.,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, image_size=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.image_size = image_size
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.act1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.act2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.act3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(image_size // 32, embed_dim, heads, output_dim)
+
+        self.init_parameters()
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def init_parameters(self):
+        if self.attnpool is not None:
+            std = self.attnpool.c_proj.in_features ** -0.5
+            nn.init.normal_(self.attnpool.q_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.k_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.v_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.c_proj.weight, std=std)
+
+        for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
+            for name, param in resnet_block.named_parameters():
+                if name.endswith("bn3.weight"):
+                    nn.init.zeros_(param)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        for param in self.parameters():
+            param.requires_grad = False
+        if freeze_bn_stats:
+            freeze_batch_norm_2d(self)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        # FIXME support for non-transformer
+        pass
+
+    def stem(self, x):
+        x = self.act1(self.bn1(self.conv1(x)))
+        x = self.act2(self.bn2(self.conv2(x)))
+        x = self.act3(self.bn3(self.conv3(x)))
+        x = self.avgpool(x)
+        return x
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x

code/C3/visual_bge/visual_bge/eva_clip/openai.py

@@ -0,0 +1,144 @@
+""" OpenAI pretrained model functions
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+
+import os
+import warnings
+from typing import List, Optional, Union
+
+import torch
+
+from .model import build_model_from_openai_state_dict, convert_weights_to_lp, get_cast_dtype
+from .pretrained import get_pretrained_url, list_pretrained_models_by_tag, download_pretrained_from_url
+
+__all__ = ["list_openai_models", "load_openai_model"]
+
+
+def list_openai_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list_pretrained_models_by_tag('openai')
+
+
+def load_openai_model(
+        name: str,
+        precision: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        jit: bool = True,
+        cache_dir: Optional[str] = None,
+):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+    precision: str
+        Model precision, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
+    device : Union[str, torch.device]
+        The device to put the loaded model
+    jit : bool
+        Whether to load the optimized JIT model (default) or more hackable non-JIT model.
+    cache_dir : Optional[str]
+        The directory to cache the downloaded model weights
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if precision is None:
+        precision = 'fp32' if device == 'cpu' else 'fp16'
+
+    if get_pretrained_url(name, 'openai'):
+        model_path = download_pretrained_from_url(get_pretrained_url(name, 'openai'), cache_dir=cache_dir)
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
+
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location=device if jit else "cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        if jit:
+            warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+            jit = False
+        state_dict = torch.load(model_path, map_location="cpu")
+
+    if not jit:
+        # Build a non-jit model from the OpenAI jitted model state dict
+        cast_dtype = get_cast_dtype(precision)
+        try:
+            model = build_model_from_openai_state_dict(state_dict or model.state_dict(), cast_dtype=cast_dtype)
+        except KeyError:
+            sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
+            model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype)
+
+        # model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
+        model = model.to(device)
+        if precision.startswith('amp') or precision == 'fp32':
+            model.float()
+        elif precision == 'bf16':
+            convert_weights_to_lp(model, dtype=torch.bfloat16)
+
+        return model
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 (typically for CPU)
+    if precision == 'fp32':
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+        model.float()
+
+    # ensure image_size attr available at consistent location for both jit and non-jit
+    model.visual.image_size = model.input_resolution.item()
+    return model

code/C3/visual_bge/visual_bge/eva_clip/pretrained.py

@@ -0,0 +1,332 @@
+import hashlib
+import os
+import urllib
+import warnings
+from functools import partial
+from typing import Dict, Union
+
+from tqdm import tqdm
+
+try:
+    from huggingface_hub import hf_hub_download
+    _has_hf_hub = True
+except ImportError:
+    hf_hub_download = None
+    _has_hf_hub = False
+
+
+def _pcfg(url='', hf_hub='', filename='', mean=None, std=None):
+    return dict(
+        url=url,
+        hf_hub=hf_hub,
+        mean=mean,
+        std=std,
+    )
+
+_VITB32 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+    laion2b_e16=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-laion2b_e16-af8dbd0c.pth"),
+    laion2b_s34b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-laion2B-s34B-b79K/')
+)
+
+_VITB32_quickgelu = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+)
+
+_VITB16 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e31-00efa78f.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e32-55e67d44.pt"),
+    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-laion2B-s34B-b88K/'),
+)
+
+_EVAB16 = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
+    eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_B_psz14to16.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
+    eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt'),
+)
+
+_VITB16_PLUS_240 = dict(
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e31-8fb26589.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"),
+)
+
+_VITL14 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e31-69988bb6.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e32-3d133497.pt"),
+    laion2b_s32b_b82k=_pcfg(
+        hf_hub='laion/CLIP-ViT-L-14-laion2B-s32B-b82K/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+)
+
+_EVAL14 = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
+    eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_L_psz14.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
+    eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt'),
+)
+
+_VITL14_336 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt"),
+)
+
+_EVAL14_336 = dict(
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
+    eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt'),
+    eva_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
+    eva02_clip_224to336=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_224to336.pt'),
+)
+
+_VITH14 = dict(
+    laion2b_s32b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-laion2B-s32B-b79K/'),
+)
+
+_VITg14 = dict(
+    laion2b_s12b_b42k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s12B-b42K/'),
+    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s34B-b88K/'),
+)
+
+_EVAg14 = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
+    eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
+    eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt'),
+)
+
+_EVAg14_PLUS = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/'),
+    eva01=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_g_psz14.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
+    eva01_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt'),
+)
+
+_VITbigG14 = dict(
+    laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
+)
+
+_EVAbigE14 = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+    eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
+    eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt'),
+)
+
+_EVAbigE14_PLUS = dict(
+    eva=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+    eva02=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_E_psz14.pt'),
+    eva_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
+    eva02_clip=_pcfg(hf_hub='QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt'),
+)
+
+
+_PRETRAINED = {
+    # "ViT-B-32": _VITB32,
+    "OpenaiCLIP-B-32": _VITB32,
+    "OpenCLIP-B-32": _VITB32,
+
+    # "ViT-B-32-quickgelu": _VITB32_quickgelu,
+    "OpenaiCLIP-B-32-quickgelu": _VITB32_quickgelu,
+    "OpenCLIP-B-32-quickgelu": _VITB32_quickgelu,
+
+    # "ViT-B-16": _VITB16,
+    "OpenaiCLIP-B-16": _VITB16,
+    "OpenCLIP-B-16": _VITB16,
+
+    "EVA02-B-16": _EVAB16,
+    "EVA02-CLIP-B-16": _EVAB16,
+
+    # "ViT-B-16-plus-240": _VITB16_PLUS_240,
+    "OpenCLIP-B-16-plus-240": _VITB16_PLUS_240,
+
+    # "ViT-L-14": _VITL14,
+    "OpenaiCLIP-L-14": _VITL14,
+    "OpenCLIP-L-14": _VITL14,
+
+    "EVA02-L-14": _EVAL14,
+    "EVA02-CLIP-L-14": _EVAL14,
+
+    # "ViT-L-14-336": _VITL14_336,
+    "OpenaiCLIP-L-14-336": _VITL14_336,
+
+    "EVA02-CLIP-L-14-336": _EVAL14_336,
+
+    # "ViT-H-14": _VITH14,
+    # "ViT-g-14": _VITg14,
+    "OpenCLIP-H-14": _VITH14,
+    "OpenCLIP-g-14": _VITg14,
+
+    "EVA01-CLIP-g-14": _EVAg14,
+    "EVA01-CLIP-g-14-plus": _EVAg14_PLUS,
+
+    # "ViT-bigG-14": _VITbigG14,
+    "OpenCLIP-bigG-14": _VITbigG14,
+
+    "EVA02-CLIP-bigE-14": _EVAbigE14,
+    "EVA02-CLIP-bigE-14-plus": _EVAbigE14_PLUS,
+}
+
+
+def _clean_tag(tag: str):
+    # normalize pretrained tags
+    return tag.lower().replace('-', '_')
+
+
+def list_pretrained(as_str: bool = False):
+    """ returns list of pretrained models
+    Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
+    """
+    return [':'.join([k, t]) if as_str else (k, t) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
+
+
+def list_pretrained_models_by_tag(tag: str):
+    """ return all models having the specified pretrain tag """
+    models = []
+    tag = _clean_tag(tag)
+    for k in _PRETRAINED.keys():
+        if tag in _PRETRAINED[k]:
+            models.append(k)
+    return models
+
+
+def list_pretrained_tags_by_model(model: str):
+    """ return all pretrain tags for the specified model architecture """
+    tags = []
+    if model in _PRETRAINED:
+        tags.extend(_PRETRAINED[model].keys())
+    return tags
+
+
+def is_pretrained_cfg(model: str, tag: str):
+    if model not in _PRETRAINED:
+        return False
+    return _clean_tag(tag) in _PRETRAINED[model]
+
+
+def get_pretrained_cfg(model: str, tag: str):
+    if model not in _PRETRAINED:
+        return {}
+    model_pretrained = _PRETRAINED[model]
+    return model_pretrained.get(_clean_tag(tag), {})
+
+
+def get_pretrained_url(model: str, tag: str):
+    cfg = get_pretrained_cfg(model, _clean_tag(tag))
+    return cfg.get('url', '')
+
+
+def download_pretrained_from_url(
+        url: str,
+        cache_dir: Union[str, None] = None,
+):
+    if not cache_dir:
+        cache_dir = os.path.expanduser("~/.cache/clip")
+    os.makedirs(cache_dir, exist_ok=True)
+    filename = os.path.basename(url)
+
+    if 'openaipublic' in url:
+        expected_sha256 = url.split("/")[-2]
+    elif 'mlfoundations' in url:
+        expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
+    else:
+        expected_sha256 = ''
+
+    download_target = os.path.join(cache_dir, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if expected_sha256:
+            if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+                return download_target
+            else:
+                warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+        else:
+            return download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.headers.get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def has_hf_hub(necessary=False):
+    if not _has_hf_hub and necessary:
+        # if no HF Hub module installed, and it is necessary to continue, raise error
+        raise RuntimeError(
+            'Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`.')
+    return _has_hf_hub
+
+
+def download_pretrained_from_hf(
+        model_id: str,
+        filename: str = 'open_clip_pytorch_model.bin',
+        revision=None,
+        cache_dir: Union[str, None] = None,
+):
+    has_hf_hub(True)
+    cached_file = hf_hub_download(model_id, filename, revision=revision, cache_dir=cache_dir)
+    return cached_file
+
+
+def download_pretrained(
+        cfg: Dict,
+        force_hf_hub: bool = False,
+        cache_dir: Union[str, None] = None,
+):
+    target = ''
+    if not cfg:
+        return target
+
+    download_url = cfg.get('url', '')
+    download_hf_hub = cfg.get('hf_hub', '')
+    if download_hf_hub and force_hf_hub:
+        # use HF hub even if url exists
+        download_url = ''
+
+    if download_url:
+        target = download_pretrained_from_url(download_url, cache_dir=cache_dir)
+    elif download_hf_hub:
+        has_hf_hub(True)
+        # we assume the hf_hub entries in pretrained config combine model_id + filename in
+        # 'org/model_name/filename.pt' form. To specify just the model id w/o filename and
+        # use 'open_clip_pytorch_model.bin' default, there must be a trailing slash 'org/model_name/'.
+        model_id, filename = os.path.split(download_hf_hub)
+        if filename:
+            target = download_pretrained_from_hf(model_id, filename=filename, cache_dir=cache_dir)
+        else:
+            target = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+
+    return target

code/C3/visual_bge/visual_bge/eva_clip/rope.py

@@ -0,0 +1,137 @@
+from math import pi
+import torch
+from torch import nn
+from einops import rearrange, repeat
+import logging
+
+def broadcat(tensors, dim = -1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation'
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim = dim)
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r = 2)
+    x1, x2 = x.unbind(dim = -1)
+    x = torch.stack((-x2, x1), dim = -1)
+    return rearrange(x, '... d r -> ... (d r)')
+
+
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len,
+        ft_seq_len=None,
+        custom_freqs = None,
+        freqs_for = 'lang',
+        theta = 10000,
+        max_freq = 10,
+        num_freqs = 1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None: ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs_h = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_h = repeat(freqs_h, '... n -> ... (n r)', r = 2)
+
+        freqs_w = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_w = repeat(freqs_w, '... n -> ... (n r)', r = 2)
+
+        freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim = -1) 
+
+        self.register_buffer("freqs_cos", freqs.cos())
+        self.register_buffer("freqs_sin", freqs.sin())
+
+        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+
+    def forward(self, t, start_index = 0):
+        rot_dim = self.freqs_cos.shape[-1]
+        end_index = start_index + rot_dim
+        assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+        t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+        t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
+
+        return torch.cat((t_left, t, t_right), dim = -1)
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len,
+        ft_seq_len=None,
+        custom_freqs = None,
+        freqs_for = 'lang',
+        theta = 10000,
+        max_freq = 10,
+        num_freqs = 1,
+        patch_dropout = 0.
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None: ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum('..., f -> ... f', t, freqs)
+        freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim = -1)
+
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.patch_dropout = patch_dropout
+
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+
+    def forward(self, t, patch_indices_keep=None):
+        if patch_indices_keep is not None:
+            batch = t.size()[0]
+            batch_indices = torch.arange(batch)
+            batch_indices = batch_indices[..., None]
+
+            freqs_cos = repeat(self.freqs_cos, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
+            freqs_sin = repeat(self.freqs_sin, 'i j -> n i m j', n=t.shape[0], m=t.shape[1])
+
+            freqs_cos = freqs_cos[batch_indices, patch_indices_keep]
+            freqs_cos = rearrange(freqs_cos, 'n i m j -> n m i j')
+            freqs_sin = freqs_sin[batch_indices, patch_indices_keep]
+            freqs_sin = rearrange(freqs_sin, 'n i m j -> n m i j')
+
+            return  t * freqs_cos + rotate_half(t) * freqs_sin
+
+        return  t * self.freqs_cos + rotate_half(t) * self.freqs_sin

code/C3/visual_bge/visual_bge/eva_clip/timm_model.py

@@ -0,0 +1,122 @@
+""" timm model adapter
+
+Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model.
+"""
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+try:
+    import timm
+    from timm.models.layers import Mlp, to_2tuple
+    try:
+        # old timm imports < 0.8.1
+        from timm.models.layers.attention_pool2d import RotAttentionPool2d
+        from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
+    except ImportError:
+        # new timm imports >= 0.8.1
+        from timm.layers import RotAttentionPool2d
+        from timm.layers import AttentionPool2d as AbsAttentionPool2d
+except ImportError:
+    timm = None
+
+from .utils import freeze_batch_norm_2d
+
+
+class TimmModel(nn.Module):
+    """ timm model adapter
+    # FIXME this adapter is a work in progress, may change in ways that break weight compat
+    """
+
+    def __init__(
+            self,
+            model_name,
+            embed_dim,
+            image_size=224,
+            pool='avg',
+            proj='linear',
+            proj_bias=False,
+            drop=0.,
+            pretrained=False):
+        super().__init__()
+        if timm is None:
+            raise RuntimeError("Please `pip install timm` to use timm models.")
+
+        self.image_size = to_2tuple(image_size)
+        self.trunk = timm.create_model(model_name, pretrained=pretrained)
+        feat_size = self.trunk.default_cfg.get('pool_size', None)
+        feature_ndim = 1 if not feat_size else 2
+        if pool in ('abs_attn', 'rot_attn'):
+            assert feature_ndim == 2
+            # if attn pooling used, remove both classifier and default pool
+            self.trunk.reset_classifier(0, global_pool='')
+        else:
+            # reset global pool if pool config set, otherwise leave as network default
+            reset_kwargs = dict(global_pool=pool) if pool else {}
+            self.trunk.reset_classifier(0, **reset_kwargs)
+        prev_chs = self.trunk.num_features
+
+        head_layers = OrderedDict()
+        if pool == 'abs_attn':
+            head_layers['pool'] = AbsAttentionPool2d(prev_chs, feat_size=feat_size, out_features=embed_dim)
+            prev_chs = embed_dim
+        elif pool == 'rot_attn':
+            head_layers['pool'] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
+            prev_chs = embed_dim
+        else:
+            assert proj, 'projection layer needed if non-attention pooling is used.'
+
+        # NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
+        if proj == 'linear':
+            head_layers['drop'] = nn.Dropout(drop)
+            head_layers['proj'] = nn.Linear(prev_chs, embed_dim, bias=proj_bias)
+        elif proj == 'mlp':
+            head_layers['mlp'] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=drop, bias=(True, proj_bias))
+
+        self.head = nn.Sequential(head_layers)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        """ lock modules
+        Args:
+            unlocked_groups (int): leave last n layer groups unlocked (default: 0)
+        """
+        if not unlocked_groups:
+            # lock full model
+            for param in self.trunk.parameters():
+                param.requires_grad = False
+            if freeze_bn_stats:
+                freeze_batch_norm_2d(self.trunk)
+        else:
+            # NOTE: partial freeze requires latest timm (master) branch and is subject to change
+            try:
+                # FIXME import here until API stable and in an official release
+                from timm.models.helpers import group_parameters, group_modules
+            except ImportError:
+                raise RuntimeError(
+                    'Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
+            matcher = self.trunk.group_matcher()
+            gparams = group_parameters(self.trunk, matcher)
+            max_layer_id = max(gparams.keys())
+            max_layer_id = max_layer_id - unlocked_groups
+            for group_idx in range(max_layer_id + 1):
+                group = gparams[group_idx]
+                for param in group:
+                    self.trunk.get_parameter(param).requires_grad = False
+            if freeze_bn_stats:
+                gmodules = group_modules(self.trunk, matcher, reverse=True)
+                gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
+                freeze_batch_norm_2d(self.trunk, gmodules)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        try:
+            self.trunk.set_grad_checkpointing(enable)
+        except Exception as e:
+            logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
+
+    def forward(self, x):
+        x = self.trunk(x)
+        x = self.head(x)
+        return x

code/C3/visual_bge/visual_bge/eva_clip/tokenizer.py

@@ -0,0 +1,201 @@
+""" CLIP tokenizer
+
+Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import gzip
+import html
+import os
+from functools import lru_cache
+from typing import Union, List
+
+import ftfy
+import regex as re
+import torch
+
+# https://stackoverflow.com/q/62691279
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        if not special_tokens:
+            special_tokens = ['<start_of_text>', '<end_of_text>']
+        else:
+            special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
+        vocab.extend(special_tokens)
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {t:t for t in special_tokens}
+        special = "|".join(special_tokens)
+        self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+        self.vocab_size = len(self.encoder)
+        self.all_special_ids = [self.encoder[t] for t in special_tokens]
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text
+
+
+_tokenizer = SimpleTokenizer()
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<start_of_text>"]
+    eot_token = _tokenizer.encoder["<end_of_text>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            tokens = tokens[:context_length]  # Truncate
+            tokens[-1] = eot_token
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+class HFTokenizer:
+    "HuggingFace tokenizer wrapper"
+    def __init__(self, tokenizer_name:str):
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+    def __call__(self, texts:Union[str, List[str]], context_length:int=77) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+        texts = [whitespace_clean(basic_clean(text)) for text in texts]
+        input_ids = self.tokenizer(texts, return_tensors='pt', max_length=context_length, padding='max_length', truncation=True).input_ids
+        return input_ids

code/C3/visual_bge/visual_bge/eva_clip/transform.py

@@ -0,0 +1,103 @@
+from typing import Optional, Sequence, Tuple
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+
+from torchvision.transforms import Normalize, Compose, RandomResizedCrop, InterpolationMode, ToTensor, Resize, \
+    CenterCrop
+
+from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+
+
+class ResizeMaxSize(nn.Module):
+
+    def __init__(self, max_size, interpolation=InterpolationMode.BICUBIC, fn='max', fill=0):
+        super().__init__()
+        if not isinstance(max_size, int):
+            raise TypeError(f"Size should be int. Got {type(max_size)}")
+        self.max_size = max_size
+        self.interpolation = interpolation
+        self.fn = min if fn == 'min' else min
+        self.fill = fill
+
+    def forward(self, img):
+        if isinstance(img, torch.Tensor):
+            height, width = img.shape[:2]
+        else:
+            width, height = img.size
+        scale = self.max_size / float(max(height, width))
+        if scale != 1.0:
+            new_size = tuple(round(dim * scale) for dim in (height, width))
+            img = F.resize(img, new_size, self.interpolation)
+            pad_h = self.max_size - new_size[0]
+            pad_w = self.max_size - new_size[1]
+            img = F.pad(img, padding=[pad_w//2, pad_h//2, pad_w - pad_w//2, pad_h - pad_h//2], fill=self.fill)
+        return img
+
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+
+# class CatGen(nn.Module):
+#     def __init__(self, num=4):
+#         self.num = num
+#     def mixgen_batch(image, text):
+#         batch_size = image.shape[0]
+#         index = np.random.permutation(batch_size)
+
+#         cat_images = []
+#         for i in range(batch_size):
+#             # image mixup
+#             image[i,:] = lam * image[i,:] + (1 - lam) * image[index[i],:]
+#             # text concat
+#             text[i] = tokenizer((str(text[i]) + " " + str(text[index[i]])))[0]
+#         text = torch.stack(text)
+#         return image, text
+
+
+def image_transform(
+        image_size: int,
+        is_train: bool,
+        mean: Optional[Tuple[float, ...]] = None,
+        std: Optional[Tuple[float, ...]] = None,
+        resize_longest_max: bool = False,
+        fill_color: int = 0,
+):
+    mean = mean or OPENAI_DATASET_MEAN
+    if not isinstance(mean, (list, tuple)):
+        mean = (mean,) * 3
+
+    std = std or OPENAI_DATASET_STD
+    if not isinstance(std, (list, tuple)):
+        std = (std,) * 3
+
+    if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
+        # for square size, pass size as int so that Resize() uses aspect preserving shortest edge
+        image_size = image_size[0]
+
+    normalize = Normalize(mean=mean, std=std)
+    if is_train:
+        return Compose([
+            RandomResizedCrop(image_size, scale=(0.9, 1.0), interpolation=InterpolationMode.BICUBIC),
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+    else:
+        if resize_longest_max:
+            transforms = [
+                ResizeMaxSize(image_size, fill=fill_color)
+            ]
+        else:
+            transforms = [
+                Resize(image_size, interpolation=InterpolationMode.BICUBIC),
+                CenterCrop(image_size),
+            ]
+        transforms.extend([
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+        return Compose(transforms)

code/C3/visual_bge/visual_bge/eva_clip/transformer.py

@@ -0,0 +1,737 @@
+import os
+import logging
+from collections import OrderedDict
+import math
+from typing import Callable, Optional, Sequence
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+try:
+    from timm.models.layers import trunc_normal_
+except:
+    from timm.layers import trunc_normal_
+    
+from .rope import VisionRotaryEmbedding, VisionRotaryEmbeddingFast
+from .utils import to_2tuple
+
+if os.getenv('ENV_TYPE') == 'deepspeed':
+    try:
+        import deepspeed
+        from deepspeed.runtime.activation_checkpointing.checkpointing import checkpoint
+    except:
+        print("Please 'pip install deepspeed'")
+        deepspeed = None
+        from torch.utils.checkpoint import checkpoint
+else:
+    from torch.utils.checkpoint import checkpoint
+
+try:
+    import xformers.ops as xops
+except ImportError:
+    xops = None
+    # print("Please 'pip install xformers'")
+
+class LayerNormFp32(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16 (by casting to float32 and back)."""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        output = F.layer_norm(
+            x.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        )
+        return output.type_as(x)
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        return x.to(orig_type)
+
+class QuickGELU(nn.Module):
+    # NOTE This is slower than nn.GELU or nn.SiLU and uses more GPU memory
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token  # exclude CLS token
+        logging.info(f"os.getenv('RoPE')={os.getenv('RoPE')}")
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        if self.training and os.getenv('RoPE') == '1':
+            return x, patch_indices_keep
+
+        return x
+
+
+def _in_projection_packed(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    w: torch.Tensor,
+    b: Optional[torch.Tensor] = None,
+    ):
+    """
+    https://github.com/pytorch/pytorch/blob/db2a237763eb8693a20788be94f8c192e762baa8/torch/nn/functional.py#L4726
+    """
+    E = q.size(-1)
+    if k is v:
+        if q is k:
+            # self-attention
+            return F.linear(q, w, b).chunk(3, dim=-1)
+        else:
+            # encoder-decoder attention
+            w_q, w_kv = w.split([E, E * 2])
+            if b is None:
+                b_q = b_kv = None
+            else:
+                b_q, b_kv = b.split([E, E * 2])
+            return (F.linear(q, w_q, b_q),) + F.linear(k, w_kv, b_kv).chunk(2, dim=-1)
+    else:
+        w_q, w_k, w_v = w.chunk(3)
+        if b is None:
+            b_q = b_k = b_v = None
+        else:
+            b_q, b_k, b_v = b.chunk(3)
+        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)
+
+class Attention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=True,
+            scaled_cosine=False,
+            scale_heads=False,
+            logit_scale_max=math.log(1. / 0.01),
+            attn_drop=0.,
+            proj_drop=0.,
+            xattn=False,
+            rope=False
+    ):
+        super().__init__()
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.logit_scale_max = logit_scale_max
+
+        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+        else:
+            self.logit_scale = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn
+        self.xattn_drop = attn_drop
+        self.rope = rope
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        L, N, C = x.shape
+        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+        if self.xattn:
+            q = q.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
+            k = k.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
+            v = v.contiguous().view(L, N, self.num_heads, -1).transpose(0, 1)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale if self.logit_scale is None else None,
+                attn_bias=xops.LowerTriangularMask() if attn_mask is not None else None,
+                )
+        else:
+            q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+            k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+            v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+
+            if self.logit_scale is not None:
+                attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
+                logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+                attn = attn.view(N, self.num_heads, L, L) * logit_scale
+                attn = attn.view(-1, L, L)
+            else:
+                q = q * self.scale
+                attn = torch.bmm(q, k.transpose(-1, -2))
+
+            if attn_mask is not None:
+                if attn_mask.dtype == torch.bool:
+                    new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+                    new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                    attn_mask = new_attn_mask
+                attn += attn_mask
+
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = torch.bmm(attn, v)
+
+        if self.head_scale is not None:
+            x = x.view(N, self.num_heads, L, C) * self.head_scale
+            x = x.view(-1, L, C)
+        x = x.transpose(0, 1).reshape(L, N, C)
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+class CustomAttention(nn.Module):
+    def __init__(
+            self,
+            dim,
+            num_heads=8,
+            qkv_bias=True,
+            scaled_cosine=True,
+            scale_heads=False,
+            logit_scale_max=math.log(1. / 0.01),
+            attn_drop=0.,
+            proj_drop=0.,
+            xattn=False
+    ):
+        super().__init__()
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.logit_scale_max = logit_scale_max
+
+        # keeping in_proj in this form (instead of nn.Linear) to match weight scheme of original
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))))
+        else:
+            self.logit_scale = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+        self.xattn = xattn
+        self.xattn_drop = attn_drop
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        q, k, v = _in_projection_packed(query, key, value, self.in_proj_weight, self.in_proj_bias)
+        N_q, B_q, C_q = q.shape
+        N_k, B_k, C_k = k.shape
+        N_v, B_v, C_v = v.shape
+        if self.xattn:
+            # B, N, C -> B, N, num_heads, C
+            q = q.permute(1, 0, 2).reshape(B_q, N_q, self.num_heads, -1)
+            k = k.permute(1, 0, 2).reshape(B_k, N_k, self.num_heads, -1)
+            v = v.permute(1, 0, 2).reshape(B_v, N_v, self.num_heads, -1)
+
+            x = xops.memory_efficient_attention(
+                q, k, v,
+                p=self.xattn_drop,
+                scale=self.scale if self.logit_scale is None else None,
+                attn_bias=xops.LowerTriangularMask() if attn_mask is not None else None
+                )
+        else:
+            # B*H, L, C
+            q = q.contiguous().view(N_q, B_q * self.num_heads, -1).transpose(0, 1)
+            k = k.contiguous().view(N_k, B_k * self.num_heads, -1).transpose(0, 1)
+            v = v.contiguous().view(N_v, B_v * self.num_heads, -1).transpose(0, 1)
+
+            if self.logit_scale is not None:
+                # B*H, N_q, N_k
+                attn = torch.bmm(F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2))
+                logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+                attn = attn.view(B_q, self.num_heads, N_q, N_k) * logit_scale
+                attn = attn.view(-1, N_q, N_k)
+            else:
+                q = q * self.scale
+                attn = torch.bmm(q, k.transpose(-1, -2))
+
+            if attn_mask is not None:
+                if attn_mask.dtype == torch.bool:
+                    new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+                    new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                    attn_mask = new_attn_mask
+                attn += attn_mask
+
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+
+            x = torch.bmm(attn, v)
+            
+        if self.head_scale is not None:
+            x = x.view(B_q, self.num_heads, N_q, C_q) * self.head_scale
+            x = x.view(-1, N_q, C_q)
+        x = x.transpose(0, 1).reshape(N_q, B_q, C_q)
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+class CustomResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            scale_cosine_attn: bool = False,
+            scale_heads: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+            cross_attn: bool = False,
+            xattn: bool = False,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        self.ln_1_k = norm_layer(d_model) if cross_attn else self.ln_1
+        self.ln_1_v = norm_layer(d_model) if cross_attn else self.ln_1
+        self.attn = CustomAttention(
+            d_model, n_head,
+            qkv_bias=True,
+            attn_drop=0.,
+            proj_drop=0.,
+            scaled_cosine=scale_cosine_attn,
+            scale_heads=scale_heads,
+            xattn=xattn
+        )
+
+        self.ln_attn = norm_layer(d_model) if scale_attn else nn.Identity()
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ('ln', norm_layer(mlp_width) if scale_fc else nn.Identity()),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        q = q + self.ls_1(self.ln_attn(self.attn(self.ln_1(q), self.ln_1_k(k), self.ln_1_v(v), attn_mask=attn_mask)))
+        q = q + self.ls_2(self.mlp(self.ln_2(q)))
+        return q
+
+class CustomTransformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            scale_cosine_attn: bool = True,
+            scale_heads: bool = False,
+            scale_attn: bool = False,
+            scale_fc: bool = False,
+            cross_attn: bool = False,
+            xattn: bool = False,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = False
+        self.xattn = xattn
+
+        self.resblocks = nn.ModuleList([
+            CustomResidualAttentionBlock(
+                width,
+                heads,
+                mlp_ratio,
+                ls_init_value=ls_init_value,
+                act_layer=act_layer,
+                norm_layer=norm_layer,
+                scale_cosine_attn=scale_cosine_attn,
+                scale_heads=scale_heads,
+                scale_attn=scale_attn,
+                scale_fc=scale_fc,
+                cross_attn=cross_attn,
+                xattn=xattn)
+            for _ in range(layers)
+        ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].mlp.c_fc.weight.dtype 
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor = None, v: torch.Tensor = None, attn_mask: Optional[torch.Tensor] = None):
+        if k is None and v is None:
+            k = v = q
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                q = checkpoint(r, q, k, v, attn_mask)
+            else:
+                q = r(q, k, v, attn_mask=attn_mask)
+        return q
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+            self,
+            d_model: int,
+            n_head: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            xattn: bool = False,
+    ):
+        super().__init__()
+
+        self.ln_1 = norm_layer(d_model)
+        if xattn:
+            self.attn = Attention(d_model, n_head, xattn=True)
+        else:
+            self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ls_1 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+
+        self.ln_2 = norm_layer(d_model)
+        mlp_width = int(d_model * mlp_ratio)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, mlp_width)),
+            ("gelu", act_layer()),
+            ("c_proj", nn.Linear(mlp_width, d_model))
+        ]))
+
+        self.ls_2 = LayerScale(d_model, ls_init_value) if ls_init_value is not None else nn.Identity()
+        self.xattn = xattn
+
+    def attention(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        attn_mask = attn_mask.to(x.dtype) if attn_mask is not None else None
+        if self.xattn:
+            return self.attn(x, attn_mask=attn_mask)
+        return self.attn(x, x, x, need_weights=False, attn_mask=attn_mask)[0]
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        x = x + self.ls_1(self.attention(self.ln_1(x), attn_mask=attn_mask))
+        x = x + self.ls_2(self.mlp(self.ln_2(x)))
+        return x
+
+class Transformer(nn.Module):
+    def __init__(
+            self,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float = 4.0,
+            ls_init_value: float = None,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            xattn: bool = False,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = False
+
+        self.resblocks = nn.ModuleList([
+            ResidualAttentionBlock(
+                width, heads, mlp_ratio, ls_init_value=ls_init_value, act_layer=act_layer, norm_layer=norm_layer, xattn=xattn)
+            for _ in range(layers)
+        ])
+
+    def get_cast_dtype(self) -> torch.dtype:
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+    def forward(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None):
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                x = checkpoint(r, x, attn_mask)
+            else:
+                x = r(x, attn_mask=attn_mask)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    def __init__(
+            self,
+            image_size: int,
+            patch_size: int,
+            width: int,
+            layers: int,
+            heads: int,
+            mlp_ratio: float,
+            ls_init_value: float = None,
+            patch_dropout: float = 0.,
+            global_average_pool: bool = False,
+            output_dim: int = 512,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            xattn: bool = False,
+    ):
+        super().__init__()
+        self.image_size = to_2tuple(image_size)
+        self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (self.image_size[0] // self.patch_size[0], self.image_size[1] // self.patch_size[1])
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width))
+
+        # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
+        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+        self.ln_pre = norm_layer(width)
+        
+        self.transformer = Transformer(
+            width,
+            layers,
+            heads,
+            mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            xattn=xattn
+        )
+
+        self.global_average_pool = global_average_pool
+        self.ln_post = norm_layer(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        for param in self.parameters():
+            param.requires_grad = False
+        
+        if unlocked_groups != 0:
+            groups = [
+                [
+                    self.conv1,
+                    self.class_embedding,
+                    self.positional_embedding,
+                    self.ln_pre,
+                ],
+                *self.transformer.resblocks[:-1],
+                [
+                    self.transformer.resblocks[-1],
+                    self.ln_post,
+                ],
+                self.proj,
+            ]
+
+            def _unlock(x):
+                if isinstance(x, Sequence):
+                    for g in x:
+                        _unlock(g)
+                else:
+                    if isinstance(x, torch.nn.Parameter):
+                        x.requires_grad = True
+                    else:
+                        for p in x.parameters():
+                            p.requires_grad = True
+
+            _unlock(groups[-unlocked_groups:])
+
+    def get_num_layers(self):
+        return self.transformer.layers
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'positional_embedding', 'class_embedding'}
+
+    def forward(self, x: torch.Tensor, return_all_features: bool=False):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat(
+            [self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device),
+             x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+
+        # a patch_dropout of 0. would mean it is disabled and this function would do nothing but return what was passed in
+        x = self.patch_dropout(x)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        if not return_all_features:
+            if self.global_average_pool:
+                x = x.mean(dim=1) #x = x[:,1:,:].mean(dim=1)
+            else:
+                x = x[:, 0]
+
+            x = self.ln_post(x)
+
+            if self.proj is not None:
+                x = x @ self.proj
+
+        return x
+
+
+class TextTransformer(nn.Module):
+    def __init__(
+            self,
+            context_length: int = 77,
+            vocab_size: int = 49408,
+            width: int = 512,
+            heads: int = 8,
+            layers: int = 12,
+            ls_init_value: float = None,
+            output_dim: int = 512,
+            act_layer: Callable = nn.GELU,
+            norm_layer: Callable = LayerNorm,
+            xattn: bool= False,
+            attn_mask: bool = True
+    ):
+        super().__init__()
+        self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.width = width
+        self.output_dim = output_dim
+
+        self.token_embedding = nn.Embedding(vocab_size, width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, width))
+        self.transformer = Transformer(
+            width=width,
+            layers=layers,
+            heads=heads,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+            xattn=xattn
+        )
+        
+        self.xattn = xattn
+        self.ln_final = norm_layer(width)
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+        if attn_mask:
+            self.register_buffer('attn_mask', self.build_attention_mask(), persistent=False)
+        else:
+            self.attn_mask = None
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+    
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        # return {'positional_embedding', 'token_embedding'}
+        return {'positional_embedding'}
+
+    def get_num_layers(self):
+        return self.transformer.layers
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    def forward(self, text, return_all_features: bool=False):
+        cast_dtype = self.transformer.get_cast_dtype()
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        # x = self.transformer(x) # no attention mask is applied
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)
+
+        if not return_all_features:
+            # x.shape = [batch_size, n_ctx, transformer.width]
+            # take features from the eot embedding (eot_token is the highest number in each sequence)
+            x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return x

code/C3/visual_bge/visual_bge/eva_clip/utils.py

@@ -0,0 +1,326 @@
+from itertools import repeat
+import collections.abc
+import logging
+import math
+import numpy as np
+
+import torch
+from torch import nn as nn
+from torchvision.ops.misc import FrozenBatchNorm2d
+import torch.nn.functional as F
+
+# open CLIP
+def resize_clip_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        align_corners=True,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed
+
+
+def resize_visual_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        align_corners=True,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['positional_embedding'] = new_pos_embed
+
+def resize_evaclip_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    all_keys = list(state_dict.keys())
+    # interpolate position embedding
+    if 'visual.pos_embed' in state_dict:
+        pos_embed_checkpoint = state_dict['visual.pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.visual.patch_embed.num_patches
+        num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            state_dict['visual.pos_embed'] = new_pos_embed
+
+            patch_embed_proj = state_dict['visual.patch_embed.proj.weight']
+            patch_size = model.visual.patch_embed.patch_size
+            state_dict['visual.patch_embed.proj.weight'] = torch.nn.functional.interpolate(
+                patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
+
+
+def resize_eva_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    all_keys = list(state_dict.keys())
+    # interpolate position embedding
+    if 'pos_embed' in state_dict:
+        pos_embed_checkpoint = state_dict['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.visual.patch_embed.num_patches
+        num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            state_dict['pos_embed'] = new_pos_embed
+
+            patch_embed_proj = state_dict['patch_embed.proj.weight']
+            patch_size = model.visual.patch_embed.patch_size
+            state_dict['patch_embed.proj.weight'] = torch.nn.functional.interpolate(
+                patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
+                
+
+def resize_rel_pos_embed(state_dict, model, interpolation: str = 'bicubic', seq_dim=1):
+    all_keys = list(state_dict.keys())
+    for key in all_keys:
+        if "relative_position_index" in key:
+            state_dict.pop(key)
+
+        if "relative_position_bias_table" in key:
+            rel_pos_bias = state_dict[key]
+            src_num_pos, num_attn_heads = rel_pos_bias.size()
+            dst_num_pos, _ = model.visual.state_dict()[key].size()
+            dst_patch_shape = model.visual.patch_embed.patch_shape
+            if dst_patch_shape[0] != dst_patch_shape[1]:
+                raise NotImplementedError()
+            num_extra_tokens = dst_num_pos - (dst_patch_shape[0] * 2 - 1) * (dst_patch_shape[1] * 2 - 1)
+            src_size = int((src_num_pos - num_extra_tokens) ** 0.5)
+            dst_size = int((dst_num_pos - num_extra_tokens) ** 0.5)
+            if src_size != dst_size:
+                print("Position interpolate for %s from %dx%d to %dx%d" % (
+                    key, src_size, src_size, dst_size, dst_size))
+                extra_tokens = rel_pos_bias[-num_extra_tokens:, :]
+                rel_pos_bias = rel_pos_bias[:-num_extra_tokens, :]
+
+                def geometric_progression(a, r, n):
+                    return a * (1.0 - r ** n) / (1.0 - r)
+
+                left, right = 1.01, 1.5
+                while right - left > 1e-6:
+                    q = (left + right) / 2.0
+                    gp = geometric_progression(1, q, src_size // 2)
+                    if gp > dst_size // 2:
+                        right = q
+                    else:
+                        left = q
+
+                # if q > 1.090307:
+                #     q = 1.090307
+
+                dis = []
+                cur = 1
+                for i in range(src_size // 2):
+                    dis.append(cur)
+                    cur += q ** (i + 1)
+
+                r_ids = [-_ for _ in reversed(dis)]
+
+                x = r_ids + [0] + dis
+                y = r_ids + [0] + dis
+
+                t = dst_size // 2.0
+                dx = np.arange(-t, t + 0.1, 1.0)
+                dy = np.arange(-t, t + 0.1, 1.0)
+
+                print("Original positions = %s" % str(x))
+                print("Target positions = %s" % str(dx))
+
+                all_rel_pos_bias = []
+
+                for i in range(num_attn_heads):
+                    z = rel_pos_bias[:, i].view(src_size, src_size).float().numpy()
+                    f = F.interpolate.interp2d(x, y, z, kind='cubic')
+                    all_rel_pos_bias.append(
+                        torch.Tensor(f(dx, dy)).contiguous().view(-1, 1).to(rel_pos_bias.device))
+
+                rel_pos_bias = torch.cat(all_rel_pos_bias, dim=-1)
+
+                new_rel_pos_bias = torch.cat((rel_pos_bias, extra_tokens), dim=0)
+                state_dict[key] = new_rel_pos_bias
+
+    # interpolate position embedding
+    if 'pos_embed' in state_dict:
+        pos_embed_checkpoint = state_dict['pos_embed']
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        num_patches = model.visual.patch_embed.num_patches
+        num_extra_tokens = model.visual.pos_embed.shape[-2] - num_patches
+        # height (== width) for the checkpoint position embedding
+        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        # height (== width) for the new position embedding
+        new_size = int(num_patches ** 0.5)
+        # class_token and dist_token are kept unchanged
+        if orig_size != new_size:
+            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
+            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
+            # only the position tokens are interpolated
+            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
+            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
+            pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
+            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
+            state_dict['pos_embed'] = new_pos_embed
+
+            patch_embed_proj = state_dict['patch_embed.proj.weight']
+            patch_size = model.visual.patch_embed.patch_size
+            state_dict['patch_embed.proj.weight'] = torch.nn.functional.interpolate(
+                patch_embed_proj.float(), size=patch_size, mode='bicubic', align_corners=False)
+
+
+def freeze_batch_norm_2d(module, module_match={}, name=''):
+    """
+    Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
+    itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
+    returned. Otherwise, the module is walked recursively and submodules are converted in place.
+
+    Args:
+        module (torch.nn.Module): Any PyTorch module.
+        module_match (dict): Dictionary of full module names to freeze (all if empty)
+        name (str): Full module name (prefix)
+
+    Returns:
+        torch.nn.Module: Resulting module
+
+    Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
+    """
+    res = module
+    is_match = True
+    if module_match:
+        is_match = name in module_match
+    if is_match and isinstance(module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)):
+        res = FrozenBatchNorm2d(module.num_features)
+        res.num_features = module.num_features
+        res.affine = module.affine
+        if module.affine:
+            res.weight.data = module.weight.data.clone().detach()
+            res.bias.data = module.bias.data.clone().detach()
+        res.running_mean.data = module.running_mean.data
+        res.running_var.data = module.running_var.data
+        res.eps = module.eps
+    else:
+        for child_name, child in module.named_children():
+            full_child_name = '.'.join([name, child_name]) if name else child_name
+            new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
+            if new_child is not child:
+                res.add_module(child_name, new_child)
+    return res
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = lambda n, x: _ntuple(n)(x)
+
+
+def is_logging(args):
+    def is_global_master(args):
+        return args.rank == 0
+
+    def is_local_master(args):
+        return args.local_rank == 0
+
+    def is_master(args, local=False):
+        return is_local_master(args) if local else is_global_master(args)
+    return is_master
+
+
+class AllGather(torch.autograd.Function):
+    """An autograd function that performs allgather on a tensor.
+    Performs all_gather operation on the provided tensors.
+    *** Warning ***: torch.distributed.all_gather has no gradient.
+    """
+
+    @staticmethod
+    def forward(ctx, tensor, rank, world_size):
+        tensors_gather = [torch.empty_like(tensor) for _ in range(world_size)]
+        torch.distributed.all_gather(tensors_gather, tensor)
+        ctx.rank = rank
+        ctx.batch_size = tensor.shape[0]
+        return torch.cat(tensors_gather, 0)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return (
+            grad_output[ctx.batch_size * ctx.rank: ctx.batch_size * (ctx.rank + 1)],
+            None,
+            None
+        )
+
+allgather = AllGather.apply

code/C3/visual_bge/visual_bge/modeling.py

@@ -0,0 +1,378 @@
+import os
+import logging
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import torch
+import torch.distributed as dist
+from torch import nn, Tensor
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+from transformers.file_utils import ModelOutput
+
+
+from .eva_clip import create_eva_vision_and_transforms
+from PIL import Image
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class EncoderOutput(ModelOutput):
+    q_reps: Optional[Tensor] = None
+    c_reps: Optional[Tensor] = None
+    loss: Optional[Tensor] = None
+    scores: Optional[Tensor] = None
+
+
+class Visualized_BGE(nn.Module):
+    def __init__(self,
+                 model_name_bge: str = None,
+                 model_weight = None, # "/path/to/your/weight/file/"
+                 normlized: bool = True,
+                 sentence_pooling_method: str = 'cls',
+                 negatives_cross_device: bool = False,
+                 temperature: float = 0.02, # 1.0
+                 from_pretrained=None, # local config file and model 
+                 ):
+        super().__init__()
+
+        assert 'bge' in model_name_bge
+        assert model_weight is not None
+        
+        self.model_name_bge = model_name_bge
+        
+        if 'bge-base-en-v1.5' in model_name_bge:
+            model_name_eva = "EVA02-CLIP-B-16"
+            self.hidden_dim = 768
+            self.depth = 12
+        elif 'bge-m3' in model_name_bge:
+            model_name_eva = "EVA02-CLIP-L-14"
+            self.hidden_dim = 1024
+            self.depth = 24
+        else:
+            raise Exception(f'Unavailable model_name {model_name_bge}')
+        
+        if not from_pretrained:
+            bge_config = AutoConfig.from_pretrained(model_name_bge)
+            bge = AutoModel.from_config(bge_config)
+        else:
+            print("Loading from local path.")
+            bge_config = AutoConfig.from_pretrained(from_pretrained, local_files_only=True)
+            bge = AutoModel.from_config(bge_config)
+        
+        self.bge_encoder = bge.encoder
+        self.bge_embeddings = bge.embeddings
+        self.bge_pooler = bge.pooler
+
+        self.model_visual, self.preprocess_train, self.preprocess_val= create_eva_vision_and_transforms(
+            model_name_eva, 
+            force_custom_clip=True)
+
+        
+        self.visual_proj = nn.Linear(self.hidden_dim, self.hidden_dim)
+
+        
+        self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
+
+        self.normlized = normlized
+        self.sentence_pooling_method = sentence_pooling_method
+        self.temperature = temperature
+        if not normlized:
+            self.temperature = 1.0
+            logger.info("reset temperature = 1.0 due to using inner product to compute similarity")
+
+        self.negatives_cross_device = negatives_cross_device
+        if self.negatives_cross_device:
+            if not dist.is_initialized():
+                raise ValueError('Distributed training has not been initialized for representation all gather.')
+
+            self.process_rank = dist.get_rank()
+            self.world_size = dist.get_world_size()
+        
+        self.load_model(model_weight)
+        
+        if not from_pretrained:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name_bge, use_fast=False)
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(from_pretrained, use_fast=False)
+
+        if torch.cuda.is_available():
+            self.device = torch.device('cuda')
+            self.to(self.device)
+        else:
+            self.device = torch.device('cpu')
+        self.dtype = next(bge.parameters()).dtype
+    
+    def load_model(self, model_weight):
+        self.load_state_dict(torch.load(model_weight, map_location='cpu'))
+    
+    def gradient_checkpointing_enable(self, **kwargs):
+        # self.bge_encoder.gradient_checkpointing_enable()
+        self.model_visual.set_grad_checkpointing(True)
+    
+    
+    
+    def encode(self, image=None, text=None):
+        # used for simple inference
+        if image is not None:
+            image = self.preprocess_val(Image.open(image)).unsqueeze(0)
+
+            if text is not None:
+                text = self.tokenizer(text, return_tensors="pt", padding=True)
+                return self.encode_mm(image.to(self.device), text.to(self.device))
+            else:
+                return self.encode_image(image.to(self.device))
+        else:
+            if text is not None:
+                text = self.tokenizer(text, return_tensors="pt", padding=True)
+                return self.encode_text(text.to(self.device))
+            else:
+                return None
+        
+    
+    def get_extended_attention_mask(
+        self, attention_mask: Tensor, input_shape: Tuple[int], device: torch.device = None, dtype: torch.float = torch.float16
+    ) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (`Tuple[int]`):
+                The shape of the input to the model.
+
+        Returns:
+            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
+        """
+        
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            
+            extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and the dtype's smallest value for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(dtype=dtype)  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(dtype).min
+        
+        return extended_attention_mask
+
+    def sentence_embedding(self, hidden_state, mask):
+        if self.sentence_pooling_method == 'mean':
+            s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
+            d = mask.sum(axis=1, keepdim=True).float()
+            return s / d
+        elif self.sentence_pooling_method == 'cls':
+            return hidden_state[:, 0]
+
+    
+    def encode_text(self, texts):
+        '''
+        encode text only
+        '''
+        input_ids = texts['input_ids']
+        attention_mask = texts['attention_mask']
+
+        input_shape = input_ids.size()
+        device = input_ids.device
+
+        token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        head_mask = [None] * self.depth
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape).to(self.dtype)
+        
+        embedding_output = self.bge_embeddings(
+            input_ids=input_ids,
+            position_ids=None,
+            token_type_ids=token_type_ids,
+            inputs_embeds=None,
+            past_key_values_length=0,
+        )
+        encoder_outputs = self.bge_encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            past_key_values=None,
+            use_cache=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+        sequence_output = encoder_outputs[0]
+        # pooled_output = self.bge_pooler(sequence_output) if self.bge_pooler is not None else None
+
+        t_reps = self.sentence_embedding(sequence_output, texts['attention_mask']) # tensor: reps with pooling
+        if self.normlized:
+            t_reps = torch.nn.functional.normalize(t_reps, dim=-1)
+        return t_reps.contiguous()
+
+    def encode_mm(self, images:torch.Tensor, texts):
+        img_token_emb = self.img_token_embedding(images) #[B, Patch_num, C]
+        img_token_emb = img_token_emb[:,1:]              # img_cls is not used here
+        img_token_emb = self.visual_proj(img_token_emb)
+        device = img_token_emb.device
+        
+        img_token_len = img_token_emb.size()[1]
+
+        # image position embedding, default position: bge_cls + img tokens + texts
+        img_token_position_ids = torch.arange(1, 1 + img_token_len).to(device=device)
+        img_position_embeddings = self.bge_embeddings.position_embeddings(img_token_position_ids)
+        img_token_emb = img_token_emb + img_position_embeddings
+
+        img_token_emb = self.bge_embeddings.LayerNorm(img_token_emb)
+
+        ### deal with prompt/text
+        prompt_input_ids = texts['input_ids']
+        prompt_attention_mask = texts['attention_mask']
+        prom_input_shape = prompt_input_ids.size()
+        
+        # bert
+        batch_size = prom_input_shape[0]
+        prompt_len = prom_input_shape[1]
+        prompt_start = 1 + img_token_len
+
+        
+        cls_id = torch.tensor([0]).to(device=device)
+        prompt_position_ids = torch.arange(prompt_start, prompt_start + prompt_len - 1).to(device=device)
+        prompt_position_ids = torch.cat([cls_id, prompt_position_ids]).to(device=device)
+
+        prompt_token_type_ids = torch.zeros(prom_input_shape, dtype=torch.long, device=device)
+        prompt_embedding_output = self.bge_embeddings(
+            input_ids=prompt_input_ids,
+            position_ids=prompt_position_ids,
+            token_type_ids=prompt_token_type_ids,
+            inputs_embeds=None,
+            past_key_values_length=0,
+        )  # [B, T, C]
+        
+        
+        cls_token = prompt_embedding_output[:, 0:1, :] # bge_cls token
+        prompt_embedding_output = prompt_embedding_output[:, 1:]
+
+        prompt_img_embedding = torch.cat([cls_token, img_token_emb, prompt_embedding_output], dim=1)
+        
+        img_attention_mask = torch.ones(batch_size, img_token_len, device=device)  
+        prom_img_attention_mask = torch.cat([img_attention_mask, prompt_attention_mask], dim=1)
+        prom_img_input_shape = prompt_img_embedding.size()
+
+        head_mask = [None] * self.depth
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(prom_img_attention_mask, prom_img_input_shape).to(self.dtype)
+        
+        
+        encoder_outputs = self.bge_encoder(
+            prompt_img_embedding,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=None,
+            encoder_attention_mask=None,
+            past_key_values=None,
+            use_cache=False,
+            output_attentions=False,
+            output_hidden_states=False,
+            return_dict=True,
+        )
+        sequence_output = encoder_outputs[0]
+        
+        prompt_img_reps = self.sentence_embedding(sequence_output, prom_img_attention_mask) # tensor: reps with pooling
+        if self.normlized:
+            prompt_img_reps = torch.nn.functional.normalize(prompt_img_reps, dim=-1)
+        return prompt_img_reps
+
+    def compute_similarity(self, q_reps, p_reps):
+        if len(p_reps.size()) == 2:
+            return torch.matmul(q_reps, p_reps.transpose(0, 1))
+        return torch.matmul(q_reps, p_reps.transpose(-2, -1))
+
+    def img_token_embedding(self, images):
+        if images is None:
+            return None
+        img_token_emb = self.model_visual.encode_image(images, normalize=False) # return_all_features=True, [B, Patch_num, C] 
+        
+        return img_token_emb.contiguous()
+    
+    def encode_image(self, images):
+        if images is None:
+            return None
+        
+        batch_size = images.shape[0]
+        prompts = [""] * batch_size
+        
+        prompts = self.tokenizer(prompts, return_tensors="pt", padding=True)        
+        prompts = prompts.to(images.device)
+        img_reps = self.encode_mm(images, prompts)
+        return img_reps
+    
+    def forward(self, mm_it_query=None, image_candidate=None, text_candidate=None, text_query=None, mm_it_candidate=None, task_type=None):
+        ### for stage-2 training
+        if task_type == "edit_image":
+            mm_query_reps = self.encode_mm(mm_it_query[0], mm_it_query[1])
+            image_candi_reps = self.encode_image(image_candidate)
+            query_reps = mm_query_reps
+            candi_reps = image_candi_reps
+                
+        elif task_type == "t2it":
+            text_query_reps = self.encode_text(text_query)
+            mmit_candi_reps = self.encode_mm(mm_it_candidate[0], mm_it_candidate[1])
+            query_reps = text_query_reps
+            candi_reps = mmit_candi_reps
+            
+        
+        if self.training:
+            if self.negatives_cross_device:
+                query_reps = self._dist_gather_tensor(query_reps)
+                candi_reps = self._dist_gather_tensor(candi_reps)
+
+            scores = self.compute_similarity(query_reps, candi_reps)
+            scores = scores / self.temperature
+            scores = scores.view(query_reps.size(0), -1)
+            
+            target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
+            target = target * (candi_reps.size(0) // query_reps.size(0))
+            
+            loss_edit = self.compute_loss(scores, target)
+            loss = loss_edit
+            
+            logging.info("task types: %s; loss: %s" %(task_type, str(loss_edit)))
+        else:
+            scores = self.compute_similarity(query_reps, candi_reps)
+            loss=None
+        return EncoderOutput(
+            loss=loss,
+            scores=scores,
+            q_reps=query_reps,
+            c_reps=candi_reps,
+        )
+
+    def compute_loss(self, scores, target):
+        return self.cross_entropy(scores, target)
+
+    def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
+        if t is None:
+            return None
+        t = t.contiguous()
+
+        all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
+        dist.all_gather(all_tensors, t)
+
+        all_tensors[self.process_rank] = t
+        all_tensors = torch.cat(all_tensors, dim=0)
+
+        return all_tensors
+
+    def save(self, output_dir: str):
+        torch.save(self.state_dict(), os.path.join(output_dir, 'Visualized_BGE.pth'))

code/C3/work_hybrid_multimodal_search.py

@@ -0,0 +1,403 @@
+import json
+import os
+from tqdm import tqdm
+import torch
+from visual_bge.visual_bge.modeling import Visualized_BGE
+from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
+from pymilvus.model.hybrid import BGEM3EmbeddingFunction
+import numpy as np
+import cv2
+from PIL import Image
+from typing import List, Dict, Any
+from dataclasses import dataclass
+
+@dataclass
+class DragonImage:
+    """龙类图像数据类"""
+    img_id: str
+    path: str
+    title: str
+    description: str
+    category: str
+    location: str
+    environment: str
+    combat_details: Dict[str, Any] = None
+    scene_info: Dict[str, Any] = None
+
+class DragonDataset:
+    """龙类图像数据集管理类"""
+    def __init__(self, data_dir: str, metadata_path: str):
+        self.data_dir = data_dir
+        self.metadata_path = metadata_path
+        self.images: List[DragonImage] = []
+        self._load_metadata()
+    
+    def _load_metadata(self):
+        """加载图像元数据"""
+        with open(self.metadata_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+            for img_data in data:
+                # 确保图片路径是完整的
+                if not img_data['path'].startswith(self.data_dir):
+                    img_data['path'] = os.path.join(self.data_dir, img_data['path'].split('/')[-1])
+                self.images.append(DragonImage(**img_data))
+    
+    def get_text_content(self, img: DragonImage) -> str:
+        """获取图像的文本描述内容"""
+        parts = [
+            img.title, img.description,
+            img.location, img.environment
+        ]
+        if img.combat_details:
+            parts.extend(img.combat_details.get('combat_style', []))
+            parts.extend(img.combat_details.get('abilities_used', []))
+        if img.scene_info:
+            parts.append(img.scene_info.get('time_of_day', ''))
+        return ' '.join(filter(None, parts))
+
+class HybridMultimodalEncoder:
+    """混合多模态编码器类"""
+    def __init__(self, visual_model_name: str, visual_model_path: str):
+        # 初始化Visual-BGE模型（用于多模态）
+        self.visual_model = Visualized_BGE(model_name_bge=visual_model_name, model_weight=visual_model_path)
+        self.visual_model.eval()
+        
+        # 初始化BGE-M3模型（用于混合检索）
+        self.bge_m3 = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
+        print(f"BGE-M3 密集向量维度: {self.bge_m3.dim['dense']}")
+
+    def encode_multimodal(self, image_path: str, text: str) -> list[float]:
+        """编码多模态内容（图像+文本）"""
+        with torch.no_grad():
+            query_emb = self.visual_model.encode(image=image_path, text=text)
+        return query_emb.tolist()[0]
+    
+    def encode_text_hybrid(self, text: str) -> dict:
+        """使用BGE-M3编码文本，返回稀疏和密集向量"""
+        embeddings = self.bge_m3([text])
+        return {
+            'sparse': embeddings["sparse"],
+            'dense': embeddings["dense"]
+        }
+    
+    def encode_query(self, image_path: str = None, text: str = None, mode: str = "multimodal") -> dict:
+        """编码查询，支持多种模式"""
+        result = {}
+        
+        if mode in ["multimodal", "all"] and image_path and text:
+            result['multimodal'] = self.encode_multimodal(image_path, text)
+        
+        if mode in ["hybrid", "dense", "sparse", "all"] and text:
+            text_embeddings = self.encode_text_hybrid(text)
+            result['dense'] = text_embeddings['dense'][0]
+            result['sparse'] = text_embeddings['sparse']._getrow(0)
+        
+        return result
+
+def visualize_results(query_image_path: str, retrieved_results: list, search_mode: str, 
+                     img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
+    """从检索到的结果创建一个全景图用于可视化"""
+    panoramic_width = img_width * row_count
+    panoramic_height = img_height * row_count
+    panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
+    query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
+
+    # 处理查询图像
+    if query_image_path and os.path.exists(query_image_path):
+        query_pil = Image.open(query_image_path).convert("RGB")
+        query_cv = np.array(query_pil)[:, :, ::-1]
+        resized_query = cv2.resize(query_cv, (img_width, img_height))
+        bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
+        query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
+        cv2.putText(query_display_area, "Query", (10, panoramic_height - 40), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2)
+        cv2.putText(query_display_area, search_mode, (10, panoramic_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 100, 0), 2)
+
+    # 处理检索到的图像
+    for i, result in enumerate(retrieved_results):
+        row, col = i // row_count, i % row_count
+        start_row, start_col = row * img_height, col * img_width
+        
+        img_path = result['image_path']
+        retrieved_pil = Image.open(img_path).convert("RGB")
+        retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
+        resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
+        bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
+        panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
+        
+        # 添加索引号和相似度
+        cv2.putText(panoramic_image, f"{i+1}", (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+        cv2.putText(panoramic_image, f"{result['distance']:.3f}", (start_col + 10, start_row + img_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+
+    return np.hstack([query_display_area, panoramic_image])
+
+class HybridMultimodalSearcher:
+    """混合多模态搜索系统"""
+    def __init__(self, data_dir: str, metadata_path: str, collection_name: str, milvus_uri: str):
+        self.data_dir = data_dir
+        self.metadata_path = metadata_path
+        self.collection_name = collection_name
+        self.milvus_uri = milvus_uri
+        
+        # 初始化数据集和编码器
+        print("--> 正在初始化数据集...")
+        self.dataset = DragonDataset(data_dir, metadata_path)
+        print(f"加载了 {len(self.dataset.images)} 张龙类图像")
+        
+        print("--> 正在初始化混合多模态编码器...")
+        self.encoder = HybridMultimodalEncoder(
+            visual_model_name="BAAI/bge-base-en-v1.5",
+            visual_model_path="../../models/bge/Visualized_base_en_v1.5.pth"
+        )
+        
+        # 连接Milvus
+        print(f"--> 正在连接到 Milvus: {milvus_uri}")
+        connections.connect(uri=milvus_uri)
+        self.milvus_client = MilvusClient(uri=milvus_uri)
+        
+        self.collection = None
+    
+    def create_collection(self):
+        """创建Collection"""
+        print(f"--> 正在创建 Collection '{self.collection_name}'")
+        if self.milvus_client.has_collection(self.collection_name):
+            self.milvus_client.drop_collection(self.collection_name)
+            print(f"已删除已存在的 Collection: '{self.collection_name}'")
+
+        # 获取向量维度
+        sample_text = self.dataset.get_text_content(self.dataset.images[0])
+        sample_path = self.dataset.images[0].path
+        multimodal_dim = len(self.encoder.encode_multimodal(sample_path, sample_text))
+        dense_dim = self.encoder.bge_m3.dim["dense"]
+        
+        print(f"多模态向量维度: {multimodal_dim}")
+        print(f"密集向量维度: {dense_dim}")
+
+        fields = [
+            FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
+            FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
+            FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
+            FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
+            FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
+            FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
+            FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
+            FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
+            # 三种向量类型
+            FieldSchema(name="multimodal_vector", dtype=DataType.FLOAT_VECTOR, dim=multimodal_dim),
+            FieldSchema(name="text_sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
+            FieldSchema(name="text_dense_vector", dtype=DataType.FLOAT_VECTOR, dim=dense_dim)
+        ]
+
+        schema = CollectionSchema(fields, description="混合多模态龙类图像检索")
+        self.collection = Collection(name=self.collection_name, schema=schema, consistency_level="Strong")
+        print("--> Collection 创建成功")
+
+        # 创建索引
+        print("--> 正在创建索引...")
+        # 多模态向量索引
+        multimodal_index = {"index_type": "HNSW", "metric_type": "COSINE", "params": {"M": 16, "efConstruction": 256}}
+        self.collection.create_index("multimodal_vector", multimodal_index)
+        print("多模态向量索引创建成功")
+        
+        # 稀疏向量索引
+        sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
+        self.collection.create_index("text_sparse_vector", sparse_index)
+        print("稀疏向量索引创建成功")
+        
+        # 密集向量索引
+        dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
+        self.collection.create_index("text_dense_vector", dense_index)
+        print("密集向量索引创建成功")
+        
+        self.collection.load()
+        print(f"--> Collection '{self.collection_name}' 已加载到内存")
+    
+    def insert_data(self):
+        """插入数据"""
+        if self.collection.is_empty:
+            print("--> Collection 为空，开始插入数据...")
+            
+            # 准备批量数据
+            img_ids, image_paths, titles, descriptions = [], [], [], []
+            categories, locations, environments = [], [], []
+            multimodal_vectors, text_sparse_vectors, text_dense_vectors = [], [], []
+            
+            for img_data in tqdm(self.dataset.images, desc="生成向量嵌入"):
+                text_content = self.dataset.get_text_content(img_data)
+                
+                # 生成多模态向量（图像+文本）
+                multimodal_vector = self.encoder.encode_multimodal(img_data.path, text_content)
+                
+                # 生成文本的混合向量（稀疏+密集）
+                text_embeddings = self.encoder.encode_text_hybrid(text_content)
+                
+                # 收集数据
+                img_ids.append(img_data.img_id)
+                image_paths.append(img_data.path)
+                titles.append(img_data.title)
+                descriptions.append(img_data.description)
+                categories.append(img_data.category)
+                locations.append(img_data.location)
+                environments.append(img_data.environment)
+                
+                multimodal_vectors.append(multimodal_vector)
+                text_sparse_vectors.append(text_embeddings['sparse']._getrow(0))
+                text_dense_vectors.append(text_embeddings['dense'][0])
+            
+            # 插入数据
+            self.collection.insert([
+                img_ids, image_paths, titles, descriptions, categories, locations, environments,
+                multimodal_vectors, text_sparse_vectors, text_dense_vectors
+            ])
+            
+            self.collection.flush()
+            print(f"--> 数据插入完成，总数: {self.collection.num_entities}")
+        else:
+            print(f"--> Collection 中已有 {self.collection.num_entities} 条数据，跳过插入")
+    
+    def search(self, query_image_path: str, query_text: str, mode: str = "hybrid", top_k: int = 5) -> list:
+        """执行搜索"""
+        search_params = {"metric_type": "IP", "params": {}}
+        cosine_params = {"metric_type": "COSINE", "params": {"ef": 128}}
+        output_fields = ["img_id", "image_path", "title", "description", "category", "location", "environment"]
+        
+        if mode == "multimodal":
+            # 多模态检索
+            query_vector = self.encoder.encode_multimodal(query_image_path, query_text)
+            results = self.collection.search(
+                [query_vector], "multimodal_vector", param=cosine_params, 
+                limit=top_k, output_fields=output_fields
+            )[0]
+            
+        elif mode == "dense":
+            # 密集向量检索
+            query_embeddings = self.encoder.encode_text_hybrid(query_text)
+            dense_vec = query_embeddings['dense'][0]
+            results = self.collection.search(
+                [dense_vec], "text_dense_vector", param=search_params,
+                limit=top_k, output_fields=output_fields
+            )[0]
+            
+        elif mode == "sparse":
+            # 稀疏向量检索
+            query_embeddings = self.encoder.encode_text_hybrid(query_text)
+            sparse_vec = query_embeddings['sparse']._getrow(0)
+            results = self.collection.search(
+                [sparse_vec], "text_sparse_vector", param=search_params,
+                limit=top_k, output_fields=output_fields
+            )[0]
+            
+        elif mode == "hybrid":
+            # 混合检索（稀疏+密集）
+            query_embeddings = self.encoder.encode_text_hybrid(query_text)
+            dense_vec = query_embeddings['dense'][0]
+            sparse_vec = query_embeddings['sparse']._getrow(0)
+            
+            # 创建RRF融合器
+            rerank = RRFRanker(k=60)
+            
+            # 创建搜索请求
+            dense_req = AnnSearchRequest([dense_vec], "text_dense_vector", search_params, limit=top_k)
+            sparse_req = AnnSearchRequest([sparse_vec], "text_sparse_vector", search_params, limit=top_k)
+            
+            # 执行混合搜索
+            results = self.collection.hybrid_search(
+                [sparse_req, dense_req], rerank=rerank, limit=top_k, output_fields=output_fields
+            )[0]
+        
+        return results
+    
+    def compare_search_modes(self, query_image_path: str, query_text: str, top_k: int = 5):
+        """对比不同搜索模式的效果"""
+        modes = ["multimodal", "dense", "sparse", "hybrid"]
+        results = {}
+        
+        print(f"\n{'='*50}")
+        print(f"查询图像: {query_image_path}")
+        print(f"查询文本: {query_text}")
+        print(f"{'='*50}")
+        
+        for mode in modes:
+            print(f"\n--- [{mode.upper()}] 搜索结果 ---")
+            search_results = self.search(query_image_path, query_text, mode, top_k)
+            
+            mode_results = []
+            for i, hit in enumerate(search_results):
+                print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+                print(f"    路径: {hit.entity.get('image_path')}")
+                print(f"    描述: {hit.entity.get('description')[:80]}...")
+                
+                mode_results.append({
+                    'image_path': hit.entity.get('image_path'),
+                    'distance': hit.distance,
+                    'title': hit.entity.get('title')
+                })
+            
+            results[mode] = mode_results
+        
+        return results
+    
+    def visualize_comparison(self, query_image_path: str, query_text: str, top_k: int = 3):
+        """可视化对比不同搜索模式"""
+        modes = ["multimodal", "dense", "sparse", "hybrid"]
+        
+        for mode in modes:
+            results = self.search(query_image_path, query_text, mode, top_k)
+            
+            retrieved_results = []
+            for hit in results:
+                retrieved_results.append({
+                    'image_path': hit.entity.get('image_path'),
+                    'distance': hit.distance
+                })
+            
+            if retrieved_results:
+                panoramic_image = visualize_results(query_image_path, retrieved_results, mode.upper())
+                output_path = f"../../data/C4/{mode}_search_result.png"
+                cv2.imwrite(output_path, panoramic_image)
+                print(f"{mode.upper()} 搜索结果已保存到: {output_path}")
+    
+    def cleanup(self):
+        """清理资源"""
+        if self.collection:
+            self.collection.release()
+            print(f"已从内存中释放 Collection: '{self.collection_name}'")
+            self.milvus_client.drop_collection(self.collection_name)
+            print(f"已删除 Collection: '{self.collection_name}'")
+
+# 主程序
+if __name__ == "__main__":
+    # 初始化设置
+    DATA_DIR = "../../data/C3/dragon"
+    METADATA_PATH = "../../data/C4/metadata/dragon.json"
+    COLLECTION_NAME = "hybrid_multimodal_dragon_demo"
+    MILVUS_URI = "http://localhost:19530"
+    
+    # 创建混合多模态搜索系统
+    searcher = HybridMultimodalSearcher(DATA_DIR, METADATA_PATH, COLLECTION_NAME, MILVUS_URI)
+    
+    try:
+        # 创建Collection并插入数据
+        searcher.create_collection()
+        searcher.insert_data()
+        
+        # 执行搜索对比
+        query_image_path = os.path.join(DATA_DIR, "query.png")
+        query_text = "悬崖上的巨龙"
+        
+        # 对比不同搜索模式
+        results = searcher.compare_search_modes(query_image_path, query_text, top_k=3)
+        
+        # 可视化结果
+        searcher.visualize_comparison(query_image_path, query_text, top_k=3)
+        
+        print(f"\n{'='*50}")
+        print("搜索模式分析:")
+        print("- MULTIMODAL: 结合图像和文本的多模态向量检索")
+        print("- DENSE: 基于语义的密集向量检索") 
+        print("- SPARSE: 基于关键词的稀疏向量检索")
+        print("- HYBRID: 稀疏+密集向量RRF融合检索")
+        print(f"{'='*50}")
+        
+    finally:
+        # 清理资源
+        searcher.cleanup() 

code/C3/work_multimodal_dragon_search.py

@@ -0,0 +1,301 @@
+import json
+import os
+from tqdm import tqdm
+from glob import glob
+import torch
+from visual_bge.visual_bge.modeling import Visualized_BGE
+from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
+import numpy as np
+import cv2
+from PIL import Image
+from typing import List, Dict, Any
+from dataclasses import dataclass
+
+@dataclass
+class DragonImage:
+    """龙类图像数据类"""
+    img_id: str
+    path: str
+    title: str
+    description: str
+    category: str
+    location: str
+    environment: str
+    combat_details: Dict[str, Any] = None
+    scene_info: Dict[str, Any] = None
+
+class DragonDataset:
+    """龙类图像数据集管理类"""
+    def __init__(self, data_dir: str, metadata_path: str):
+        self.data_dir = data_dir
+        self.metadata_path = metadata_path
+        self.images: List[DragonImage] = []
+        self._load_metadata()
+    
+    def _load_metadata(self):
+        """加载图像元数据"""
+        with open(self.metadata_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+            for img_data in data:
+                # 确保图片路径是完整的
+                if not img_data['path'].startswith(self.data_dir):
+                    img_data['path'] = os.path.join(self.data_dir, img_data['path'].split('/')[-1])
+                self.images.append(DragonImage(**img_data))
+    
+    def get_image_paths(self) -> List[str]:
+        """获取所有图像路径"""
+        return [img.path for img in self.images]
+    
+    def get_metadata_by_path(self, path: str) -> DragonImage:
+        """根据路径获取元数据"""
+        for img in self.images:
+            if img.path == path:
+                return img
+        return None
+    
+    def get_text_content(self, img: DragonImage) -> str:
+        """获取图像的文本描述内容"""
+        parts = [
+            img.title, img.description,
+            img.location, img.environment
+        ]
+        if img.combat_details:
+            parts.extend(img.combat_details.get('combat_style', []))
+            parts.extend(img.combat_details.get('abilities_used', []))
+        if img.scene_info:
+            parts.append(img.scene_info.get('time_of_day', ''))
+        return ' '.join(filter(None, parts))
+
+class Encoder:
+    """编码器类，用于将图像和文本编码为向量"""
+    def __init__(self, model_name: str, model_path: str):
+        self.model = Visualized_BGE(model_name_bge=model_name, model_weight=model_path)
+        self.model.eval()
+
+    def encode_query(self, image_path: str = None, text: str = None) -> list[float]:
+        """编码查询（支持图像+文本或仅文本）"""
+        with torch.no_grad():
+            if image_path and text:
+                query_emb = self.model.encode(image=image_path, text=text)
+            elif image_path:
+                query_emb = self.model.encode(image=image_path)
+            elif text:
+                query_emb = self.model.encode(text=text)
+            else:
+                raise ValueError("必须提供图像路径或文本内容")
+        return query_emb.tolist()[0]
+
+    def encode_multimodal(self, image_path: str, text: str) -> list[float]:
+        """编码多模态内容（图像+文本）"""
+        with torch.no_grad():
+            query_emb = self.model.encode(image=image_path, text=text)
+        return query_emb.tolist()[0]
+
+def visualize_results(query_image_path: str, retrieved_results: list, img_height: int = 300, img_width: int = 300, row_count: int = 3) -> np.ndarray:
+    """从检索到的结果创建一个全景图用于可视化"""
+    panoramic_width = img_width * row_count
+    panoramic_height = img_height * row_count
+    panoramic_image = np.full((panoramic_height, panoramic_width, 3), 255, dtype=np.uint8)
+    query_display_area = np.full((panoramic_height, img_width, 3), 255, dtype=np.uint8)
+
+    # 处理查询图像
+    if query_image_path and os.path.exists(query_image_path):
+        query_pil = Image.open(query_image_path).convert("RGB")
+        query_cv = np.array(query_pil)[:, :, ::-1]
+        resized_query = cv2.resize(query_cv, (img_width, img_height))
+        bordered_query = cv2.copyMakeBorder(resized_query, 10, 10, 10, 10, cv2.BORDER_CONSTANT, value=(255, 0, 0))
+        query_display_area[img_height * (row_count - 1):, :] = cv2.resize(bordered_query, (img_width, img_height))
+        cv2.putText(query_display_area, "Query", (10, panoramic_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+
+    # 处理检索到的图像
+    for i, result in enumerate(retrieved_results):
+        row, col = i // row_count, i % row_count
+        start_row, start_col = row * img_height, col * img_width
+        
+        img_path = result['image_path']
+        retrieved_pil = Image.open(img_path).convert("RGB")
+        retrieved_cv = np.array(retrieved_pil)[:, :, ::-1]
+        resized_retrieved = cv2.resize(retrieved_cv, (img_width - 4, img_height - 4))
+        bordered_retrieved = cv2.copyMakeBorder(resized_retrieved, 2, 2, 2, 2, cv2.BORDER_CONSTANT, value=(0, 0, 0))
+        panoramic_image[start_row:start_row + img_height, start_col:start_col + img_width] = bordered_retrieved
+        
+        # 添加索引号和相似度
+        cv2.putText(panoramic_image, f"{i+1}", (start_col + 10, start_row + 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+        cv2.putText(panoramic_image, f"{result['distance']:.3f}", (start_col + 10, start_row + img_height - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+
+    return np.hstack([query_display_area, panoramic_image])
+
+# 1. 初始化设置
+MODEL_NAME = "BAAI/bge-base-en-v1.5"
+MODEL_PATH = "../../models/bge/Visualized_base_en_v1.5.pth"
+DATA_DIR = "../../data/C3/dragon"
+METADATA_PATH = "../../data/C4/metadata/dragon.json"
+COLLECTION_NAME = "multimodal_dragon_demo"
+MILVUS_URI = "http://localhost:19530"
+
+# 2. 初始化数据集和编码器
+print("--> 正在初始化数据集...")
+dataset = DragonDataset(DATA_DIR, METADATA_PATH)
+print(f"加载了 {len(dataset.images)} 张龙类图像")
+
+print("--> 正在初始化编码器和Milvus客户端...")
+encoder = Encoder(MODEL_NAME, MODEL_PATH)
+milvus_client = MilvusClient(uri=MILVUS_URI)
+
+# 3. 创建 Milvus Collection
+print(f"\n--> 正在创建 Collection '{COLLECTION_NAME}'")
+if milvus_client.has_collection(COLLECTION_NAME):
+    milvus_client.drop_collection(COLLECTION_NAME)
+    print(f"已删除已存在的 Collection: '{COLLECTION_NAME}'")
+
+# 获取向量维度
+sample_text = dataset.get_text_content(dataset.images[0])
+sample_path = dataset.images[0].path
+dim = len(encoder.encode_multimodal(sample_path, sample_text))
+
+fields = [
+    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
+    FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dim),
+    FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
+    FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
+    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
+    FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
+    FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
+    FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
+    FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
+]
+
+schema = CollectionSchema(fields, description="多模态龙类图像检索")
+
+# 创建集合
+milvus_client.create_collection(collection_name=COLLECTION_NAME, schema=schema)
+print(f"成功创建 Collection: '{COLLECTION_NAME}'")
+
+# 4. 准备并插入数据
+print(f"\n--> 正在向 '{COLLECTION_NAME}' 插入数据")
+data_to_insert = []
+for img_data in tqdm(dataset.images, desc="生成多模态嵌入"):
+    # 结合图像和文本信息生成向量
+    text_content = dataset.get_text_content(img_data)
+    vector = encoder.encode_multimodal(img_data.path, text_content)
+    
+    data_to_insert.append({
+        "vector": vector,
+        "img_id": img_data.img_id,
+        "image_path": img_data.path,
+        "title": img_data.title,
+        "description": img_data.description,
+        "category": img_data.category,
+        "location": img_data.location,
+        "environment": img_data.environment
+    })
+
+if data_to_insert:
+    result = milvus_client.insert(collection_name=COLLECTION_NAME, data=data_to_insert)
+    print(f"成功插入 {result['insert_count']} 条数据。")
+
+# 5. 创建索引
+print(f"\n--> 正在为 '{COLLECTION_NAME}' 创建索引")
+index_params = milvus_client.prepare_index_params()
+index_params.add_index(
+    field_name="vector",
+    index_type="HNSW",
+    metric_type="COSINE",
+    params={"M": 16, "efConstruction": 256}
+)
+milvus_client.create_index(collection_name=COLLECTION_NAME, index_params=index_params)
+print("成功为向量字段创建 HNSW 索引。")
+milvus_client.load_collection(collection_name=COLLECTION_NAME)
+print("已加载 Collection 到内存中。")
+
+# 6. 执行多模态检索
+print(f"\n--> 正在 '{COLLECTION_NAME}' 中执行多模态检索")
+
+# 示例1：图像+文本查询
+query_image_path = os.path.join(DATA_DIR, "query.png")
+query_text = "悬崖上的巨龙"
+query_vector = encoder.encode_query(image_path=query_image_path, text=query_text)
+
+print(f"\n=== 多模态查询（图像+文本）===")
+print(f"查询图像: {query_image_path}")
+print(f"查询文本: {query_text}")
+
+search_results = milvus_client.search(
+    collection_name=COLLECTION_NAME,
+    data=[query_vector],
+    output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
+    limit=6,
+    search_params={"metric_type": "COSINE", "params": {"ef": 128}}
+)[0]
+
+retrieved_results = []
+print("检索结果:")
+for i, hit in enumerate(search_results):
+    print(f"  Top {i+1}: ID={hit['id']}, 距离={hit['distance']:.4f}")
+    print(f"    标题: {hit['entity']['title']}")
+    print(f"    描述: {hit['entity']['description'][:100]}...")
+    print(f"    类别: {hit['entity']['category']}")
+    print(f"    路径: {hit['entity']['image_path']}")
+    print("-" * 50)
+    retrieved_results.append({
+        'image_path': hit['entity']['image_path'],
+        'distance': hit['distance']
+    })
+
+# 示例2：纯文本查询
+print(f"\n=== 纯文本查询 ===")
+text_query = "悬崖上的巨龙"
+text_query_vector = encoder.encode_query(text=text_query)
+
+print(f"查询文本: {text_query}")
+
+text_search_results = milvus_client.search(
+    collection_name=COLLECTION_NAME,
+    data=[text_query_vector],
+    output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
+    limit=3,
+    search_params={"metric_type": "COSINE", "params": {"ef": 128}}
+)[0]
+
+print("文本检索结果:")
+for i, hit in enumerate(text_search_results):
+    print(f"  Top {i+1}: {hit['entity']['title']} (距离: {hit['distance']:.4f})")
+    print(f"    描述: {hit['entity']['description'][:80]}...")
+
+# 示例3：纯图像查询
+print(f"\n=== 纯图像查询 ===")
+image_query_path = os.path.join(DATA_DIR, "query.png")
+image_query_vector = encoder.encode_query(image_path=image_query_path)
+
+print(f"查询图像: {image_query_path}")
+
+image_search_results = milvus_client.search(
+    collection_name=COLLECTION_NAME,
+    data=[image_query_vector],
+    output_fields=["img_id", "image_path", "title", "description", "category", "location", "environment"],
+    limit=3,
+    search_params={"metric_type": "COSINE", "params": {"ef": 128}}
+)[0]
+
+print("图像检索结果:")
+for i, hit in enumerate(image_search_results):
+    print(f"  Top {i+1}: {hit['entity']['title']} (距离: {hit['distance']:.4f})")
+    print(f"    类别: {hit['entity']['category']}")
+    print(f"    描述: {hit['entity']['description'][:80]}...")
+    print(f"    路径: {hit['entity']['image_path']}")
+    print("-" * 30)
+
+# 7. 可视化与清理
+print(f"\n--> 正在可视化结果并清理资源")
+if retrieved_results:
+    panoramic_image = visualize_results(query_image_path, retrieved_results)
+    combined_image_path = "../../data/C4/multimodal_search.png"
+    cv2.imwrite(combined_image_path, panoramic_image)
+    print(f"结果图像已保存到: {combined_image_path}")
+    # Image.open(combined_image_path).show()
+
+# 8. 清理资源
+milvus_client.release_collection(collection_name=COLLECTION_NAME)
+print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
+milvus_client.drop_collection(COLLECTION_NAME)
+print(f"已删除 Collection: '{COLLECTION_NAME}'") 

code/C4/01_hybrid_search.py

@@ -0,0 +1,209 @@
+import json
+import os
+import numpy as np
+from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
+from pymilvus.model.hybrid import BGEM3EmbeddingFunction
+
+# 1. 初始化设置
+COLLECTION_NAME = "dragon_hybrid_demo"
+MILVUS_URI = "http://localhost:19530"  # 服务器模式
+DATA_PATH = "../../data/C4/metadata/dragon.json"  # 相对路径
+BATCH_SIZE = 50
+
+# 2. 连接 Milvus 并初始化嵌入模型
+print(f"--> 正在连接到 Milvus: {MILVUS_URI}")
+connections.connect(uri=MILVUS_URI)
+
+print("--> 正在初始化 BGE-M3 嵌入模型...")
+ef = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
+print(f"--> 嵌入模型初始化完成。密集向量维度: {ef.dim['dense']}")
+
+# 3. 创建 Collection
+milvus_client = MilvusClient(uri=MILVUS_URI)
+if milvus_client.has_collection(COLLECTION_NAME):
+    print(f"--> 正在删除已存在的 Collection '{COLLECTION_NAME}'...")
+    milvus_client.drop_collection(COLLECTION_NAME)
+
+fields = [
+    FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
+    FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
+    FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=256),
+    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
+    FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
+    FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
+    FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
+    FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
+    FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
+    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=ef.dim["dense"])
+]
+
+# 如果集合不存在，则创建它及索引
+if not milvus_client.has_collection(COLLECTION_NAME):
+    print(f"--> 正在创建 Collection '{COLLECTION_NAME}'...")
+    schema = CollectionSchema(fields, description="关于龙的混合检索示例")
+    # 创建集合
+    collection = Collection(name=COLLECTION_NAME, schema=schema, consistency_level="Strong")
+    print("--> Collection 创建成功。")
+
+    # 4. 创建索引
+    print("--> 正在为新集合创建索引...")
+    sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
+    collection.create_index("sparse_vector", sparse_index)
+    print("稀疏向量索引创建成功。")
+
+    dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
+    collection.create_index("dense_vector", dense_index)
+    print("密集向量索引创建成功。")
+
+collection = Collection(COLLECTION_NAME)
+
+# 5. 加载数据并插入
+collection.load()
+print(f"--> Collection '{COLLECTION_NAME}' 已加载到内存。")
+
+if collection.is_empty:
+    print(f"--> Collection 为空，开始插入数据...")
+    if not os.path.exists(DATA_PATH):
+        raise FileNotFoundError(f"数据文件未找到: {DATA_PATH}")
+    with open(DATA_PATH, 'r', encoding='utf-8') as f:
+        dataset = json.load(f)
+
+    docs, metadata = [], []
+    for item in dataset:
+        parts = [
+            item.get('title', ''),
+            item.get('description', ''),
+            item.get('location', ''),
+            item.get('environment', ''),
+            # *item.get('combat_details', {}).get('combat_style', []),
+            # *item.get('combat_details', {}).get('abilities_used', []),
+            # item.get('scene_info', {}).get('time_of_day', '')
+        ]
+        docs.append(' '.join(filter(None, parts)))
+        metadata.append(item)
+    print(f"--> 数据加载完成，共 {len(docs)} 条。")
+
+    print("--> 正在生成向量嵌入...")
+    embeddings = ef(docs)
+    print("--> 向量生成完成。")
+
+    print("--> 正在分批插入数据...")
+    # 为每个字段准备批量数据
+    img_ids = [doc["img_id"] for doc in metadata]
+    paths = [doc["path"] for doc in metadata]
+    titles = [doc["title"] for doc in metadata]
+    descriptions = [doc["description"] for doc in metadata]
+    categories = [doc["category"] for doc in metadata]
+    locations = [doc["location"] for doc in metadata]
+    environments = [doc["environment"] for doc in metadata]
+    
+    # 获取向量
+    sparse_vectors = embeddings["sparse"]
+    dense_vectors = embeddings["dense"]
+    
+    # 插入数据
+    collection.insert([
+        img_ids,
+        paths,
+        titles,
+        descriptions,
+        categories,
+        locations,
+        environments,
+        sparse_vectors,
+        dense_vectors
+    ])
+    
+    collection.flush()
+    print(f"--> 数据插入完成，总数: {collection.num_entities}")
+else:
+    print(f"--> Collection 中已有 {collection.num_entities} 条数据，跳过插入。")
+
+# 6. 执行搜索
+search_query = "悬崖上的巨龙"
+search_filter = 'category in ["western_dragon", "chinese_dragon", "movie_character"]'
+top_k = 5
+
+print(f"\n{'='*20} 开始混合搜索 {'='*20}")
+print(f"查询: '{search_query}'")
+print(f"过滤器: '{search_filter}'")
+
+query_embeddings = ef([search_query])
+dense_vec = query_embeddings["dense"][0]
+sparse_vec = query_embeddings["sparse"]._getrow(0)
+
+# 打印向量信息
+print("\n=== 向量信息 ===")
+print(f"密集向量维度: {len(dense_vec)}")
+print(f"密集向量前5个元素: {dense_vec[:5]}")
+print(f"密集向量范数: {np.linalg.norm(dense_vec):.4f}")
+
+print(f"\n稀疏向量维度: {sparse_vec.shape[1]}")
+print(f"稀疏向量非零元素数量: {sparse_vec.nnz}")
+print("稀疏向量前5个非零元素:")
+for i in range(min(5, sparse_vec.nnz)):
+    print(f"  - 索引: {sparse_vec.indices[i]}, 值: {sparse_vec.data[i]:.4f}")
+density = (sparse_vec.nnz / sparse_vec.shape[1] * 100)
+print(f"\n稀疏向量密度: {density:.8f}%")
+
+# 定义搜索参数
+search_params = {"metric_type": "IP", "params": {}}
+
+# 先执行单独的搜索
+print("\n--- [单独] 密集向量搜索结果 ---")
+dense_results = collection.search(
+    [dense_vec],
+    anns_field="dense_vector",
+    param=search_params,
+    limit=top_k,
+    expr=search_filter,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+for i, hit in enumerate(dense_results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+print("\n--- [单独] 稀疏向量搜索结果 ---")
+sparse_results = collection.search(
+    [sparse_vec],
+    anns_field="sparse_vector",
+    param=search_params,
+    limit=top_k,
+    expr=search_filter,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+for i, hit in enumerate(sparse_results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+print("\n--- [混合] 稀疏+密集向量搜索结果 ---")
+# 创建 RRF 融合器
+rerank = RRFRanker(k=60)
+
+# 创建搜索请求
+dense_req = AnnSearchRequest([dense_vec], "dense_vector", search_params, limit=top_k)
+sparse_req = AnnSearchRequest([sparse_vec], "sparse_vector", search_params, limit=top_k)
+
+# 执行混合搜索
+results = collection.hybrid_search(
+    [sparse_req, dense_req],
+    rerank=rerank,
+    limit=top_k,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+# 打印最终结果
+for i, hit in enumerate(results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+# 7. 清理资源
+milvus_client.release_collection(collection_name=COLLECTION_NAME)
+print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
+milvus_client.drop_collection(COLLECTION_NAME)
+print(f"已删除 Collection: '{COLLECTION_NAME}'")

code/C4/01_hybrid_search_v2.py

@@ -0,0 +1,329 @@
+import json
+import os
+import numpy as np
+import torch
+from transformers import AutoModel, AutoProcessor
+from sklearn.feature_extraction.text import TfidfVectorizer
+from scipy.sparse import csr_matrix
+from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection, AnnSearchRequest, RRFRanker
+
+# 1. 初始化设置
+COLLECTION_NAME = "dragon_siglip_demo"
+MILVUS_URI = "http://localhost:19530"  # 服务器模式
+DATA_PATH = "../../data/C4/metadata/dragon.json"  # 相对路径
+BATCH_SIZE = 50
+
+# 2. 自定义SigLIP嵌入函数类
+class SigLIPEmbeddingFunction:
+    def __init__(self, model_name="google/siglip-base-patch16-256-multilingual", device="cpu"):
+        """
+        初始化SigLIP嵌入函数
+        Args:
+            model_name: SigLIP模型名称
+            device: 设备类型 ("cpu" 或 "cuda")
+        """
+        self.model_name = model_name
+        self.device = device
+        
+        print(f"--> 正在加载 SigLIP 模型: {model_name}")
+        self.model = AutoModel.from_pretrained(model_name)
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model.to(device)
+        self.model.eval()
+        
+        # 初始化TF-IDF作为稀疏向量生成器
+        self.tfidf_vectorizer = TfidfVectorizer(
+            max_features=10000,  # 限制词汇表大小以节省空间
+            stop_words='english',
+            ngram_range=(1, 2)
+        )
+        self.tfidf_fitted = False
+        
+        # 获取文本编码器的输出维度
+        with torch.no_grad():
+            dummy_text = ["test"]
+            inputs = self.processor(text=dummy_text, padding="max_length", return_tensors="pt")
+            outputs = self.model.text_model(**{k: v.to(device) for k, v in inputs.items() if k != 'pixel_values'})
+            self.dense_dim = outputs.pooler_output.shape[-1]
+        
+        print(f"--> SigLIP 模型加载完成。密集向量维度: {self.dense_dim}")
+    
+    @property
+    def dim(self):
+        """返回维度信息，兼容原BGE-M3接口"""
+        return {
+            "dense": self.dense_dim,
+            "sparse": self.tfidf_vectorizer.max_features if self.tfidf_fitted else 10000
+        }
+    
+    def fit_sparse(self, docs):
+        """拟合稀疏向量模型（TF-IDF）"""
+        print("--> 正在拟合 TF-IDF 模型...")
+        self.tfidf_vectorizer.fit(docs)
+        self.tfidf_fitted = True
+        print(f"--> TF-IDF 模型拟合完成。词汇表大小: {len(self.tfidf_vectorizer.vocabulary_)}")
+    
+    def encode_text_dense(self, texts):
+        """使用SigLIP编码文本为密集向量"""
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        dense_vectors = []
+        batch_size = 8  # 减小批次大小以节省内存
+        
+        with torch.no_grad():
+            for i in range(0, len(texts), batch_size):
+                batch_texts = texts[i:i + batch_size]
+                inputs = self.processor(text=batch_texts, padding="max_length", truncation=True, return_tensors="pt")
+                inputs = {k: v.to(self.device) for k, v in inputs.items() if k != 'pixel_values'}
+                
+                outputs = self.model.text_model(**inputs)
+                embeddings = outputs.pooler_output
+                
+                # 归一化向量
+                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+                dense_vectors.extend(embeddings.cpu().numpy())
+        
+        return np.array(dense_vectors)
+    
+    def encode_text_sparse(self, texts):
+        """使用TF-IDF编码文本为稀疏向量"""
+        if not self.tfidf_fitted:
+            raise ValueError("请先调用 fit_sparse() 方法拟合TF-IDF模型")
+        
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        sparse_matrix = self.tfidf_vectorizer.transform(texts)
+        return sparse_matrix
+    
+    def __call__(self, texts):
+        """主调用方法，返回密集和稀疏向量"""
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        # 如果还没有拟合稀疏模型，先拟合
+        if not self.tfidf_fitted:
+            self.fit_sparse(texts)
+        
+        dense_vectors = self.encode_text_dense(texts)
+        sparse_vectors = self.encode_text_sparse(texts)
+        
+        return {
+            "dense": dense_vectors,
+            "sparse": sparse_vectors
+        }
+
+# 3. 连接 Milvus 并初始化嵌入模型
+print(f"--> 正在连接到 Milvus: {MILVUS_URI}")
+connections.connect(uri=MILVUS_URI)
+
+print("--> 正在初始化 SigLIP 嵌入模型...")
+ef = SigLIPEmbeddingFunction(device="cpu")  # 如果有GPU可以改为"cuda"
+
+# 4. 创建 Collection
+milvus_client = MilvusClient(uri=MILVUS_URI)
+if milvus_client.has_collection(COLLECTION_NAME):
+    print(f"--> 正在删除已存在的 Collection '{COLLECTION_NAME}'...")
+    milvus_client.drop_collection(COLLECTION_NAME)
+
+fields = [
+    FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
+    FieldSchema(name="img_id", dtype=DataType.VARCHAR, max_length=100),
+    FieldSchema(name="path", dtype=DataType.VARCHAR, max_length=256),
+    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
+    FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=4096),
+    FieldSchema(name="category", dtype=DataType.VARCHAR, max_length=64),
+    FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=128),
+    FieldSchema(name="environment", dtype=DataType.VARCHAR, max_length=64),
+    FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
+    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=ef.dim["dense"])
+]
+
+# 如果集合不存在，则创建它及索引
+if not milvus_client.has_collection(COLLECTION_NAME):
+    print(f"--> 正在创建 Collection '{COLLECTION_NAME}'...")
+    schema = CollectionSchema(fields, description="使用SigLIP的龙混合检索示例")
+    # 创建集合
+    collection = Collection(name=COLLECTION_NAME, schema=schema, consistency_level="Strong")
+    print("--> Collection 创建成功。")
+
+    # 5. 创建索引
+    print("--> 正在为新集合创建索引...")
+    sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
+    collection.create_index("sparse_vector", sparse_index)
+    print("稀疏向量索引创建成功。")
+
+    dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
+    collection.create_index("dense_vector", dense_index)
+    print("密集向量索引创建成功。")
+
+collection = Collection(COLLECTION_NAME)
+
+# 6. 加载数据并插入
+collection.load()
+print(f"--> Collection '{COLLECTION_NAME}' 已加载到内存。")
+
+if collection.is_empty:
+    print(f"--> Collection 为空，开始插入数据...")
+    if not os.path.exists(DATA_PATH):
+        raise FileNotFoundError(f"数据文件未找到: {DATA_PATH}")
+    with open(DATA_PATH, 'r', encoding='utf-8') as f:
+        dataset = json.load(f)
+
+    docs, metadata = [], []
+    for item in dataset:
+        parts = [
+            item.get('title', ''),
+            item.get('description', ''),
+            item.get('location', ''),
+            item.get('environment', ''),
+            # *item.get('combat_details', {}).get('combat_style', []),
+            # *item.get('combat_details', {}).get('abilities_used', []),
+            # item.get('scene_info', {}).get('time_of_day', '')
+        ]
+        docs.append(' '.join(filter(None, parts)))
+        metadata.append(item)
+    print(f"--> 数据加载完成，共 {len(docs)} 条。")
+
+    print("--> 正在生成向量嵌入...")
+    embeddings = ef(docs)
+    print("--> 向量生成完成。")
+
+    print("--> 正在分批插入数据...")
+    # 为每个字段准备批量数据
+    img_ids = [doc["img_id"] for doc in metadata]
+    paths = [doc["path"] for doc in metadata]
+    titles = [doc["title"] for doc in metadata]
+    descriptions = [doc["description"] for doc in metadata]
+    categories = [doc["category"] for doc in metadata]
+    locations = [doc["location"] for doc in metadata]
+    environments = [doc["environment"] for doc in metadata]
+    
+    # 获取向量 - 注意SigLIP返回的格式与BGE-M3不同
+    sparse_vectors = []
+    dense_vectors = embeddings["dense"].tolist()
+    
+    # 将稀疏矩阵转换为Milvus可接受的格式
+    sparse_matrix = embeddings["sparse"]
+    for i in range(sparse_matrix.shape[0]):
+        row = sparse_matrix.getrow(i)
+        # 创建稀疏向量字典格式
+        sparse_dict = {}
+        for j in range(row.nnz):
+            sparse_dict[row.indices[j]] = float(row.data[j])
+        sparse_vectors.append(sparse_dict)
+    
+    # 插入数据
+    collection.insert([
+        img_ids,
+        paths,
+        titles,
+        descriptions,
+        categories,
+        locations,
+        environments,
+        sparse_vectors,
+        dense_vectors
+    ])
+    
+    collection.flush()
+    print(f"--> 数据插入完成，总数: {collection.num_entities}")
+else:
+    print(f"--> Collection 中已有 {collection.num_entities} 条数据，跳过插入。")
+
+# 7. 执行搜索
+search_query = "悬崖上的巨龙"
+search_filter = 'category in ["western_dragon", "chinese_dragon", "movie_character"]'
+top_k = 5
+
+print(f"\n{'='*20} 开始混合搜索 {'='*20}")
+print(f"查询: '{search_query}'")
+print(f"过滤器: '{search_filter}'")
+
+# 生成查询向量
+query_embeddings = ef([search_query])
+dense_vec = query_embeddings["dense"][0].tolist()
+
+# 处理稀疏向量
+sparse_matrix = query_embeddings["sparse"]
+sparse_row = sparse_matrix.getrow(0)
+sparse_dict = {}
+for j in range(sparse_row.nnz):
+    sparse_dict[sparse_row.indices[j]] = float(sparse_row.data[j])
+
+# 打印向量信息
+print("\n=== 向量信息 ===")
+print(f"密集向量维度: {len(dense_vec)}")
+print(f"密集向量前5个元素: {dense_vec[:5]}")
+print(f"密集向量范数: {np.linalg.norm(dense_vec):.4f}")
+
+print(f"\n稀疏向量维度: {sparse_matrix.shape[1]}")
+print(f"稀疏向量非零元素数量: {sparse_row.nnz}")
+print("稀疏向量前5个非零元素:")
+for i, (idx, val) in enumerate(list(sparse_dict.items())[:5]):
+    print(f"  - 索引: {idx}, 值: {val:.4f}")
+density = (sparse_row.nnz / sparse_matrix.shape[1] * 100)
+print(f"\n稀疏向量密度: {density:.8f}%")
+
+# 定义搜索参数
+search_params = {"metric_type": "IP", "params": {}}
+
+# 先执行单独的搜索
+print("\n--- [单独] 密集向量搜索结果 ---")
+dense_results = collection.search(
+    [dense_vec],
+    anns_field="dense_vector",
+    param=search_params,
+    limit=top_k,
+    expr=search_filter,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+for i, hit in enumerate(dense_results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+print("\n--- [单独] 稀疏向量搜索结果 ---")
+sparse_results = collection.search(
+    [sparse_dict],
+    anns_field="sparse_vector",
+    param=search_params,
+    limit=top_k,
+    expr=search_filter,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+for i, hit in enumerate(sparse_results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+print("\n--- [混合] 稀疏+密集向量搜索结果 ---")
+# 创建 RRF 融合器
+rerank = RRFRanker(k=60)
+
+# 创建搜索请求
+dense_req = AnnSearchRequest([dense_vec], "dense_vector", search_params, limit=top_k)
+sparse_req = AnnSearchRequest([sparse_dict], "sparse_vector", search_params, limit=top_k)
+
+# 执行混合搜索
+results = collection.hybrid_search(
+    [sparse_req, dense_req],
+    rerank=rerank,
+    limit=top_k,
+    output_fields=["title", "path", "description", "category", "location", "environment"]
+)[0]
+
+# 打印最终结果
+for i, hit in enumerate(results):
+    print(f"{i+1}. {hit.entity.get('title')} (Score: {hit.distance:.4f})")
+    print(f"    路径: {hit.entity.get('path')}")
+    print(f"    描述: {hit.entity.get('description')[:100]}...")
+
+# 8. 清理资源
+milvus_client.release_collection(collection_name=COLLECTION_NAME)
+print(f"已从内存中释放 Collection: '{COLLECTION_NAME}'")
+milvus_client.drop_collection(COLLECTION_NAME)
+print(f"已删除 Collection: '{COLLECTION_NAME}'")

code/C4/02_text_to_metadata_filter.py

@@ -0,0 +1,111 @@
+import os
+from langchain_deepseek import ChatDeepSeek 
+from langchain_community.document_loaders import BiliBiliLoader
+from langchain.chains.query_constructor.base import AttributeInfo
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain_community.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+# 1. 初始化视频数据
+video_urls = [
+    "https://www.bilibili.com/video/BV1Bo4y1A7FU", 
+    "https://www.bilibili.com/video/BV1ug4y157xA",
+    "https://www.bilibili.com/video/BV1yh411V7ge",
+]
+
+bili = []
+try:
+    loader = BiliBiliLoader(video_urls=video_urls)
+    docs = loader.load()
+    
+    for doc in docs:
+        original = doc.metadata
+        
+        # 提取基本元数据字段
+        metadata = {
+            'title': original.get('title', '未知标题'),
+            'author': original.get('owner', {}).get('name', '未知作者'),
+            'source': original.get('bvid', '未知ID'),
+            'view_count': original.get('stat', {}).get('view', 0),
+            'length': original.get('duration', 0),
+        }
+        
+        doc.metadata = metadata
+        bili.append(doc)
+        
+except Exception as e:
+    print(f"加载BiliBili视频失败: {str(e)}")
+
+if not bili:
+    print("没有成功加载任何视频，程序退出")
+    exit()
+
+# 2. 创建向量存储
+embed_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+vectorstore = Chroma.from_documents(bili, embed_model)
+
+# 3. 配置元数据字段信息
+metadata_field_info = [
+    AttributeInfo(
+        name="title",
+        description="视频标题（字符串）",
+        type="string", 
+    ),
+    AttributeInfo(
+        name="author",
+        description="视频作者（字符串）",
+        type="string",
+    ),
+    AttributeInfo(
+        name="view_count",
+        description="视频观看次数（整数）",
+        type="integer",
+    ),
+    AttributeInfo(
+        name="length",
+        description="视频长度（整数）",
+        type="integer"
+    )
+]
+
+# 4. 创建自查询检索器
+llm = ChatDeepSeek(
+    model="deepseek-chat", 
+    temperature=0, 
+    api_key=os.getenv("DEEPSEEK_API_KEY")
+    )
+
+retriever = SelfQueryRetriever.from_llm(
+    llm=llm,
+    vectorstore=vectorstore,
+    document_contents="记录视频标题、作者、观看次数等信息的视频元数据",
+    metadata_field_info=metadata_field_info,
+    enable_limit=True,
+    verbose=True
+)
+
+# 5. 执行查询示例
+queries = [
+    "时间最短的视频",
+    "时长大于600秒的视频"
+]
+
+for query in queries:
+    print(f"\n--- 查询: '{query}' ---")
+    results = retriever.invoke(query)
+    if results:
+        for doc in results:
+            title = doc.metadata.get('title', '未知标题')
+            author = doc.metadata.get('author', '未知作者')
+            view_count = doc.metadata.get('view_count', '未知')
+            length = doc.metadata.get('length', '未知')
+            print(f"标题: {title}")
+            print(f"作者: {author}")
+            print(f"观看次数: {view_count}")
+            print(f"时长: {length}秒")
+            print("="*50)
+    else:
+        print("未找到匹配的视频")

code/C4/03_text2sql_demo.py

@@ -0,0 +1,220 @@
+import os
+import sys
+import sqlite3
+
+# 添加text2sql模块路径
+sys.path.append(os.path.join(os.path.dirname(__file__), 'text2sql'))
+
+from text2sql.text2sql_agent import SimpleText2SQLAgent
+
+
+def setup_demo():
+    """设置演示环境"""
+    print("=== Text2SQL框架演示 ===\n")
+    
+    # 检查API密钥
+    api_key = os.getenv("DEEPSEEK_API_KEY")
+    if not api_key:
+        print("先设置DEEPSEEK_API_KEY环境变量")
+        return None
+    
+    # 创建演示数据库
+    print("创建演示数据库...")
+    db_path = create_demo_database()
+    
+    # 初始化Text2SQL代理
+    print("初始化Text2SQL代理...")
+    agent = SimpleText2SQLAgent(api_key=api_key)
+    
+    # 连接数据库
+    print("连接数据库...")
+    if not agent.connect_database(db_path):
+        print("数据库连接失败!")
+        return None
+    
+    # 加载知识库
+    print("加载知识库...")
+    try:
+        agent.load_knowledge_base()
+        print("知识库加载成功!")
+    except Exception as e:
+        print(f"知识库加载失败: {str(e)}")
+        return None
+    
+    return agent, db_path
+
+
+def create_demo_database():
+    """创建演示数据库"""
+    db_path = "text2sql_demo.db"
+    
+    if os.path.exists(db_path):
+        os.remove(db_path)
+    
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    
+    # 创建用户表
+    cursor.execute("""
+        CREATE TABLE users (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            email TEXT UNIQUE,
+            age INTEGER,
+            city TEXT
+        )
+    """)
+    
+    # 创建产品表
+    cursor.execute("""
+        CREATE TABLE products (
+            id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            category TEXT,
+            price REAL,
+            stock INTEGER
+        )
+    """)
+    
+    # 创建订单表
+    cursor.execute("""
+        CREATE TABLE orders (
+            id INTEGER PRIMARY KEY,
+            user_id INTEGER,
+            product_id INTEGER,
+            quantity INTEGER,
+            order_date TEXT,
+            total_price REAL,
+            FOREIGN KEY (user_id) REFERENCES users(id),
+            FOREIGN KEY (product_id) REFERENCES products(id)
+        )
+    """)
+    
+    # 插入示例数据
+    users_data = [
+        (1, '张三', 'zhangsan@email.com', 25, '北京'),
+        (2, '李四', 'lisi@email.com', 32, '上海'),
+        (3, '王五', 'wangwu@email.com', 28, '广州'),
+        (4, '赵六', 'zhaoliu@email.com', 35, '深圳'),
+        (5, '陈七', 'chenqi@email.com', 29, '杭州'),
+    ]
+    
+    products_data = [
+        (1, 'iPhone 15', '电子产品', 7999.0, 50),
+        (2, 'MacBook Pro', '电子产品', 12999.0, 20),
+        (3, 'Nike运动鞋', '服装', 599.0, 100),
+        (4, '办公椅', '家具', 899.0, 30),
+        (5, '台灯', '家具', 199.0, 80),
+        (6, 'iPad', '电子产品', 3999.0, 40),
+        (7, 'Adidas外套', '服装', 399.0, 60),
+    ]
+    
+    orders_data = [
+        (1, 1, 1, 1, '2024-01-15', 7999.0),
+        (2, 2, 3, 2, '2024-01-16', 1198.0),
+        (3, 3, 5, 1, '2024-01-17', 199.0),
+        (4, 1, 2, 1, '2024-01-18', 12999.0),
+        (5, 4, 4, 1, '2024-01-19', 899.0),
+        (6, 5, 6, 1, '2024-01-20', 3999.0),
+        (7, 2, 7, 1, '2024-01-21', 399.0),
+    ]
+    
+    cursor.executemany("INSERT INTO users VALUES (?, ?, ?, ?, ?)", users_data)
+    cursor.executemany("INSERT INTO products VALUES (?, ?, ?, ?, ?)", products_data)
+    cursor.executemany("INSERT INTO orders VALUES (?, ?, ?, ?, ?, ?)", orders_data)
+    
+    conn.commit()
+    conn.close()
+    
+    print(f"演示数据库已创建: {db_path}")
+    return db_path
+
+
+def run_demo_queries(agent):
+    """运行演示查询"""
+    demo_questions = [
+        "查询所有用户的姓名和邮箱",
+        "年龄大于30的用户有哪些",
+        "哪些产品的库存少于50",
+        "查询来自北京的用户的所有订单",
+        "统计每个城市的用户数量",
+        "查询价格在500-8000之间的产品"
+    ]
+    
+    print("\n开始运行演示查询...\n")
+    
+    success_count = 0
+    
+    for i, question in enumerate(demo_questions, 1):
+        print(f"问题 {i}: {question}")
+        print("-" * 60)
+        
+        try:
+            result = agent.query(question)
+            
+            if result["success"]:
+                print(f"成功! SQL: {result['sql']}")
+                
+                if isinstance(result["results"], dict) and "rows" in result["results"]:
+                    count = result["results"]["count"]
+                    print(f"返回 {count} 行数据")
+                    
+                    # 显示前2行数据
+                    if count > 0:
+                        for j, row in enumerate(result["results"]["rows"][:2]):
+                            row_str = " | ".join(f"{k}: {v}" for k, v in row.items())
+                            print(f"  {j+1}. {row_str}")
+                        
+                        if count > 2:
+                            print(f"  ... 还有 {count - 2} 行")
+                else:
+                    print(f"结果: {result['results']}")
+                
+                success_count += 1
+                
+            else:
+                print(f"失败: {result['error']}")
+                print(f"SQL: {result['sql']}")
+                
+        except Exception as e:
+            print(f"执行错误: {str(e)}")
+        
+        print()
+    
+    # 输出统计
+    total_count = len(demo_questions)
+
+
+def cleanup(agent, db_path):
+    """清理资源"""
+    print("\n清理资源...")
+    
+    if agent:
+        agent.cleanup()
+    
+    if os.path.exists(db_path):
+        os.remove(db_path)
+        print(f"已删除演示数据库: {db_path}")
+
+
+def main():
+    """主函数"""
+    # 设置演示环境
+    setup_result = setup_demo()
+    
+    if setup_result is None:
+        return
+    
+    agent, db_path = setup_result
+    
+    try:
+        # 运行演示查询
+        run_demo_queries(agent)
+        
+    finally:
+        # 清理资源
+        cleanup(agent, db_path)
+
+
+if __name__ == "__main__":
+    main() 

code/C4/03_text2sql_demo_v2.py

@@ -0,0 +1,377 @@
+import os
+import json
+import sqlite3
+import numpy as np
+from typing import List, Dict, Any
+from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+from pymilvus import connections, MilvusClient, FieldSchema, CollectionSchema, DataType, Collection
+
+
+class BGESmallEmbeddingFunction:
+    """BGE-Small中文嵌入函数，用于Text2SQL知识库向量化"""
+    
+    def __init__(self, model_name="BAAI/bge-small-zh-v1.5", device="cpu"):
+        self.model_name = model_name
+        self.device = device
+        self.model = SentenceTransformer(model_name, device=device)
+        self.dense_dim = self.model.get_sentence_embedding_dimension()
+    
+    def encode_text(self, texts):
+        """编码文本为密集向量"""
+        if isinstance(texts, str):
+            texts = [texts]
+        
+        embeddings = self.model.encode(
+            texts,
+            normalize_embeddings=True,
+            batch_size=16,
+            convert_to_numpy=True
+        )
+        
+        return embeddings
+    
+    @property
+    def dim(self):
+        """返回向量维度"""
+        return self.dense_dim
+
+
+class SimpleKnowledgeBase:
+    """简化的知识库，使用BGE-Small进行向量检索"""
+    
+    def __init__(self, milvus_uri: str = "http://localhost:19530"):
+        self.milvus_uri = milvus_uri
+        self.collection_name = "text2sql_knowledge_base"
+        self.milvus_client = None
+        self.collection = None
+        
+        self.embedding_function = BGESmallEmbeddingFunction(
+            model_name="BAAI/bge-small-zh-v1.5",
+            device="cpu"
+        )
+        
+        self.sql_examples = []
+        self.table_schemas = []
+        self.data_loaded = False
+    
+    def connect_milvus(self):
+        """连接Milvus数据库"""
+        connections.connect(uri=self.milvus_uri)
+        self.milvus_client = MilvusClient(uri=self.milvus_uri)
+        return True
+    
+    def create_collection(self):
+        """创建Milvus集合"""
+        if not self.milvus_client:
+            self.connect_milvus()
+        
+        if self.milvus_client.has_collection(self.collection_name):
+            self.milvus_client.drop_collection(self.collection_name)
+        
+        fields = [
+            FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
+            FieldSchema(name="content_type", dtype=DataType.VARCHAR, max_length=50),
+            FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=1000),
+            FieldSchema(name="sql", dtype=DataType.VARCHAR, max_length=2000),
+            FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=1000),
+            FieldSchema(name="table_name", dtype=DataType.VARCHAR, max_length=100),
+            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim)
+        ]
+        
+        schema = CollectionSchema(fields, description="Text2SQL知识库")
+        self.collection = Collection(name=self.collection_name, schema=schema, consistency_level="Strong")
+        
+        index_params = {"index_type": "AUTOINDEX", "metric_type": "IP", "params": {}}
+        self.collection.create_index("embedding", index_params)
+        
+        return True
+    
+    def load_data(self):
+        """加载知识库数据"""        
+        data_dir = os.path.join(os.path.dirname(__file__), "data")
+        
+        self.load_sql_examples(data_dir)
+        self.load_table_schemas(data_dir)
+        self.vectorize_and_store()
+        
+        self.data_loaded = True
+    
+    def load_sql_examples(self, data_dir: str):
+        """加载SQL示例"""
+        sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
+        
+        default_examples = [
+            {"question": "查询所有用户信息", "sql": "SELECT * FROM users", "description": "获取用户记录", "database": "sqlite"},
+            {"question": "年龄大于30的用户", "sql": "SELECT * FROM users WHERE age > 30", "description": "年龄筛选", "database": "sqlite"},
+            {"question": "统计用户总数", "sql": "SELECT COUNT(*) as user_count FROM users", "description": "用户计数", "database": "sqlite"},
+            {"question": "查询库存不足的产品", "sql": "SELECT * FROM products WHERE stock < 50", "description": "库存筛选", "database": "sqlite"},
+            {"question": "查询用户订单信息", "sql": "SELECT u.name, p.name, o.quantity FROM orders o JOIN users u ON o.user_id = u.id JOIN products p ON o.product_id = p.id", "description": "订单详情", "database": "sqlite"},
+            {"question": "按城市统计用户", "sql": "SELECT city, COUNT(*) as count FROM users GROUP BY city", "description": "城市分组", "database": "sqlite"}
+        ]
+        
+        if os.path.exists(sql_examples_path):
+            with open(sql_examples_path, 'r', encoding='utf-8') as f:
+                self.sql_examples = json.load(f)
+        else:
+            self.sql_examples = default_examples
+            os.makedirs(data_dir, exist_ok=True)
+            with open(sql_examples_path, 'w', encoding='utf-8') as f:
+                json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
+    
+    def load_table_schemas(self, data_dir: str):
+        """加载表结构信息"""
+        schema_path = os.path.join(data_dir, "table_schemas.json")
+        
+        default_schemas = [
+            {
+                "table_name": "users",
+                "description": "用户信息表",
+                "columns": [
+                    {"name": "id", "type": "INTEGER", "description": "用户ID"},
+                    {"name": "name", "type": "VARCHAR", "description": "用户姓名"},
+                    {"name": "age", "type": "INTEGER", "description": "用户年龄"},
+                    {"name": "email", "type": "VARCHAR", "description": "邮箱地址"},
+                    {"name": "city", "type": "VARCHAR", "description": "所在城市"},
+                    {"name": "created_at", "type": "DATETIME", "description": "创建时间"}
+                ]
+            },
+            {
+                "table_name": "products",
+                "description": "产品信息表",
+                "columns": [
+                    {"name": "id", "type": "INTEGER", "description": "产品ID"},
+                    {"name": "product_name", "type": "VARCHAR", "description": "产品名称"},
+                    {"name": "category", "type": "VARCHAR", "description": "产品类别"},
+                    {"name": "price", "type": "DECIMAL", "description": "产品价格"},
+                    {"name": "stock", "type": "INTEGER", "description": "库存数量"},
+                    {"name": "description", "type": "TEXT", "description": "产品描述"}
+                ]
+            },
+            {
+                "table_name": "orders",
+                "description": "订单信息表",
+                "columns": [
+                    {"name": "id", "type": "INTEGER", "description": "订单ID"},
+                    {"name": "user_id", "type": "INTEGER", "description": "用户ID"},
+                    {"name": "product_id", "type": "INTEGER", "description": "产品ID"},
+                    {"name": "quantity", "type": "INTEGER", "description": "购买数量"},
+                    {"name": "total_price", "type": "DECIMAL", "description": "总价格"},
+                    {"name": "order_date", "type": "DATETIME", "description": "订单日期"}
+                ]
+            }
+        ]
+        
+        if os.path.exists(schema_path):
+            with open(schema_path, 'r', encoding='utf-8') as f:
+                self.table_schemas = json.load(f)
+        else:
+            self.table_schemas = default_schemas
+            os.makedirs(data_dir, exist_ok=True)
+            with open(schema_path, 'w', encoding='utf-8') as f:
+                json.dump(self.table_schemas, f, ensure_ascii=False, indent=2)
+    
+    def vectorize_and_store(self):
+        """向量化数据并存储到Milvus"""
+        self.create_collection()
+        
+        all_texts = []
+        all_metadata = []
+        
+        for example in self.sql_examples:
+            text = f"问题: {example['question']} SQL: {example['sql']} 描述: {example.get('description', '')}"
+            all_texts.append(text)
+            all_metadata.append({
+                "content_type": "sql_example",
+                "question": example['question'],
+                "sql": example['sql'],
+                "description": example.get('description', ''),
+                "table_name": ""
+            })
+        
+        for schema in self.table_schemas:
+            columns_desc = ", ".join([f"{col['name']} ({col['type']}): {col.get('description', '')}" 
+                                    for col in schema['columns']])
+            text = f"表 {schema['table_name']}: {schema['description']} 字段: {columns_desc}"
+            all_texts.append(text)
+            all_metadata.append({
+                "content_type": "table_schema",
+                "question": "",
+                "sql": "",
+                "description": schema['description'],
+                "table_name": schema['table_name']
+            })
+        
+        embeddings = self.embedding_function.encode_text(all_texts)
+        
+        insert_data = []
+        for i, (embedding, metadata) in enumerate(zip(embeddings, all_metadata)):
+            insert_data.append([
+                metadata["content_type"],
+                metadata["question"],
+                metadata["sql"],
+                metadata["description"],
+                metadata["table_name"],
+                embedding.tolist()
+            ])
+        
+        self.collection.insert(insert_data)
+        self.collection.flush()
+        self.collection.load()
+    
+    def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        """搜索相关的知识库信息"""
+        if not self.data_loaded:
+            self.load_data()
+        
+        query_embedding = self.embedding_function.encode_text([query])[0]
+        
+        search_params = {"metric_type": "IP", "params": {}}
+        results = self.collection.search(
+            [query_embedding.tolist()],
+            anns_field="embedding",
+            param=search_params,
+            limit=top_k,
+            output_fields=["content_type", "question", "sql", "description", "table_name"]
+        )[0]
+        
+        formatted_results = []
+        for hit in results:
+            result = {
+                "score": float(hit.distance),
+                "content_type": hit.entity.get("content_type"),
+                "question": hit.entity.get("question"),
+                "sql": hit.entity.get("sql"),
+                "description": hit.entity.get("description"),
+                "table_name": hit.entity.get("table_name")
+            }
+            formatted_results.append(result)
+        
+        return formatted_results
+    
+    def _fallback_search(self, query: str, top_k: int) -> List[Dict[str, Any]]:
+        """降级搜索方法（简单文本匹配）"""
+        results = []
+        query_lower = query.lower()
+        
+        for example in self.sql_examples:
+            question_lower = example['question'].lower()
+            sql_lower = example['sql'].lower()
+            
+            score = 0
+            for word in query_lower.split():
+                if word in question_lower:
+                    score += 2
+                if word in sql_lower:
+                    score += 1
+            
+            if score > 0:
+                results.append({
+                    "score": score,
+                    "content_type": "sql_example",
+                    "question": example['question'],
+                    "sql": example['sql'],
+                    "description": example.get('description', ''),
+                    "table_name": ""
+                })
+        
+        results.sort(key=lambda x: x['score'], reverse=True)
+        return results[:top_k]
+    
+    def add_sql_example(self, question: str, sql: str, description: str = ""):
+        """添加新的SQL示例"""
+        new_example = {
+            "question": question,
+            "sql": sql,
+            "description": description,
+            "database": "sqlite"
+        }
+        self.sql_examples.append(new_example)
+        
+        data_dir = os.path.join(os.path.dirname(__file__), "data")
+        sql_examples_path = os.path.join(data_dir, "qsql_examples.json")
+        
+        with open(sql_examples_path, 'w', encoding='utf-8') as f:
+            json.dump(self.sql_examples, f, ensure_ascii=False, indent=2)
+        
+        if self.collection and self.data_loaded:
+            text = f"问题: {question} SQL: {sql} 描述: {description}"
+            embedding = self.embedding_function.encode_text([text])[0]
+            
+            insert_data = [[
+                "sql_example",
+                question,
+                sql,
+                description,
+                "",
+                embedding.tolist()
+            ]]
+            
+            self.collection.insert(insert_data)
+            self.collection.flush()
+    
+    def cleanup(self):
+        """清理资源"""
+        if self.collection:
+            self.collection.release()
+        
+        if self.milvus_client and self.milvus_client.has_collection(self.collection_name):
+            self.milvus_client.drop_collection(self.collection_name)
+
+
+def demo():
+    """简单演示"""
+    # 模型测试
+    embedding_function = BGESmallEmbeddingFunction()
+    test_texts = ["查询用户", "统计数据"]
+    embeddings = embedding_function.encode_text(test_texts)
+    print(f"向量维度: {embeddings.shape}")
+    
+    # 数据库查询演示
+    db_path = "demo.db"
+    
+    if os.path.exists(db_path):
+        os.remove(db_path)
+    
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    
+    cursor.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, name TEXT, age INTEGER, city TEXT)")
+    
+    users_data = [(1, '张三', 25, '北京'), (2, '李四', 32, '上海'), (3, '王五', 35, '深圳')]
+    cursor.executemany("INSERT INTO users VALUES (?, ?, ?, ?)", users_data)
+    
+    conn.commit()
+    
+    # 执行查询
+    test_sqls = [
+        ("查询所有用户", "SELECT * FROM users"),
+        ("年龄大于30的用户", "SELECT * FROM users WHERE age > 30"),
+        ("统计用户总数", "SELECT COUNT(*) FROM users")
+    ]
+    
+    for i, (question, sql) in enumerate(test_sqls, 1):
+        print(f"\n问题 {i}: {question}")
+        print("-" * 40)
+        print(f"SQL: {sql}")
+        
+        cursor.execute(sql)
+        rows = cursor.fetchall()
+        
+        if rows:
+            print(f"返回 {len(rows)} 行数据")
+            for j, row in enumerate(rows[:2], 1):
+                print(f"  {j}. {row}")
+            
+            if len(rows) > 2:
+                print(f"  ... 还有 {len(rows) - 2} 行")
+        else:
+            print("无数据返回")
+    
+    conn.close()
+    os.remove(db_path)
+
+
+if __name__ == "__main__":
+    demo()

code/C4/04_text_to_metadata_filter_v2.py

@@ -0,0 +1,149 @@
+import os
+from langchain_deepseek import ChatDeepSeek 
+from langchain_community.document_loaders import BiliBiliLoader
+from langchain.chains.query_constructor.base import AttributeInfo
+from openai import OpenAI
+from langchain_community.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+# 1. 初始化视频数据
+video_urls = [
+    "https://www.bilibili.com/video/BV1Bo4y1A7FU", 
+    "https://www.bilibili.com/video/BV1ug4y157xA",
+    "https://www.bilibili.com/video/BV1yh411V7ge",
+]
+
+bili = []
+try:
+    loader = BiliBiliLoader(video_urls=video_urls)
+    docs = loader.load()
+    
+    for doc in docs:
+        original = doc.metadata
+        
+        # 提取基本元数据字段
+        metadata = {
+            'title': original.get('title', '未知标题'),
+            'author': original.get('owner', {}).get('name', '未知作者'),
+            'source': original.get('bvid', '未知ID'),
+            'view_count': original.get('stat', {}).get('view', 0),
+            'length': original.get('duration', 0),
+        }
+        
+        doc.metadata = metadata
+        bili.append(doc)
+        
+except Exception as e:
+    print(f"加载BiliBili视频失败: {str(e)}")
+
+if not bili:
+    print("没有成功加载任何视频，程序退出")
+    exit()
+
+# 2. 创建向量存储
+embed_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+vectorstore = Chroma.from_documents(bili, embed_model)
+
+# 3. 配置元数据字段信息
+metadata_field_info = [
+    AttributeInfo(
+        name="title",
+        description="视频标题（字符串）",
+        type="string", 
+    ),
+    AttributeInfo(
+        name="author",
+        description="视频作者（字符串）",
+        type="string",
+    ),
+    AttributeInfo(
+        name="view_count",
+        description="视频观看次数（整数）",
+        type="integer",
+    ),
+    AttributeInfo(
+        name="length",
+        description="视频长度（整数）",
+        type="integer"
+    )
+]
+
+# 4. 初始化LLM客户端
+client = OpenAI(
+    base_url="https://api.deepseek.com",
+    api_key=os.getenv("DEEPSEEK_API_KEY")
+)
+
+# 5. 获取所有文档用于排序
+all_documents = vectorstore.similarity_search("", k=len(bili)) 
+
+# 6. 执行查询示例
+queries = [
+    "时间最短的视频",
+    "播放量最高的视频"
+]
+
+for query in queries:
+    print(f"\n--- 原始查询: '{query}' ---")
+
+    # 使用大模型将自然语言转换为排序指令
+    prompt = f"""你是一个智能助手，请将用户的问题转换成一个用于排序视频的JSON指令。
+
+你需要识别用户想要排序的字段和排序方向。
+- 排序字段必须是 'view_count' (观看次数) 或 'length' (时长) 之一。
+- 排序方向必须是 'asc' (升序) 或 'desc' (降序) 之一。
+
+例如:
+- '时间最短的视频' 或 '哪个视频时间最短' 应转换为 {{"sort_by": "length", "order": "asc"}}
+- '播放量最高的视频' 或 '哪个视频最火' 应转换为 {{"sort_by": "view_count", "order": "desc"}}
+
+请根据以下问题生成JSON指令:
+原始问题: "{query}"
+
+JSON指令:"""
+    
+    response = client.chat.completions.create(
+        model="deepseek-chat",
+        messages=[
+            {"role": "user", "content": prompt}
+        ],
+        temperature=0,
+        response_format={"type": "json_object"}
+    )
+    
+    try:
+        import json
+        instruction_str = response.choices[0].message.content
+        instruction = json.loads(instruction_str)
+        print(f"--- 生成的排序指令: {instruction} ---")
+
+        sort_by = instruction.get('sort_by')
+        order = instruction.get('order')
+
+        if sort_by in ['length', 'view_count'] and order in ['asc', 'desc']:
+            # 在代码中执行排序
+            reverse_order = (order == 'desc')
+            sorted_docs = sorted(all_documents, key=lambda doc: doc.metadata.get(sort_by, 0), reverse=reverse_order)
+            
+            # 获取排序后的第一个结果
+            if sorted_docs:
+                doc = sorted_docs[0]
+                title = doc.metadata.get('title', '未知标题')
+                author = doc.metadata.get('author', '未知作者')
+                view_count = doc.metadata.get('view_count', '未知')
+                length = doc.metadata.get('length', '未知')
+                print(f"标题: {title}")
+                print(f"作者: {author}")
+                print(f"观看次数: {view_count}")
+                print(f"时长: {length}秒")
+                print("="*50)
+            else:
+                print("没有找到任何视频")
+        else:
+            print("生成的指令无效，无法执行排序")
+
+    except (json.JSONDecodeError, KeyError) as e:
+        print(f"解析或执行指令失败: {e}")

code/C4/05_llm_based_routing.py

@@ -0,0 +1,72 @@
+import os
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_deepseek import ChatDeepSeek
+from langchain_core.runnables import RunnableBranch
+
+llm = ChatDeepSeek(
+    model="deepseek-chat", 
+    temperature=0, 
+    api_key=os.getenv("DEEPSEEK_API_KEY")
+    )
+
+# 1. 设置不同菜系的处理链
+sichuan_prompt = ChatPromptTemplate.from_template(
+    "你是一位川菜大厨。请用正宗的川菜做法，回答关于「{question}」的问题。"
+)
+sichuan_chain = sichuan_prompt | llm | StrOutputParser()
+
+cantonese_prompt = ChatPromptTemplate.from_template(
+    "你是一位粤菜大厨。请用经典的粤菜做法，回答关于「{question}」的问题。"
+)
+cantonese_chain = cantonese_prompt | llm | StrOutputParser()
+
+# 定义备用通用链
+general_prompt = ChatPromptTemplate.from_template(
+    "你是一个美食助手。请回答关于「{question}」的问题。"
+)
+general_chain = general_prompt | llm | StrOutputParser()
+
+
+# 2. 创建路由链
+classifier_prompt = ChatPromptTemplate.from_template(
+    """根据用户问题中提到的菜品，将其分类为：['川菜', '粤菜', 或 '其他']。
+    不要解释你的理由，只返回一个单词的分类结果。
+    问题: {question}"""
+)
+classifier_chain = classifier_prompt | llm | StrOutputParser()
+
+# 定义路由分支
+router_branch = RunnableBranch(
+    (lambda x: "川菜" in x["topic"], sichuan_chain),
+    (lambda x: "粤菜" in x["topic"], cantonese_chain),
+    general_chain  # 默认选项
+)
+
+# 组合成完整路由链
+full_router_chain = {"topic": classifier_chain, "question": lambda x: x["question"]} | router_branch
+print("完整的路由链创建成功。\n")
+
+
+# 3. 运行演示查询
+demo_questions = [
+    {"question": "麻婆豆腐怎么做？"},      # 应该路由到川菜
+    {"question": "白切鸡的正宗做法是什么？"}, # 应该路由到粤菜
+    {"question": "番茄炒蛋需要放糖吗？"}      # 应该路由到其他
+]
+
+for i, item in enumerate(demo_questions, 1):
+    question = item["question"]
+    print(f"\n--- 问题 {i}: {question} ---")
+    
+    try:
+        # 获取路由决策
+        topic = classifier_chain.invoke({"question": question})
+        print(f"路由决策: {topic}")
+
+        # 执行完整链
+        result = full_router_chain.invoke(item)
+        print(f"回答: {result}")
+    except Exception as e:
+        print(f"执行错误: {e}")
+

code/C4/06_embedding_based_routing.py

@@ -0,0 +1,83 @@
+import os
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_deepseek import ChatDeepSeek
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough, RunnablePassthrough
+from langchain_community.utils.math import cosine_similarity
+import numpy as np
+
+# 1. 定义路由描述
+sichuan_route_prompt = "你是一位处理川菜的专家。用户的问题是关于麻辣、辛香、重口味的菜肴，例如水煮鱼、麻婆豆腐、鱼香肉丝、宫保鸡丁、花椒、海椒等。"
+cantonese_route_prompt = "你是一位处理粤菜的专家。用户的问题是关于清淡、鲜美、原汁原味的菜肴，例如白切鸡、老火靓汤、虾饺、云吞面等。"
+
+route_prompts = [sichuan_route_prompt, cantonese_route_prompt]
+route_names = ["川菜", "粤菜"]
+
+# 初始化嵌入模型，并对路由描述进行向量化
+embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh-v1.5")
+route_prompt_embeddings = embeddings.embed_documents(route_prompts)
+print(f"已定义 {len(route_names)} 个路由: {', '.join(route_names)}")
+
+# 2. 定义不同路由的目标链
+llm = ChatDeepSeek(
+    model="deepseek-chat", 
+    temperature=0, 
+    api_key=os.getenv("DEEPSEEK_API_KEY")
+    )
+
+# 定义川菜和粤菜处理链
+sichuan_chain = (
+    PromptTemplate.from_template("你是一位川菜大厨。请用正宗的川菜做法，回答关于「{query}」的问题。")
+    | llm
+    | StrOutputParser()
+)
+cantonese_chain = (
+    PromptTemplate.from_template("你是一位粤菜大厨。请用经典的粤菜做法，回答关于「{query}」的问题。")
+    | llm
+    | StrOutputParser()
+)
+
+route_map = { "川菜": sichuan_chain, "粤菜": cantonese_chain }
+print("川菜和粤菜的处理链创建成功。\n")
+
+# 3. 创建路由函数
+def route(info):
+    # 对用户查询进行嵌入
+    query_embedding = embeddings.embed_query(info["query"])
+    
+    # 计算与各路由提示的余弦相似度
+    similarity_scores = cosine_similarity([query_embedding], route_prompt_embeddings)[0]
+    
+    # 找到最相似的路由
+    chosen_route_index = np.argmax(similarity_scores)
+    chosen_route_name = route_names[chosen_route_index]
+    
+    print(f"路由决策: 检测到问题与“{chosen_route_name}”最相似。")
+    
+    # 获取对应的处理链
+    chosen_chain = route_map[chosen_route_name]
+    
+    # 直接调用选中的链并返回结果
+    return chosen_chain.invoke(info)
+
+# 创建完整的路由链
+full_chain = RunnableLambda(route)
+
+
+# 4. 运行演示查询
+demo_queries = [
+    "水煮鱼怎么做才嫩？",        # 应该路由到川菜
+    "如何做一碗清淡的云吞面？",    # 应该路由到粤菜
+    "麻婆豆腐的核心调料是什么？",  # 应该路由到川菜
+]
+
+for i, query in enumerate(demo_queries, 1):
+    print(f"\n--- 问题 {i}: {query} ---")
+    try:
+        # 传入字典，full_chain 会直接返回最终答案
+        result = full_chain.invoke({"query": query})
+        print(f"回答: {result}")
+    except Exception as e:
+        print(f"执行错误: {e}")
+

code/C4/07_rerank_and_refine.py

@@ -0,0 +1,186 @@
+import os
+from langchain_community.vectorstores import FAISS
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import LLMChainExtractor
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+from langchain_deepseek import ChatDeepSeek
+
+# 导入ColBERT重排器需要的模块
+from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
+from langchain.retrievers.document_compressors import DocumentCompressorPipeline
+from langchain_core.documents import Document
+from typing import Sequence
+import torch
+from transformers import AutoTokenizer, AutoModel
+import torch.nn.functional as F
+
+class ColBERTReranker(BaseDocumentCompressor):
+    """ColBERT重排器"""
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        model_name = "bert-base-uncased"
+
+        # 加载模型和分词器
+        object.__setattr__(self, 'tokenizer', AutoTokenizer.from_pretrained(model_name))
+        object.__setattr__(self, 'model', AutoModel.from_pretrained(model_name))
+        self.model.eval()
+        print(f"ColBERT模型加载完成")
+
+    def encode_text(self, texts):
+        """ColBERT文本编码"""
+        inputs = self.tokenizer(
+            texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=128
+        )
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+
+        embeddings = outputs.last_hidden_state
+        embeddings = F.normalize(embeddings, p=2, dim=-1)
+
+        return embeddings
+
+    def calculate_colbert_similarity(self, query_emb, doc_embs, query_mask, doc_masks):
+        """ColBERT相似度计算（MaxSim操作）"""
+        scores = []
+
+        for i, doc_emb in enumerate(doc_embs):
+            doc_mask = doc_masks[i:i+1]
+
+            # 计算相似度矩阵
+            similarity_matrix = torch.matmul(query_emb, doc_emb.unsqueeze(0).transpose(-2, -1))
+
+            # 应用文档mask
+            doc_mask_expanded = doc_mask.unsqueeze(1)
+            similarity_matrix = similarity_matrix.masked_fill(~doc_mask_expanded.bool(), -1e9)
+
+            # MaxSim操作
+            max_sim_per_query_token = similarity_matrix.max(dim=-1)[0]
+
+            # 应用查询mask
+            query_mask_expanded = query_mask.unsqueeze(0)
+            max_sim_per_query_token = max_sim_per_query_token.masked_fill(~query_mask_expanded.bool(), 0)
+
+            # 求和得到最终分数
+            colbert_score = max_sim_per_query_token.sum(dim=-1).item()
+            scores.append(colbert_score)
+
+        return scores
+
+    def compress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks=None,
+    ) -> Sequence[Document]:
+        """对文档进行ColBERT重排序"""
+        if len(documents) == 0:
+            return documents
+
+        # 编码查询
+        query_inputs = self.tokenizer(
+            [query],
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=128
+        )
+
+        with torch.no_grad():
+            query_outputs = self.model(**query_inputs)
+            query_embeddings = F.normalize(query_outputs.last_hidden_state, p=2, dim=-1)
+
+        # 编码文档
+        doc_texts = [doc.page_content for doc in documents]
+        doc_inputs = self.tokenizer(
+            doc_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=128
+        )
+
+        with torch.no_grad():
+            doc_outputs = self.model(**doc_inputs)
+            doc_embeddings = F.normalize(doc_outputs.last_hidden_state, p=2, dim=-1)
+
+        # 计算ColBERT相似度
+        scores = self.calculate_colbert_similarity(
+            query_embeddings,
+            doc_embeddings,
+            query_inputs['attention_mask'],
+            doc_inputs['attention_mask']
+        )
+
+        # 排序并返回前5个
+        scored_docs = list(zip(documents, scores))
+        scored_docs.sort(key=lambda x: x[1], reverse=True)
+        reranked_docs = [doc for doc, _ in scored_docs[:5]]
+
+        return reranked_docs
+
+
+
+# 初始化配置
+hf_bge_embeddings = HuggingFaceBgeEmbeddings(
+    model_name="BAAI/bge-large-zh-v1.5"
+)
+
+llm = ChatDeepSeek(
+    model="deepseek-chat", 
+    temperature=0.1, 
+    api_key=os.getenv("DEEPSEEK_API_KEY")
+)
+
+# 1. 加载和处理文档
+loader = TextLoader("../../data/C4/txt/ai.txt", encoding="utf-8")
+documents = loader.load()
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
+docs = text_splitter.split_documents(documents)
+
+# 2. 创建向量存储和基础检索器
+vectorstore = FAISS.from_documents(docs, hf_bge_embeddings)
+base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
+
+# 3. 设置ColBERT重排序器
+reranker = ColBERTReranker()
+
+# 4. 设置LLM压缩器
+compressor = LLMChainExtractor.from_llm(llm)
+
+# 5. 使用DocumentCompressorPipeline组装压缩管道
+# 流程: ColBERT重排 -> LLM压缩
+pipeline_compressor = DocumentCompressorPipeline(
+    transformers=[reranker, compressor]
+)
+
+# 6. 创建最终的压缩检索器
+final_retriever = ContextualCompressionRetriever(
+    base_compressor=pipeline_compressor,
+    base_retriever=base_retriever
+)
+
+# 7. 执行查询并展示结果
+query = "AI还有哪些缺陷需要克服？"
+print(f"\n{'='*20} 开始执行查询 {'='*20}")
+print(f"查询: {query}\n")
+
+# 7.1 基础检索结果
+print(f"--- (1) 基础检索结果 (Top 20) ---")
+base_results = base_retriever.get_relevant_documents(query)
+for i, doc in enumerate(base_results):
+    print(f"  [{i+1}] {doc.page_content[:100]}...\n")
+
+# 7.2 使用管道压缩器的最终结果
+print(f"\n--- (2) 管道压缩后结果 (ColBERT重排 + LLM压缩) ---")
+final_results = final_retriever.get_relevant_documents(query)
+for i, doc in enumerate(final_results):
+    print(f"  [{i+1}] {doc.page_content}\n")

code/C4/text2sql/__init__.py

@@ -0,0 +1,17 @@
+"""
+简化的Text2SQL框架
+基于RAGFlow方案实现的Text2SQL框架
+"""
+
+__version__ = "1.0.0"
+__author__ = "RAG Team"
+
+from .knowledge_base import SimpleKnowledgeBase
+from .sql_generator import SimpleSQLGenerator
+from .text2sql_agent import SimpleText2SQLAgent
+
+__all__ = [
+    "SimpleKnowledgeBase",
+    "SimpleSQLGenerator", 
+    "SimpleText2SQLAgent"
+] 

code/C4/text2sql/data/db_descriptions.json

@@ -0,0 +1,57 @@
+[
+  {
+    "table_name": "users",
+    "table_description": "用户信息表，存储注册用户的基本信息",
+    "columns": [
+      {"name": "id", "description": "用户唯一标识符，主键", "type": "INT"},
+      {"name": "name", "description": "用户姓名，不能为空", "type": "VARCHAR(100)"},
+      {"name": "email", "description": "用户邮箱地址，必须唯一", "type": "VARCHAR(150)"},
+      {"name": "age", "description": "用户年龄", "type": "INT"},
+      {"name": "created_at", "description": "用户注册时间", "type": "TIMESTAMP"}
+    ]
+  },
+  {
+    "table_name": "orders",
+    "table_description": "订单表，记录用户的购买订单信息",
+    "columns": [
+      {"name": "id", "description": "订单唯一标识符，主键", "type": "INT"},
+      {"name": "user_id", "description": "下单用户的ID，外键关联users表", "type": "INT"},
+      {"name": "product_name", "description": "购买的产品名称", "type": "VARCHAR(200)"},
+      {"name": "quantity", "description": "购买数量", "type": "INT"},
+      {"name": "price", "description": "订单总价格", "type": "DECIMAL(10,2)"},
+      {"name": "order_date", "description": "下单时间", "type": "TIMESTAMP"}
+    ]
+  },
+  {
+    "table_name": "products",
+    "table_description": "产品表，存储商城中所有产品的信息",
+    "columns": [
+      {"name": "id", "description": "产品唯一标识符，主键", "type": "INT"},
+      {"name": "name", "description": "产品名称", "type": "VARCHAR(200)"},
+      {"name": "category", "description": "产品分类", "type": "VARCHAR(100)"},
+      {"name": "price", "description": "产品单价", "type": "DECIMAL(10,2)"},
+      {"name": "stock", "description": "库存数量", "type": "INT"},
+      {"name": "description", "description": "产品详细描述", "type": "TEXT"}
+    ]
+  },
+  {
+    "table_name": "categories",
+    "table_description": "产品分类表，定义产品的分类信息",
+    "columns": [
+      {"name": "id", "description": "分类唯一标识符，主键", "type": "INT"},
+      {"name": "name", "description": "分类名称，必须唯一", "type": "VARCHAR(100)"},
+      {"name": "description", "description": "分类描述", "type": "TEXT"}
+    ]
+  },
+  {
+    "table_name": "order_items",
+    "table_description": "订单明细表，存储订单中包含的具体商品信息",
+    "columns": [
+      {"name": "id", "description": "订单明细唯一标识符，主键", "type": "INT"},
+      {"name": "order_id", "description": "关联的订单ID，外键", "type": "INT"},
+      {"name": "product_id", "description": "关联的产品ID，外键", "type": "INT"},
+      {"name": "quantity", "description": "该商品在订单中的数量", "type": "INT"},
+      {"name": "unit_price", "description": "该商品的单价（下单时的价格）", "type": "DECIMAL(10,2)"}
+    ]
+  }
+] 

code/C4/text2sql/data/ddl_examples.json

@@ -0,0 +1,27 @@
+[
+  {
+    "table_name": "users",
+    "ddl_statement": "CREATE TABLE users (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL, email VARCHAR(150) UNIQUE, age INT, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
+    "description": "用户信息表，存储用户基本信息"
+  },
+  {
+    "table_name": "orders",
+    "ddl_statement": "CREATE TABLE orders (id INT PRIMARY KEY AUTO_INCREMENT, user_id INT, product_name VARCHAR(200), quantity INT, price DECIMAL(10,2), order_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP, FOREIGN KEY (user_id) REFERENCES users(id))",
+    "description": "订单表，存储用户订单信息"
+  },
+  {
+    "table_name": "products",
+    "ddl_statement": "CREATE TABLE products (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(200) NOT NULL, category VARCHAR(100), price DECIMAL(10,2), stock INT DEFAULT 0, description TEXT)",
+    "description": "产品表，存储产品基本信息"
+  },
+  {
+    "table_name": "categories",
+    "ddl_statement": "CREATE TABLE categories (id INT PRIMARY KEY AUTO_INCREMENT, name VARCHAR(100) NOT NULL UNIQUE, description TEXT)",
+    "description": "产品分类表"
+  },
+  {
+    "table_name": "order_items",
+    "ddl_statement": "CREATE TABLE order_items (id INT PRIMARY KEY AUTO_INCREMENT, order_id INT, product_id INT, quantity INT, unit_price DECIMAL(10,2), FOREIGN KEY (order_id) REFERENCES orders(id), FOREIGN KEY (product_id) REFERENCES products(id))",
+    "description": "订单明细表，存储订单中的具体商品信息"
+  }
+] 

code/C4/text2sql/data/qsql_examples.json

@@ -0,0 +1,62 @@
+[
+  {
+    "question": "查询所有用户的姓名和邮箱",
+    "sql": "SELECT name, email FROM users",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查找年龄大于25岁的用户",
+    "sql": "SELECT * FROM users WHERE age > 25",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询每个用户的订单数量",
+    "sql": "SELECT u.name, COUNT(o.id) as order_count FROM users u LEFT JOIN orders o ON u.id = o.user_id GROUP BY u.id, u.name",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查找最近7天的订单",
+    "sql": "SELECT * FROM orders WHERE order_date >= DATE_SUB(NOW(), INTERVAL 7 DAY)",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询总销售额最高的前5个产品",
+    "sql": "SELECT p.name, SUM(oi.quantity * oi.unit_price) as total_sales FROM products p JOIN order_items oi ON p.id = oi.product_id GROUP BY p.id, p.name ORDER BY total_sales DESC LIMIT 5",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询某个用户的所有订单",
+    "sql": "SELECT o.*, u.name as user_name FROM orders o JOIN users u ON o.user_id = u.id WHERE u.name = '张三'",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询价格在100到500之间的产品",
+    "sql": "SELECT * FROM products WHERE price BETWEEN 100 AND 500",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询库存少于10的产品",
+    "sql": "SELECT * FROM products WHERE stock < 10",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询每个分类的产品数量",
+    "sql": "SELECT category, COUNT(*) as product_count FROM products GROUP BY category",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询订单总金额大于1000的订单",
+    "sql": "SELECT o.*, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id HAVING total_amount > 1000",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询没有下过订单的用户",
+    "sql": "SELECT u.* FROM users u LEFT JOIN orders o ON u.id = o.user_id WHERE o.id IS NULL",
+    "database": "ecommerce"
+  },
+  {
+    "question": "查询平均订单金额",
+    "sql": "SELECT AVG(total_amount) as avg_order_amount FROM (SELECT o.id, SUM(oi.quantity * oi.unit_price) as total_amount FROM orders o JOIN order_items oi ON o.id = oi.order_id GROUP BY o.id) as order_totals",
+    "database": "ecommerce"
+  }
+] 

code/C4/text2sql/knowledge_base.py

@@ -0,0 +1,184 @@
+import json
+import os
+from typing import List, Dict, Any
+from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
+from pymilvus.model.hybrid import BGEM3EmbeddingFunction
+
+
+class SimpleKnowledgeBase:
+    """知识库"""
+    
+    def __init__(self, milvus_uri: str = "http://localhost:19530"):
+        self.milvus_uri = milvus_uri
+        self.client = MilvusClient(uri=milvus_uri)
+        self.embedding_function = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
+        self.collection_name = "text2sql_kb"
+        self._setup_collection()
+    
+    def _setup_collection(self):
+        """设置集合"""
+        if self.client.has_collection(self.collection_name):
+            self.client.drop_collection(self.collection_name)
+        
+        # 定义字段
+        fields = [
+            FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
+            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=4096),
+            FieldSchema(name="type", dtype=DataType.VARCHAR, max_length=32),  # ddl, qsql, description
+            FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=self.embedding_function.dim["dense"])
+        ]
+        
+        schema = CollectionSchema(fields, description="Text2SQL知识库")
+        
+        # 创建集合
+        self.client.create_collection(
+            collection_name=self.collection_name,
+            schema=schema,
+            consistency_level="Strong"
+        )
+        
+        # 创建索引
+        index_params = self.client.prepare_index_params()
+        index_params.add_index(
+            field_name="dense_vector",
+            index_type="AUTOINDEX",
+            metric_type="IP"
+        )
+        
+        self.client.create_index(
+            collection_name=self.collection_name,
+            index_params=index_params
+        )
+    
+    def load_data(self):
+        """加载所有知识库数据"""
+        data_dir = os.path.join(os.path.dirname(__file__), "data")
+        
+        # 加载DDL数据
+        ddl_path = os.path.join(data_dir, "ddl_examples.json")
+        if os.path.exists(ddl_path):
+            with open(ddl_path, 'r', encoding='utf-8') as f:
+                ddl_data = json.load(f)
+            self._add_ddl_data(ddl_data)
+        
+        # 加载Q->SQL数据
+        qsql_path = os.path.join(data_dir, "qsql_examples.json")
+        if os.path.exists(qsql_path):
+            with open(qsql_path, 'r', encoding='utf-8') as f:
+                qsql_data = json.load(f)
+            self._add_qsql_data(qsql_data)
+        
+        # 加载描述数据
+        desc_path = os.path.join(data_dir, "db_descriptions.json")
+        if os.path.exists(desc_path):
+            with open(desc_path, 'r', encoding='utf-8') as f:
+                desc_data = json.load(f)
+            self._add_description_data(desc_data)
+        
+        # 加载集合到内存
+        self.client.load_collection(collection_name=self.collection_name)
+        print("知识库数据加载完成")
+    
+    def _add_ddl_data(self, data: List[Dict]):
+        """添加DDL数据"""
+        contents = []
+        types = []
+        
+        for item in data:
+            content = f"表名: {item.get('table_name', '')}\n"
+            content += f"DDL: {item.get('ddl_statement', '')}\n"
+            content += f"描述: {item.get('description', '')}"
+            
+            contents.append(content)
+            types.append("ddl")
+        
+        self._insert_data(contents, types)
+    
+    def _add_qsql_data(self, data: List[Dict]):
+        """添加Q->SQL数据"""
+        contents = []
+        types = []
+        
+        for item in data:
+            content = f"问题: {item.get('question', '')}\n"
+            content += f"SQL: {item.get('sql', '')}"
+            
+            contents.append(content)
+            types.append("qsql")
+        
+        self._insert_data(contents, types)
+    
+    def _add_description_data(self, data: List[Dict]):
+        """添加描述数据"""
+        contents = []
+        types = []
+        
+        for item in data:
+            content = f"表名: {item.get('table_name', '')}\n"
+            content += f"表描述: {item.get('table_description', '')}\n"
+            
+            columns = item.get('columns', [])
+            if columns:
+                content += "字段信息:\n"
+                for col in columns:
+                    content += f"  - {col.get('name', '')}: {col.get('description', '')} ({col.get('type', '')})\n"
+            
+            contents.append(content)
+            types.append("description")
+        
+        self._insert_data(contents, types)
+    
+    def _insert_data(self, contents: List[str], types: List[str]):
+        """插入数据"""
+        if not contents:
+            return
+        
+        # 生成嵌入
+        embeddings = self.embedding_function(contents)
+        
+        # 构建插入数据，每一行是一个字典
+        data_to_insert = []
+        for i in range(len(contents)):
+            data_to_insert.append({
+                "content": contents[i],
+                "type": types[i],
+                "dense_vector": embeddings["dense"][i]
+            })
+        
+        # 插入数据
+        result = self.client.insert(
+            collection_name=self.collection_name,
+            data=data_to_insert
+        )
+    
+    def search(self, query: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        """搜索相关内容"""
+        self.client.load_collection(collection_name=self.collection_name)
+            
+        query_embeddings = self.embedding_function([query])
+        
+        search_results = self.client.search(
+            collection_name=self.collection_name,
+            data=query_embeddings["dense"],
+            anns_field="dense_vector",
+            search_params={"metric_type": "IP"},
+            limit=top_k,
+            output_fields=["content", "type"]
+        )
+        
+        results = []
+        for hit in search_results[0]:
+            results.append({
+                "content": hit["entity"]["content"],
+                "type": hit["entity"]["type"],
+                "score": hit["distance"]
+            })
+        
+        return results
+    
+    def cleanup(self):
+        """清理资源"""
+        try:
+            self.client.drop_collection(self.collection_name)
+        except:
+            pass