Pytorch Documentation Search

🚀 PyTorch Documentation Search Tool (Project on Hold)

A semantic search prototype for PyTorch documentation with command-line functionality.

🚧 Current Status (April 19, 2025)

⚠️ Important Note

This project is currently on hold due to a major refactoring.

This tool provides a basic command-line search interface for querying PyTorch documentation, but it requires substantial improvements in multiple aspects. Although the core embedding and search functions work at a basic level, further development is needed in terms of relevance and MCP integration.

📄 Example Output

$ 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
Excerpt: # 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
Excerpt: https://github.com/pytorch/vision/tree/main/gallery/...

✨ Features

✅ Basic Semantic Search: Document query with a command-line interface
✅ Vector Database: ChromaDB integration for storing and querying embeddings
✅ Content Differentiation: Differentiate between code and text content
✅ Interactive Mode: Option to run continuous interactive query sessions

🛠️ Areas for Improvement

❌ Relevance Quality: Medium similarity scores (0.35 - 0.37) indicate suboptimal results
❌ Content Coverage: Some specialized topics are underrepresented in the database
❌ Chunking Strategy: The current method splits documents at arbitrary points
❌ Result Presentation: Excerpts are too short and lack sufficient context
❌ MCP Integration: Connection timeout issues prevent Claude Code integration

🚀 Quick Start

🌐 Environment Setup

Create a conda environment with all dependencies:

conda env create -f environment.yml
conda activate pytorch_docs_search

🔑 API Key Setup

The tool requires an OpenAI API key for embedding generation:

export OPENAI_API_KEY=your_key_here

💻 Usage Examples

📋 Basic Usage

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

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

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

🏗️ Project Architecture

• ptsearch/core/: Core search functions (database, embedding, search)
• ptsearch/config/: Configuration management
• ptsearch/utils/: Utility functions and logging
• scripts/: Command-line tools
• data/: Embedded documents and database
• ptsearch/protocol/: MCP protocol handling (currently unused)
• ptsearch/instructions/: MCP instruction set (currently unused)

❓ Why on Hold?

🛑 Reasons for Pause

• Data Quality: The quality of the original documents is insufficient to generate high-quality embeddings.
• Chunking Strategy: More sophisticated logic is needed to handle content in technical documents.
• MCP Integration: Agreement on interfaces and protocols with the PyTorch team has not been reached.

🚀 Future Plans

📈 Planned Improvements

1. Data Quality

   • Clean the original documents using preprocessing steps.
   • Integrate external knowledge bases to enrich the context.

2. Chunking Strategy

   • Develop more intelligent logic to handle technical terms and code snippets.
   • Implement a chunking method based on semantic similarity.

3. Result Presentation

   • Enhance context information to make search results more readable.
   • Support multi-format output (HTML, Markdown).

4. MCP Integration

   • Collaborate with the PyTorch team to define a unified interface.
   • Implement a two-way communication mechanism to support real-time feedback.

🛠️ Development Guide

📦 Install Dependencies

Install the required libraries using the following command:

pip install -r requirements.txt

🚀 Start the Script

Run the following command to start the search tool:

python scripts/run.py

🤝 Contribution Guidelines

1. Code Style

   • Follow the PEP8 coding style.
   • Provide clear documentation comments.

2. Testing Requirements

   • Write unit tests for each functional module.
   • Run the test suite regularly to ensure code quality.

3. Submission Guidelines

   • Conduct code reviews in branches.
   • Attach detailed change logs when submitting.