AI导哈

PyTorch文档语义搜索工具

一个用于PyTorch文档的语义搜索原型工具，提供命令行搜索功能，支持交互式查询和结果过滤。

  • by seanmichaelmcgee
  • Python
  • 本地部署
  • 2 星标
  • 更新于
在 GitHub 查看 seanmichaelmcgee/pytorch-docs-refactored

使用说明

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许可证

相关 MCP