PyTorch Documentation Search Tool

A semantic search tool for PyTorch documentation with MCP integration for Claude Code.

Overview

This tool provides semantic search capabilities over PyTorch documentation, allowing users to find relevant documentation, APIs, code examples, and error messages. It utilizes vector embeddings and semantic similarity to provide high-quality search results.

Features

• Semantic search for PyTorch documentation
• Code-aware search results (differentiates between code and text)
• Easy integration with Claude Code via MCP
• Multiple transport options (STDIO, SSE, UVX)
• Configurable search parameters and result formatting

Installation

Environment Setup

Create a conda environment with all dependencies:

conda env create -f environment.yml
conda activate pytorch_docs_search

For a minimal environment:

conda env create -f minimal_env.yml
conda activate pytorch_docs_search_min

API Key Setup

The tool requires an OpenAI API key for generating embeddings:

export OPENAI_API_KEY=your_key_here

MCP Integration

The tool can be integrated with Claude Code in three ways:

1. Direct STDIO Integration (Local Development)

# Register with Claude CLI
./register_mcp.sh

# This runs:
# claude mcp add search_pytorch_docs stdio ./run_mcp.sh

2. SSE Integration (Server Deployment)

# Start the server
python -m ptsearch.server --transport sse --host 0.0.0.0 --port 5000

# Register with Claude CLI
claude mcp add search_pytorch_docs http://localhost:5000/events --transport sse

3. UVX Integration (Packaged Distribution)

# Run with UVX
./run_mcp_uvx.sh

# This executes:
# uvx mcp-server-pytorch --transport sse --host 127.0.0.1 --port 5000 --data-dir ./data

Usage

Once registered with Claude Code, you can use the tool by asking questions about PyTorch:

How do I implement a custom dataset in PyTorch?

Claude Code will automatically use the PyTorch Documentation Search Tool to find relevant documentation.

Direct CLI Usage

You can also use the tool directly:

# Search from command line
python -m ptsearch.server --transport stdio --data-dir ./data

Architecture

• ptsearch/server.py: Unified server implementation
• ptsearch/protocol/: MCP protocol handling
• ptsearch/transport/: Transport implementations (STDIO, SSE)
• ptsearch/core/: Core search functionality

Development

Running Tests

pytest -v tests/

Format Code

black .

License

MIT License