Cách Papers with Code tối ưu hóa tìm kiếm nhờ Hugging Face Inference Endpoints, Jobs và Buckets

3 tháng trước, chúng tôi bắt đầu công cuộc hồi sinh Papers with Code (xem thêm tweet thông báo). Mục tiêu là làm cho các nghiên cứu AI mã nguồn mở trở nên dễ tiếp cận và dễ hiểu hơn, giúp mọi người dễ dàng tìm thấy các tài nguyên liên quan đến một bài báo, tìm kiếm trạng thái kỹ thuật tốt nhất (SOTA) trong các lĩnh vực AI khác nhau, chia sẻ các nghiên cứu thú vị và phát triển dựa trên công trình của nhau. Nói cách khác, mục tiêu là thúc đẩy làn sóng nghiên cứu dẫn đến mô hình Transformer tiếp theo.

Tất nhiên, việc làm cho nghiên cứu AI trở nên dễ tiếp cận đòi hỏi một công cụ tìm kiếm mạnh mẽ, để con người và các tác nhân (agents) có thể nhanh chóng tìm thấy các công trình liên quan, thông qua trang web hoặc lệnh CLI tìm kiếm pwc mà các tác nhân có thể sử dụng thông qua Skill.

Cần lưu ý rằng tìm kiếm nghiên cứu không hoàn toàn giống với tìm kiếm văn bản thông thường. Một công cụ tìm kiếm bài báo hữu ích cần tìm được tiêu đề chính xác hoặc mã định danh arXiv, nhưng cũng phải hiểu được các truy vấn như “small language models for code generation” ngay cả khi các từ đó không xuất hiện cùng nhau trong một bài báo. Nó cần nhận diện được “the original BERT paper” là một yêu cầu điều hướng, chấp nhận tiêu đề không đầy đủ hoặc lỗi chính tả, và vẫn phản hồi nhanh chóng ngay cả khi dịch vụ mô hình đang ở trạng thái "nguội" (cold) hoặc tạm thời không khả dụng.

Đối với Papers with Code, chúng tôi đã xây dựng hệ thống này dưới dạng tìm kiếm lai (hybrid search). Điều này dựa trên kinh nghiệm trước đây của chúng tôi tại ML6, nơi chúng tôi đã phát triển các hệ thống dựa trên RAG cho khách hàng. Kết quả cho thấy tìm kiếm lai thường vượt trội hơn các hệ thống tìm kiếm dựa trên từ khóa và vector, vì nó kết hợp được ưu điểm của cả hai (xem thêm blog này để biết thêm thông tin). Tìm kiếm từ khóa giúp tìm các đề cập chính xác, trong khi tìm kiếm vector giúp tìm các thuật ngữ gần đúng và tương đồng về ngữ nghĩa. Lưu ý rằng các bộ sắp xếp lại (rerankers, còn gọi là cross-encoders) có thể cải thiện kết quả hơn nữa, mặc dù chúng cũng đi kèm với chi phí bổ sung và độ trễ.

Papers with Code dựa trên cơ sở dữ liệu PostgreSQL, do đó khả năng tìm kiếm toàn văn (full-text search) của nó cung cấp một nền tảng từ vựng nhanh chóng. Đối với các dense embedding, pgvector được sử dụng để bổ sung khả năng truy hồi ngữ nghĩa, và thuật toán reciprocal rank fusion (RRF) kết hợp cả hai. Ba dịch vụ Hugging Face được sử dụng cho các dense embedding:

Hiện nay, hệ thống duy trì embedding cho hơn 110.000 bài báo mới nhất được lấy từ arXiv và Daily Papers. Bài viết này giải thích về kiến trúc, các quyết định thiết kế đằng sau nó và những bài học chúng tôi rút ra khi đưa hệ thống vào vận hành thực tế.

Tóm tắt (TL;DR)

