Pytorch Documentation Search

🚀 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 集成用於存儲和查詢嵌入
✅ 內容區分：區分代碼和文本內容
✅ 交互模式：運行連續交互式查詢會話的選項

📚 詳細文檔

當前狀態（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/...

待改進之處

❌ 相關性質量：中等相似度得分（0.35 - 0.37）表明結果不夠理想
❌ 內容覆蓋範圍：某些專業主題在數據庫中的表示不足
❌ 分塊策略：當前方法以任意點分割文檔
❌ 結果呈現：摘錄太短，缺乏足夠上下文
❌ MCP 集成：連接超時問題阻止 Claude Code 整合

項目架構

• 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. 提交指南
   • 在分支中執行代碼審查。
   • 提交時附帶詳細的變更日誌。