What is Pytorch Docs Refactored

pytorch-docs-refactored is a semantic search tool designed for PyTorch documentation, integrating with MCP for Claude Code. It enables users to efficiently find relevant documentation, APIs, code examples, and error messages using advanced semantic search capabilities.

Use cases

Use cases include searching for specific PyTorch documentation, finding code examples, troubleshooting error messages, and enhancing the development workflow by quickly accessing relevant resources.

How to use

To use pytorch-docs-refactored, set up a conda environment, obtain an OpenAI API key, and integrate it with Claude Code via one of three methods: Direct STDIO, SSE, or UVX. After registration, users can query PyTorch-related questions directly.

Key features

Key features include semantic search for PyTorch documentation, code-aware search results, easy integration with Claude Code, multiple transport options (STDIO, SSE, UVX), and configurable search parameters and result formatting.

Where to use

pytorch-docs-refactored can be used in various fields such as machine learning, data science, and software development, particularly for those working with PyTorch for deep learning applications.

PyTorch Documentation Search Tool (Project Paused)

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

Current Status (April 19, 2025)

⚠️ This project is currently paused for significant redesign.

The tool provides a basic command-line search interface for PyTorch documentation but requires substantial improvements in several areas. While the core embedding and search functionality works at a basic level, both relevance quality and MCP integration require additional development.

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
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/...

What Works

✅ Basic Semantic Search: Command-line interface for querying PyTorch documentation
✅ Vector Database: Functional ChromaDB integration for storing and querying embeddings
✅ Content Differentiation: Distinguishes between code and text content
✅ Interactive Mode: Option to run continuous interactive queries in a session

What Needs Improvement

❌ Relevance Quality: Moderate similarity scores (0.35-0.37) indicate suboptimal results
❌ Content Coverage: Specialized topics may have insufficient representation in the database
❌ Chunking Strategy: Current approach breaks documentation at arbitrary points
❌ Result Presentation: Snippets are too short and lack sufficient context
❌ MCP Integration: Connection timeout issues prevent Claude Code integration

Getting Started

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

Command-line Usage

# 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

Project Architecture

• ptsearch/core/: Core search functionality (database, embedding, search)
• ptsearch/config/: Configuration management
• ptsearch/utils/: Utility functions and logging
• scripts/: Command-line tools
• data/: Embedded documentation and database
• ptsearch/protocol/: MCP protocol handling (currently unused)
• ptsearch/transport/: Transport implementations (STDIO, SSE) (currently unused)

Why This Project Is Paused

After evaluating the current implementation, we’ve identified several challenges that require significant redesign:

1. Data Quality Issues: The current embedding approach doesn’t capture semantic relationships between PyTorch concepts effectively enough. Relevance scores around 0.35-0.37 are too low for a quality user experience.

2. Chunking Limitations: Our current method divides documentation into chunks based on character count rather than conceptual boundaries, leading to fragmented results.

3. MCP Integration Problems: Despite multiple implementation approaches, we encountered persistent timeout issues when attempting to integrate with Claude Code:

   • STDIO integration failed at connection establishment
   • Flask server with SSE transport couldn’t maintain stable connections
   • UVX deployment experienced similar timeout issues

Future Roadmap

When development resumes, we plan to focus on:

1. Improved Chunking Strategy: Implement semantic chunking that preserves conceptual boundaries
2. Enhanced Result Formatting: Provide more context and better snippet selection
3. Expanded Documentation Coverage: Ensure comprehensive representation of all PyTorch topics
4. MCP Integration Redesign: Work with the Claude team to resolve timeout issues

Development

Running Tests

pytest -v tests/

Format Code

black .

License

MIT License