PyTorch 文档搜索工具（项目暂停）

一个具有命令行功能的PyTorch文档语义搜索原型。

当前状态（2025年4月19日）

⚠️ 该项目目前因重大重构而暂停。

该工具提供了一个基本的命令行搜索界面，用于查询PyTorch文档，但需要在多个方面进行实质性改进。尽管核心嵌入和搜索功能在基本层面上正常工作，但在相关性和MCP集成方面仍需进一步开发。

示例输出

$ python scripts/search.py "如何在 PyTorch 中绘制多注意力头？"

找到 5 个结果，针对 'How are multi-attention heads plotted out in PyTorch?'：

--- 结果1（代码） ---
标题：plot_visualization_utils.py
来源：plot_visualization_utils.py
得分：0.3714
摘录：# models. Let's start by analyzing the output of a Mask-RCNN model. Note that...

--- 结果2（代码） ---
标题：plot_transforms_getting_started.py
来源：plot_transforms_getting_started.py
得分：0.3571
摘录：https://github.com/pytorch/vision/tree/main/gallery/...

已实现的功能

✅ 基本语义搜索：带有命令行界面的文档查询
✅ 向量数据库：ChromaDB集成用于存储和查询嵌入
✅ 内容区分：区分代码和文本内容
✅ 交互模式：运行连续交互式查询会话的选项

待改进之处

❌ 相关性质量：中等相似度得分（0.35-0.37）表明结果不够理想
❌ 内容覆盖范围：某些专业主题在数据库中的表示不足
❌ 分块策略：当前方法以任意点分割文档
❌ 结果呈现：摘录太短，缺乏足够上下文
❌ MCP集成：连接超时问题阻止Claude Code整合

入门指南

环境设置

创建包含所有依赖项的conda环境：

conda env create -f environment.yml
conda activate pytorch_docs_search

API 密钥设置

工具需要OpenAI API密钥用于嵌入生成：

export OPENAI_API_KEY=your_key_here

命令行用法

# 使用直接查询搜索
python scripts/search.py "your search query here"

# 运行交互模式
python scripts/search.py --interactive

# 其他选项
python scripts/search.py "query" --results 5  # 限制为 5 个结果
python scripts/search.py "query" --filter code  # 只显示代码结果
python scripts/search.py "query" --json  # 输出 JSON 格式

项目架构

• ptsearch/core/：核心搜索功能（数据库、嵌入、搜索）
• ptsearch/config/：配置管理
• ptsearch/utils/：实用函数和日志记录
• scripts/：命令行工具
• data/：嵌入文档和数据库
• ptsearch/protocol/：MCP协议处理（当前未使用）
• ptsearch/instructions/：MCP指令集（当前未使用）

为什么暂停？

暂停原因

• 数据质量：原始文档的质量不足以生成高质量的嵌入。
• 分块策略：需要更复杂的逻辑来处理技术文档中的内容。
• MCP集成：与PyTorch团队在接口和协议上尚未达成一致。

未来计划

计划改进

1. 数据质量

   • 使用预处理步骤清理原始文档。
   • 集成外部知识库以丰富上下文。

2. 分块策略

   • 开发更智能的逻辑来处理技术术语和代码片段。
   • 实现基于语义相似度的分块方法。

3. 结果呈现

   • 增强上下文信息，使搜索结果更具可读性。
   • 支持多格式输出（HTML、Markdown）。

4. MCP集成

   • 与PyTorch团队合作定义统一接口。
   • 实现双向通信机制以支持实时反馈。

开发指南

安装依赖项

使用以下命令安装所需库：

pip install -r requirements.txt

启动脚本

运行以下命令启动搜索工具：

python scripts/run.py

贡献指南

1. 代码规范

   • 遵循PEP8编码风格。
   • 提供清晰的文档注释。

2. 测试要求

   • 为每个功能模块编写单元测试。
   • 定期运行测试套件以确保代码质量。

3. 提交指南

   • 在分支中执行代码审查。
   • 提交时附带详细的变更日志。