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

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

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

⚠️ 该项目目前因重大重新设计而暂停。

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

示例输出

$ python scripts/search.py "How are multi-attention heads plotted out in PyTorch?"

Found 5 results for 'How are multi-attention heads plotted out in PyTorch?':

--- Result 1 (code) ---
Title: plot_visualization_utils.py
Source: plot_visualization_utils.py
Score: 0.3714
Snippet: # models. Let's start by analyzing the output of a Mask-RCNN model. Note that...

--- Result 2 (code) ---
Title: plot_transforms_getting_started.py
Source: plot_transforms_getting_started.py
Score: 0.3571
Snippet: https://github.com/pytorch/vision/tree/main/gallery/...

已实现的功能

✅ 基本语义搜索：命令行界面用于查询 PyTorch 文档
✅ 向量数据库：集成了 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

命令行用法

# Search with a direct query
python scripts/search.py "your search query here"

# Run in interactive mode
python scripts/search.py --interactive

# Additional options
python scripts/search.py "query" --results 5  # Limit to 5 results
python scripts/search.py "query" --filter code  # Only code results
python scripts/search.py "query" --json  # Output in JSON format

项目架构

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

为什么暂停此项目

在评估当前实现后，我们发现了几个需要重大重新设计的挑战：

1. 数据质量问题：当前的嵌入方法不能有效地捕捉 PyTorch 概念之间的语义关系。相关性分数在 0.35-0.37 左右对于高质量用户体验来说太低了。

2. 分块限制：我们当前的方法是基于字符数而不是概念边界来划分文档，导致结果碎片化。

3. MCP 集成问题：尽管尝试了多种实现方法，我们在试图与 Claude Code 集成时遇到了持续的超时问题：

   • STDIO 集成在建立连接时失败
   • 使用 SSE 传输的 Flask 服务器无法维持稳定的连接
   • UVX 部署也遇到了类似的超时问题

未来路线图

当开发恢复时，我们计划专注于：

1. 改进的分块策略：实现保留概念边界的语义分块
2. 增强的结果格式化：提供更多上下文和更好的代码片段选择
3. 扩展文档覆盖范围：确保涵盖所有 PyTorch 主题
4. MCP 集成重新设计：与 Claude 团队合作解决超时问题

开发

运行测试

pytest -v tests/

格式化代码

black .

许可证

MIT 许可证