Chúng tôi chủ động tách biệt việc tìm kiếm thành hai phần: xây dựng kho dữ liệu ngoại tuyến (offline corpus build) và dịch vụ tìm kiếm trực tuyến (online search service):

Các công việc tốn kém, tập trung vào thông lượng (throughput) được chạy dưới dạng Jobs. Các tài nguyên bền vững nằm trong một Bucket. Chỉ bước nhúng truy vấn (query-embedding) nhỏ nằm trên đường dẫn yêu cầu, phía sau một Inference Endpoint được bảo vệ, để cung cấp năng lượng cho tìm kiếm trực tuyến. Nếu endpoint đó ở trạng thái nguội, bận hoặc không ổn định, tìm kiếm sẽ ngay lập tức chuyển sang truy hồi toàn văn. Sự tách biệt này giúp hệ thống vừa mạnh mẽ vừa nhanh chóng.

Bắt đầu với một hợp đồng embedding nghiêm ngặt

Các đường ống (pipeline) embedding thường thất bại theo những cách tinh vi: phiên bản mô hình thay đổi, các prompt truy vấn và tài liệu bị trộn lẫn, các vector bị cắt ngắn khác nhau, hoặc một bản tóm tắt được cập nhật không còn khớp với vector đã lưu trữ của nó.

Chúng tôi tránh điều này bằng cách coi định dạng embedding như một API có phiên bản. Mỗi bài báo được mã hóa dưới dạng:

Đối với mỗi lần tạo vector, chúng tôi ghi lại:

Quá trình tạo dữ liệu thực tế của chúng tôi sử dụng Qwen/Qwen3-Embedding-0.6B, được ghim vào một phiên bản chính xác, với các vector 256 chiều được chuẩn hóa L2. Lưu ý rằng các mô hình embedding mới hơn như Qwen3 cho phép 2 tính năng mới:

Hợp đồng này theo sát một embedding từ khâu xuất dữ liệu, qua suy luận GPU, vào PostgreSQL, và cuối cùng là truy hồi trực tuyến.

Jobs biến ảnh chụp nhanh cơ sở dữ liệu thành một kho vector

Việc nhúng toàn bộ kho dữ liệu là một khối lượng công việc hàng loạt (batch workload) điển hình. Nó cần GPU trong một khoảng thời gian tương đối ngắn, hưởng lợi từ thông lượng cao và không nên tiêu tốn tài nguyên giữa các lần chạy. Hugging Face Jobs phù hợp với hình thức này: một Job được định nghĩa bởi một lệnh, loại phần cứng và tùy chọn là một Docker image, và có thể chạy các tập lệnh uv với các phụ thuộc được khai báo trực tiếp.

Quá trình xây dựng kho dữ liệu của chúng tôi bắt đầu bằng việc xuất phiên bản mới nhất của mọi bài báo từ một ảnh chụp nhanh PostgreSQL (repeatable-read). Bộ xuất dữ liệu truyền các hàng thay vì tải toàn bộ danh mục vào bộ nhớ, ghi các tệp JSONL phân đoạn và tạo một bản kê khai chứa số lượng hàng và mã kiểm tra SHA-256.

Chúng tôi đồng bộ thư mục chạy bất biến đó vào một Storage Bucket riêng và gắn (mount) Bucket trực tiếp (sử dụng hf-mount) vào một Job l4x1 (NVIDIA L4 GPU, có 24GB VRAM). Từ góc độ của worker, nó chỉ đơn giản là một hệ thống tệp:

Worker:

Mỗi phân đoạn hoàn thành đều có dấu hiệu riêng, vì vậy một Job bị khởi động lại có thể bỏ qua các công việc đã xác minh. Điều này rất hữu ích cho một kho dữ liệu lớn: việc thử lại chỉ nên tiếp tục công việc thay vì ghi đè lên các embedding hiện có.

