JinDaGe - PyTorch Documentation Search MCP Details

article

README

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

这是一个具备命令行功能的 PyTorch 文档语义搜索原型，能够帮助用户在 PyTorch 文档中进行语义查询，但目前因重大重构而暂停开发。

🚀 快速开始

环境设置

创建包含所有依赖项的 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 格式

✨ 主要特性

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

📦 安装指南

安装依赖项

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

pip install -r requirements.txt

启动脚本

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

python scripts/run.py

💻 使用示例

基础用法

$ 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/...

📚 详细文档

项目架构

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

为什么暂停？

暂停原因

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

未来计划

计划改进

数据质量
- 使用预处理步骤清理原始文档。
- 集成外部知识库以丰富上下文。
分块策略
- 开发更智能的逻辑来处理技术术语和代码片段。
- 实现基于语义相似度的分块方法。
结果呈现
- 增强上下文信息，使搜索结果更具可读性。
- 支持多格式输出（HTML、Markdown）。
MCP 集成
- 与 PyTorch 团队合作定义统一接口。
- 实现双向通信机制以支持实时反馈。

开发指南

贡献指南

代码规范
- 遵循 PEP8 编码风格。
- 提供清晰的文档注释。
测试要求
- 为每个功能模块编写单元测试。
- 定期运行测试套件以确保代码质量。
提交指南
- 在分支中执行代码审查。
- 提交时附带详细的变更日志。

🔧 技术细节

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

📄 许可证

文档未提及相关信息，故跳过该章节。

⚠️ 重要提示

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

💡 使用建议

使用前请确保已正确设置环境和 API 密钥。

PyTorch Documentation Search