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_searchAPI密钥设置
该工具需要一个用于嵌入生成的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）（目前未使用）
为什么这个项目被暂停
在评估了当前的实施后，我们发现了几个需要重新设计的挑战：
- 数据质量问题：当前的嵌入方法不能有效地捕获PyTorch概念之间的语义关系。相关性得分在0.35-0.37左右，对于高质量的用户体验来说太低了。
- 分块限制：我们目前的方法根据字符数而不是概念边界将文档划分为块，导致结果支离破碎。
- MCP集成问题：尽管有多种实现方法，但在尝试与Claude Code集成时，我们遇到了持续的超时问题：
- STDIO集成在连接建立时失败 - 具有SSE传输的Flask服务器无法保持稳定的连接 - UVX部署也遇到了类似的超时问题
未来路线图
当开发恢复时，我们计划重点关注：
- 改进的分块策略：实现保留概念边界的语义组块
- 增强的结果格式：提供更多上下文和更好的代码段选择
- 扩大文档覆盖范围：确保全面呈现所有PyTorch主题
- MCP集成重新设计：与Claude团队合作解决超时问题
发展
运行测试
pytest -v tests/格式码
black .许可证
MIT许可证