Trong chương trình thí điểm với 5.000 bài báo, Job Qwen đã mã hóa khoảng 75 bài báo mỗi giây ở 1024 chiều trên GPU L4. Cùng một lượt chạy đó có thể được hiện thực hóa một cách xác định ở 256 và 512 chiều, vì vậy chúng tôi có thể so sánh sự đánh đổi giữa lưu trữ và truy hồi mà không phải trả thêm chi phí cho suy luận.

Buckets là sợi dây liên kết

Storage Buckets là bộ lưu trữ đối tượng có thể thay đổi, giống S3 trên Hub, được tối ưu hóa cho các khối lượng công việc AI. Chúng có thể được truy cập thông qua các đường dẫn hf://buckets/... và được gắn quyền đọc-ghi trong các Jobs mà không cần xây dựng tích hợp lưu trữ riêng biệt.

Đối với chúng tôi, Bucket không chỉ là nơi đặt các vector. Nó là ranh giới giữa ba hệ thống với các vòng đời khác nhau:

Chúng tôi tổ chức các tài nguyên theo các tiền tố chạy bất biến:

Bản thân các Buckets có thể thay đổi, vì vậy tính bất biến là một quy tắc ở cấp độ ứng dụng: ID của một lần chạy không bao giờ bị ghi đè, và mọi tài nguyên đều được bao phủ bởi một bản kê khai và mã kiểm tra.

Điều này mang lại cho chúng tôi một số thuộc tính hữu ích:

Chỉ sau khi bộ nhập dữ liệu kiểm tra lại các lược đồ, mã kiểm tra, kích thước, chuẩn hóa, ID bài báo duy nhất và mã băm nội dung hiện tại, chúng tôi mới tải các vector vào PostgreSQL. Sau đó, chúng tôi xây dựng một chỉ mục HNSW riêng cho thế hệ mới và đánh dấu nó là hoạt động một cách nguyên tử chỉ khi mọi bài báo đủ điều kiện hiện tại đã được bao phủ (HNSW là thuật toán dựa trên đồ thị cho phép tìm kiếm vector nhanh chóng).

Inference Endpoints đưa tìm kiếm ngữ nghĩa vào đường dẫn yêu cầu

Các embedding hàng loạt giải quyết phía tài liệu của việc truy hồi. Một truy vấn của người dùng vẫn cần được nhúng tại thời điểm yêu cầu bằng cách sử dụng cùng một hợp đồng mô hình.

Chúng tôi triển khai mô hình đã ghim dưới dạng một Inference Endpoint được xác thực, được hỗ trợ bởi Text Embeddings Inference (TEI). Endpoint chấp nhận văn bản truy vấn và trả về một vector 256 chiều đã chuẩn hóa bằng cách sử dụng prompt truy vấn của mô hình. Lưu ý rằng người ta cũng có thể tận dụng vLLM hoặc SGLang ở đây.

Sau đó, API thực hiện tìm kiếm khoảng cách cosine trên thế hệ pgvector đang hoạt động:

Chỉ mục HNSW giữ cho việc tra cứu này nhanh chóng. Trong chương trình thí điểm 5.000 bài báo, chỉ mục Qwen 256 chiều đạt Recall@20 là 0,9955 so với tìm kiếm chính xác, với độ trễ tra cứu HNSW p50 là 1,31 ms và p95 là 2,21 ms. Bảng và chỉ mục của nó sử dụng khoảng 27% dung lượng lưu trữ so với phiên bản 1024 chiều trong khi vẫn giữ nguyên khả năng truy hồi ANN trong thử nghiệm đó.

Endpoint được cấu hình với tối đa một bản sao và có thể mở rộng về 0 khi nhàn rỗi. Đó là một đòn bẩy chi phí hữu ích, vì điều này có nghĩa là bạn không phải trả tiền khi không có người sử dụng. Tuy nhiên, điều này cũng có nghĩa là khởi động nguội (cold starts) phải là một phần của thiết kế ứng dụng thay vì được coi là một sự kiện ngoại lệ, vì cần một khoảng thời gian để endpoint khởi động và phục vụ lưu lượng truy cập.