Herramienta de búsqueda de documentación de PyTorch

Herramienta de búsqueda de documentación de PyTorch (proyecto en pausa)

Un prototipo de búsqueda semántica para la documentación de PyTorch con capacidades de línea de comandos.

Estado actual (19 de abril de 2025)

⚠️ Este proyecto se encuentra actualmente en pausa por un rediseño significativo.

La herramienta proporciona una interfaz básica de búsqueda por línea de comandos para la documentación de PyTorch, pero requiere mejoras sustanciales en varias áreas. Si bien las funciones principales de incrustación y búsqueda funcionan a un nivel básico, tanto la calidad de la relevancia como la integración con MCP requieren desarrollo adicional.

Ejemplo de salida

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

Lo que funciona

✅ Búsqueda semántica básica : interfaz de línea de comandos para consultar la documentación de PyTorch
✅ Base de datos vectorial : Integración funcional de ChromaDB para almacenar y consultar incrustaciones
✅ Diferenciación de contenido : distingue entre código y contenido de texto.
✅ Modo interactivo : Opción para ejecutar consultas interactivas continuas en una sesión

Qué necesita mejorar

❌ Calidad de relevancia : los puntajes de similitud moderados (0,35-0,37) indican resultados subóptimos
❌ Cobertura de contenido : Los temas especializados pueden tener una representación insuficiente en la base de datos.
❌ Estrategia de fragmentación : el enfoque actual interrumpe la documentación en puntos arbitrarios
❌ Presentación de resultados : los fragmentos son demasiado cortos y carecen de contexto suficiente
❌ Integración de MCP : Los problemas de tiempo de espera de conexión impiden la integración de Claude Code

Empezando

Configuración del entorno

Cree un entorno conda con todas las dependencias:

conda env create -f environment.yml
conda activate pytorch_docs_search

Configuración de la clave API

La herramienta requiere una clave API de OpenAI para la generación de incrustaciones:

export OPENAI_API_KEY=your_key_here

Uso de la línea de comandos

# 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

Arquitectura del proyecto

ptsearch/core/ : Funcionalidad de búsqueda principal (base de datos, incrustación, búsqueda)
ptsearch/config/ : Gestión de configuración
ptsearch/utils/ : Funciones de utilidad y registro
scripts/ : Herramientas de línea de comandos
data/ : Documentación y base de datos integradas
ptsearch/protocol/ : manejo del protocolo MCP (actualmente sin uso)
ptsearch/transport/ : Implementaciones de transporte (STDIO, SSE) (actualmente sin uso)

¿Por qué este proyecto está en pausa?

Después de evaluar la implementación actual, hemos identificado varios desafíos que requieren un rediseño significativo:

Problemas de calidad de los datos : El enfoque de integración actual no captura las relaciones semánticas entre los conceptos de PyTorch con la suficiente eficacia. Las puntuaciones de relevancia de entre 0,35 y 0,37 son demasiado bajas para una experiencia de usuario de calidad.
Limitaciones de la fragmentación : nuestro método actual divide la documentación en fragmentos según el número de caracteres en lugar de límites conceptuales, lo que genera resultados fragmentados.
Problemas de integración de MCP : a pesar de los múltiples enfoques de implementación, encontramos problemas de tiempo de espera persistentes al intentar integrar con Claude Code:
- La integración de STDIO falló al establecer la conexión
- El servidor Flask con transporte SSE no pudo mantener conexiones estables
- La implementación de UVX experimentó problemas de tiempo de espera similares

Hoja de ruta futura

Cuando se reanude el desarrollo, planeamos centrarnos en:

Estrategia de fragmentación mejorada : implementar una fragmentación semántica que preserve los límites conceptuales
Formato de resultados mejorado : proporciona más contexto y una mejor selección de fragmentos
Cobertura de documentación ampliada : garantiza una representación completa de todos los temas de PyTorch
Rediseño de la integración de MCP : trabajar con el equipo de Claude para resolver problemas de tiempo de espera

Desarrollo

Ejecución de pruebas

pytest -v tests/

Código de formato

black .

Licencia

Licencia MIT

PyTorch Documentation Search Tool

Integrations