El Reto de Buscar en la Investigación Académica
Buscar artículos de investigación en IA no es lo mismo que buscar contenido web normal. Un motor de búsqueda robusto para investigación debe manejar consultas exactas, como un ID específico de arXiv, pero también necesita entender consultas conceptuales. Por ejemplo, un usuario que busca "small language models for code generation" espera resultados incluso si esas palabras exactas no aparecen en el texto de un artículo. También debe reconocer solicitudes de navegación como "the original BERT paper" y ser tolerante a errores tipográficos o títulos incompletos.
Para cumplir con estas demandas, el equipo detrás de Papers with Code no se basó en un solo método. Construyeron un sistema de búsqueda híbrida sofisticado. Una parte clave de este sistema se apoya en varias herramientas y mejores prácticas de Python para gestionar datos y cómputo.
Este post detalla la arquitectura que construyeron, enfocándose en las decisiones de diseño que la hacen poderosa y rápida.
Una Historia de Dos Métodos de Recuperación: Búsqueda por Palabras Clave y por Vectores
La idea central es que la búsqueda por palabras clave y la búsqueda por vectores son complementarias. La búsqueda por palabras clave (como la búsqueda de texto completo de PostgreSQL) es excelente para encontrar coincidencias exactas, identificadores y nombres raros. Por otro lado, la búsqueda por vectores utiliza embeddings densos para encontrar términos semánticamente similares, capturando el significado 'difuso' detrás de una consulta.
La Arquitectura: Offline vs. Online
El sistema está deliberadamente dividido en una construcción de corpus offline y un servicio de búsqueda online. Esta separación es clave para equilibrar el costo, el rendimiento y la confiabilidad.
- Construcción de Corpus Offline (Jobs): El trabajo costoso y orientado a la producción de incrustar todo el corpus de artículos (más de 110,000 artículos) se realiza como un trabajo por lotes en una GPU. Esto es impulsado por Hugging Face Jobs, que proporciona cómputo 'burstable'.
- Almacenamiento Duradero (Buckets): La salida de estos trabajos (los embeddings de vectores) se almacena en un Hugging Face Storage Bucket. Esto actúa como el tejido conectivo, asegurando una transferencia duradera entre el cómputo efímero y la base de datos de producción.
- Servicio de Búsqueda Online (Inference Endpoints): Solo el pequeño trabajo sensible a la latencia de incrustar una consulta de usuario en vivo está en la ruta de la solicitud. Esto es manejado por un Hugging Face Inference Endpoint, diseñado para respuestas de baja latencia.
Aquí hay una vista simplificada del proceso de incrustación de consultas:
# Pseudocódigo para el proceso de incrustación de consultas en línea
import requests
# 1. Envía la consulta del usuario al Inference Endpoint
response = requests.post(
"YOUR_INFERENCE_ENDPOINT_URL",
headers={"Authorization": "Bearer YOUR_TOKEN"},
json={"inputs": query_text}
)
# 2. Procesa la respuesta para obtener el vector de la consulta
query_vector = response.json()["embedding"]
# 3. Usa este vector para realizar una búsqueda de distancia de coseno en PostgreSQL (pgvector)
# SELECT paper_id, embedding <=> CAST(:query_vector AS halfvec(256)) AS distance
# FROM paper_embeddings
# WHERE generation_id = :active_generation
# ORDER BY embedding <=> CAST(:query_vector AS halfvec(256))
# LIMIT 50;
Patrones de Diseño Clave
- Contrato de Embedding Estricto: Tratan el formato de embedding como una API versionada. Esto incluye la revisión del modelo, las dimensiones de salida, el formato de entrada y el método de normalización. Esto evita fallas sutiles cuando los modelos se actualizan o el código se refactoriza.
- Trabajos Reanudables y Reproducibles: Al organizar los artefactos bajo prefijos de ejecución inmutables y usar sumas de verificación, pueden reintentar de manera segura trabajos fallidos o rastrear cualquier resultado hasta su fuente de datos y revisión de modelo exactas.
- Degradación Graceful: El sistema está diseñado para arranques en frío. Si el Inference Endpoint está escalando o se agota el tiempo, el cliente de consulta cae inmediatamente a la búsqueda léxica. Esto asegura que los usuarios siempre obtengan resultados, incluso si no son tan ricos semánticamente.
Lecciones Aprendidas y el Camino a Seguir
La experiencia del equipo ofrece lecciones valiosas para cualquiera que esté construyendo funciones de búsqueda similares con IA.
- Separa la Producción de la Latencia: La incrustación por lotes y la incrustación de consultas son problemas de infraestructura diferentes. Optimiza para costo y producción en el primero, y para disponibilidad y velocidad en el segundo.
- Haz que el Almacenamiento sea el Contrato: Usa un bucket de almacenamiento como un límite explícito y auditable entre tu cómputo y los sistemas de producción.
- Fija Más que el Nombre del Modelo: La revisión, la dimensión y el prompt afectan la calidad final de la recuperación. Almacena y valida todos ellos.
- Diseña para Arranques en Frío: Si usas escala a cero, debes tener un fallback rápido. La búsqueda híbrida proporciona esto de forma natural.
- Los Vectores Más Pequeños Pueden Ser una Característica: El Aprendizaje de Representación Matryoshka (MRL) te permite intercambiar calidad por velocidad y almacenamiento. En su piloto, 256 dimensiones preservaron la recuperación usando una fracción del almacenamiento.
- La Activación Debe Ser Aburrida: Las nuevas generaciones de embeddings deben importarse y validarse junto a la actual antes de activarse atómicamente. Esto hace que la reversión sea un simple cambio de configuración.
Limitaciones y Consideraciones
Aunque esta arquitectura es poderosa, es importante notar que la búsqueda híbrida no siempre es la mejor solución. El equipo recomienda comenzar con la búsqueda por palabras clave como una línea base barata y solo agregar búsqueda semántica y/o híbrida cuando veas un aumento razonable en la calidad de la recuperación. Agregar un re-ranker también puede mejorar los resultados, pero introduce latencia y sobrecarga adicionales.
Conclusión
El sistema de búsqueda híbrida que impulsa Papers with Code es una clase magistral en diseño de sistemas pragmático. Al separar las cargas de trabajo en línea y fuera de línea, usar un bucket de almacenamiento como contrato y diseñar para la degradación gradual, han construido un sistema que es poderoso y confiable. El uso de Jobs, Buckets e Inference Endpoints de Hugging Face proporciona un modelo para construir infraestructura de IA escalable.
Para aquellos que buscan optimizar aún más sus sistemas de recuperación, explorar las capacidades de los modelos más nuevos es un excelente próximo paso. También puedes leer sobre otros desarrollos emocionantes en el ecosistema de Python, como las nuevas características en el próximo lanzamiento de Python 3.15.