Tìm kiếm hàng triệu trang Wikipedia ngoại tuyến bằng ‘vibe’

 (leebutterman.com)
1 điểm bởi GN⁺ 2023-09-03 | 1 bình luận | Chia sẻ qua WhatsApp
  • Triển khai tìm kiếm embedding thời gian thực ngoại tuyến ngay trong trình duyệt mà không cần máy chủ, cho phép khám phá tương tác cục bộ 6 triệu tài liệu English Wikipedia
  • Embedding gốc 384 chiều dạng float32 có dung lượng khoảng 9GB nên không phù hợp với ứng dụng trình duyệt; vì vậy kích thước được giảm dựa trên all-minilm-l6-v2 và embedding trung bình của các đoạn trang
  • Áp dụng Product Quantization theo các không gian con 8 chiều để giảm bộ dữ liệu xuống mức 288MB, và có thể chứa khoảng 2 triệu embedding trong 96MB
  • Việc tìm kiếm thực hiện phép tính khoảng cách ngay trong miền nén mà không giải nén, và dùng opcode GatherElements của ONNX để sửa kết quả chuyển đổi, giúp tăng tốc tính toán khoảng 4 lần
  • Trên thiết bị mới, suy luận MiniLM và tính khoảng cách cho 100k embedding hoàn tất trong thời gian ngắn, cho phép giao diện React cập nhật kết quả theo chu kỳ 100~300ms

Tìm kiếm vector ngoại tuyến chạy ngay trong trình duyệt

  • Mục tiêu là chứng minh rằng có thể tìm kiếm thời gian thực hàng triệu trang Wikipedia ngay trong trình duyệt mà không cần hạ tầng tìm kiếm máy chủ phức tạp
  • Thiết kế cốt lõi được cô đọng thành ba ý
    • Tìm kiếm vector có thể đủ nhanh chỉ với Product Quantization và quét tuyến tính
    • Dùng ONNX Runtime, nhưng những điểm kém hiệu quả phát sinh từ quá trình chuyển đổi PyTorch có thể được sửa ở cấp opcode
    • Trình duyệt đủ nhanh để chạy suy luận thời gian thực bằng WASM, và WebGPU vẫn là một hướng tiềm năng trong tương lai
  • Điều kiện tìm kiếm kết hợp độ tương đồng vector của embedding nội dung với điều kiện bằng nhau của chữ cái đầu tiên trong tiêu đề
    • Điều này gần với trường hợp truy vấn cơ sở dữ liệu, trong đó một số cột được tìm theo độ tương đồng vector còn các cột khác theo điều kiện bằng nhau
  • Do ràng buộc là ứng dụng trình duyệt ngoại tuyến, embedding phải nhỏ, mô hình embedding phải là mã nguồn mở, đồng thời kích thước mô hình và lượng tính toán cũng phải nhẹ
    • Theo tiêu chuẩn năm 2023, mô hình nhẹ thường có nghĩa là dưới 100MB

Dữ liệu English Wikipedia và kích thước embedding

  • Bộ dữ liệu là English Wikipedia, được embedding bằng mô hình all-minilm-l6-v2
  • Có khoảng 6 triệu tài liệu, và sau khi chia trang thành các đoạn nhỏ thì lấy trung bình embedding của từng đoạn để tạo embedding cấp trang
  • Các trang được sắp xếp từ dài đến ngắn trước
    • Mục đích là khi tải cơ sở dữ liệu dần dần, các trang dài hơn — có khả năng đã được con người đầu tư công sức nhiều hơn — sẽ được hiển thị trước
  • Chiều embedding là 384 chiều
    • Nếu lưu dưới dạng float32 thì 6M * 384 * 4 = 9GB
    • Trong 96MB chỉ chứa được khoảng 64k embedding
  • Cách mã hóa float32 gốc quá lớn để phù hợp với tìm kiếm ngoại tuyến trên trình duyệt

Nén embedding bằng Product Quantization

  • Product Quantization là cách không lưu trực tiếp giá trị số thực dấu phẩy động, mà tạo nhiều bảng mã/palette rồi thay giá trị đầu vào bằng chỉ số của palette
  • Palette thường có không quá 256 giá trị nên mỗi chỉ số có thể chỉ tốn tối đa 1 byte
    • Palette ngầm định dùng một trục số như -127~127 hoặc 0~255 rồi dịch chuyển và co giãn để sử dụng
    • Palette tường minh lưu lần lượt 256 con số
    • Product Quantization sử dụng kiểu palette tường minh
  • Với cấu hình tối thiểu, nếu chuyển từng chiều trong 384 chiều thành chỉ số 1 byte thì bộ dữ liệu sẽ là 6M * 384 = 2.25GB
    • Trong 96MB có thể lưu khoảng 256k embedding
    • Cách này vẫn chưa đủ hiệu quả so với mục tiêu
  • Nếu lượng tử hóa theo từng điểm 2 chiều thì embedding 384 chiều có thể được biểu diễn bằng 192 chỉ số, giúp tiết kiệm không gian gấp 2 lần
  • Cấu hình thực tế dùng đơn vị điểm 8 chiều
    • 384 chiều được biểu diễn bằng 48 chỉ số
    • Kích thước palette được giữ ở mức khoảng 384KB
    • Kích thước bộ dữ liệu là 6M * 48 = 288M
    • Trong 96MB có thể chứa khoảng 2 triệu embedding
  • Ở bất kỳ mức lượng tử hóa nào cũng có thể đánh giá precision và recall

Tìm kiếm trực tiếp mà không giải nén

  • Ưu điểm của Product Quantization là có thể tìm kiếm ngay khi dữ liệu vẫn đang nén mà không cần giải nén
  • Nếu gom n chiều thành n/2 điểm 2 chiều, có thể tính trước khoảng cách giữa từng điểm trong palette và điểm 2 chiều tương ứng của truy vấn
  • Sau đó, mỗi embedding chỉ cần tra các giá trị khoảng cách bằng chỉ số palette rồi cộng lại để ra khoảng cách tổng
    • Cách này ít việc hơn so với việc khôi phục toàn bộ điểm n chiều rồi tính khoảng cách trên từng chiều
  • Trong cấu hình điểm 8 chiều thực tế, phép tính khoảng cách chỉ cần khối lượng công việc ở mức 1/8 so với khi xử lý embedding không nén

Dùng Arrow làm định dạng trao đổi dữ liệu

  • Vì chi phí di chuyển và chuyển đổi dữ liệu lớn, nên sẽ có lợi hơn nếu tuần tự hóa embedding nén theo định dạng có thể dùng ngay lập tức
  • Arrow là định dạng hướng cột phù hợp cho mục đích này
    • Embedding và tiêu đề trang được lưu dưới dạng bảng Arrow
    • Chỉ số palette được xử lý như mảng 2 chiều, còn tiêu đề như mảng chuỗi 1 chiều
    • Có thể tập trung vào việc sao chép bit về đúng vị trí hơn là phân tích cú pháp hay nạp dữ liệu
  • Định dạng mảng Arrow chỉ lưu dữ liệu 1 chiều, nên để xử lý 48 chiều embedding cần dùng hai schema
    • Schema metadata có đơn vị 100 nghìn hàng
    • Schema embedding có đơn vị 100 nghìn * 48 hàng
    • Khi nạp sẽ reshape embedding trở lại
  • safetensors cũng theo cùng nguyên tắc thiết kế
  • JSON tuần tự hóa mảng 48 phần tử thành chuỗi ASCII độ dài biến thiên, khiến giai đoạn nạp trở nên phức tạp
  • Protocol Buffers dùng số nguyên độ rộng biến thiên base 128 cho các số nguyên nhỏ hơn 32 bit, và định dạng số nguyên này hiện không được compute kernel hỗ trợ tốt

Phân tích Wikipedia và tạo embedding

  • Để phân tích cú pháp markup của Wikipedia, tác giả dùng mediawiki parser from hell, tức mwparserfromhell
  • Wikipedia dataset của OLM lấy bản dump dữ liệu mới nhất của Wikipedia theo từng ngôn ngữ rồi phân tích thành các hàng gồm tiêu đề và nội dung
    • Bộ dữ liệu này giống mã cần được thực thi hơn là dữ liệu tuần tự hóa vô hại, nên cần có sự tin cậy khi chạy
  • Việc phân tích được thực hiện trên tất cả các lõi của một máy
    • Máy hiện đại có hàng chục lõi, nên phù hợp với hệ thống dưới 100 triệu tài liệu như English Wikipedia hiện tại ở quy mô 6 triệu trang
  • Mô hình embedding dùng all-minilm-l6-v2
    • Đây là sentence transformer nhẹ nhưng có hiệu năng tốt cho tiếng Anh
    • Truy vấn và tài liệu được embedding vào cùng một không gian
    • Mô hình rất nhẹ với 22 triệu tham số
  • Để chạy JavaScript trong trình duyệt, tác giả cũng tận dụng phiên bản lượng tử hóa 8 bit thông qua ONNX và transformers.js
  • Mô hình được huấn luyện với chuỗi 128 token, nên cửa sổ ngữ cảnh ngắn hơn rất nhiều so với độ dài trang trung bình
    • Vì vậy trang được chia thành các đoạn nhỏ và lấy trung bình embedding của từng đoạn để tạo embedding trung bình cấp trang

Tìm kiếm facet dựa trên quét tuyến tính với pq.js

  • Khi tìm kiếm cục bộ trên hàng triệu tài liệu, không nhất thiết phải cần hệ thống lập chỉ mục phức tạp
  • Mục tiêu là lấy ra vài kết quả gần nhất với một điểm cụ thể, ví dụ 12 khoảng cách nhỏ nhất
  • Khi có 10 triệu phần tử khoảng cách, 10 triệu giá trị cột facet và một giá trị facet, có thể lọc bằng cách cộng 0 nếu khớp điều kiện, ngược lại cộng Infinity, rồi tìm top-k
    • Trên điện thoại hiện đại, có thể tìm 100 kết quả đầu trong dưới 10ms
    • Cài đặt nằm trong filtered-topk và pq.js
  • Việc tính khoảng cách cho embedding nén bằng Product Quantization có thể được viết bằng PyTorch
    • Có subspaceCount palette, và mỗi palette có codewordCount điểm subspaceDim chiều
  • Chỉ mục trong mô hình ONNX được xuất từ PyTorch có dạng hơi vụng về
    • Opcode GatherElements của ONNX thực hiện trực tiếp thao tác cần thiết
    • Có thể dùng công cụ như ONNX-modifier để thêm hoặc xóa các nút trong đồ thị luồng dữ liệu của mô hình ONNX đã xuất
  • Khi thay nhiều bước đánh chỉ mục bằng đúng một opcode duy nhất, việc tính khoảng cách nhanh hơn khoảng 4 lần

Tính toán streaming và giao diện tương tác

  • Số lượng embedding cần tính khoảng cách không cố định
    • Có thể tính streaming khoảng cách giữa truy vấn và một phần embedding
    • Sau khi đủ thời gian, chạy top-k trên các khoảng cách đã tính đến thời điểm đó để cập nhật kết quả tìm kiếm
  • Vì toàn bộ embedding đều nằm cục bộ, độ trễ tới hạ tầng tìm kiếm về cơ bản là 0ms
  • Mức phản hồi giao diện mục tiêu là hiển thị kết quả trong 100~300ms sau thao tác
    • MiniLM có thể chạy trong dưới 100ms
    • Trên một số thiết bị edge còn có thể dưới 15ms
    • Tính khoảng cách cho 100k embedding có thể chạy trong khoảng 10ms
  • Nếu sắp xếp bộ dữ liệu từ bài dài đến bài ngắn, phần lớn kết quả tìm kiếm top cuối cùng sẽ xuất hiện sớm trong kết quả streaming ban đầu
    • Vì các bài dài thường nhận được nhiều sự chú ý và công sức hơn khi viết, đồng thời thường khớp truy vấn tốt hơn so với vô số bài stub
  • Sau khi tính khoảng cách toàn cục, chỉ cần áp dụng facet filter và top-k; vì thế khi thay đổi giá trị facet hay số lượng kết quả trên cùng một truy vấn, chỉ cần chạy lại bước lọc
    • Bước lọc này dưới 10ms nên cho cảm giác gần như tức thì

Các thành phần pq.js có thể tái sử dụng

  • Nhiều hàm thư viện trong toàn bộ ứng dụng tìm kiếm Wikipedia có thể được chuyển thành các thành phần pq.js có thể tái sử dụng
  • Hiện tại nhiều ONNX shape vẫn được cố định sẵn
  • Nếu hỗ trợ nhiều mức lượng tử hóa và nhiều chiều embedding khác nhau, khả năng tái sử dụng sẽ rộng hơn

Bài viết liên quan

1 bình luận

 
GN⁺ 2023-09-03
Các ý kiến trên Hacker News

  • Chắc chắn là thú vị, nhưng khi thử mô tả vài thuật ngữ triết học và tâm lý học, tất cả các mục tôi tìm đều chỉ xuất hiện quanh vị trí thứ 20
    Các mục nổi tiếng hơn nhưng kém chính xác hơn nằm phía trên; ví dụ dù nhập gì để định nghĩa một trường phái trị liệu tâm lý cụ thể, “psychotherapy” luôn đứng số 1
    Ngược lại, tôi từng dùng ChatGPT để tìm một phân nhánh hẹp mà mình không nhớ tên, và lần nào nó cũng đoán đúng
    Ý tưởng về một dịch vụ AI tìm tên sự vật bằng mô tả thì hay, nhưng tôi không chắc cách giới hạn vào Wikipedia hay tiêu đề bài Wikipedia có đúng không; các mô hình ngôn ngữ lớn đa dụng có vẻ đã làm khá tốt rồi
    Dù vậy, đây là một proof of concept và việc có thể chạy cục bộ trong trình duyệt thật sự rất tuyệt

    • Mục tiêu là trình diễn công cụ cơ sở dữ liệu, và demo rằng sau khi trình duyệt tải mọi thứ xuống, nó vẫn có thể hoạt động trong một môi trường hoàn toàn cô lập
      Có vẻ còn nhiều tham số để tinh chỉnh. Chẳng hạn dùng chỉ đoạn đầu của bài hay toàn bộ bài, hoặc tìm trong phạm vi gần với một bài cụ thể; hiện vẫn chưa động đến các phần đó
      Wikipedia là một bộ dữ liệu demo tuyệt vời, và tôi cũng muốn thử thêm các bộ dữ liệu khác. Ví dụ như khi tìm “mountain” trong iPhoto thì ảnh có núi hiện ra; sẽ thú vị nếu tìm kiếm nhiều bộ dữ liệu bằng mô hình đa phương thức như CLIP
    • Tôi từng dùng theo kiểu mô tả ý tưởng và mục tiêu rồi nhận lại tên lĩnh vực hoặc các từ khóa đáng tra cứu
      Mô hình ngôn ngữ lớn trông như công cụ tìm kiếm mờ tốt nhất, hoạt động theo một cách khá độc đáo nhưng bổ trợ cho các công cụ tìm kiếm truyền thống

  • Tôi thích ý tưởng này, nhưng kết quả không tốt lắm
    Tôi nhập “weird looking monkey” và kỳ vọng các kết quả như khỉ vòi hay khỉ mũi hếch vàng, nhưng chỉ nhận được các bài như “Pet monkey”, “List of individual monkeys”, “Ethnoprimatology”, “Monkey”
    Khi nhập cùng truy vấn vào Google thì ra đúng kết quả mong đợi, nên hơi tiếc; tôi muốn tìm hiểu những loài khỉ trông kỳ lạ mà mình chưa biết

    • Đây chỉ là dùng nguyên một mô hình sentence-transformer có sẵn đã hơn 1 năm tuổi
      Trọng tâm của demo là trình diễn cơ sở dữ liệu embedding, nhưng bản thân embedding cũng hữu ích phần nào
      Vì không lưu lại bất kỳ dữ liệu phân tích nào về việc người dùng tìm được hay không tìm được gì trên trang, nên hiện chưa sẵn sàng để cải thiện kết quả tìm kiếm
    • Các biên tập viên và hướng dẫn của Wikipedia nhìn chung không thích tính từ mang tính chủ quan, và “weird looking” trong truy vấn nghe như một cách diễn đạt mà bài viết Wikipedia sẽ tránh dùng
      Vì vậy sẽ khó lấy được kết quả tốt từ kho tri thức này

  • Phần triển khai thật sự tốt, và việc có thể làm chuyện này ngoại tuyến rất tuyệt. Tuy nhiên chất lượng embedding có vẻ vẫn còn thiếu
    Một mẹo có thể hữu ích là không embedding toàn bộ bài Wikipedia, mà chỉ embedding các câu gần với định nghĩa, hoặc thường là câu đầu/đoạn đầu. Tôi không rõ hiện đang dùng phần nào
    Trang OneLook của tôi cũng đã cung cấp một tính năng tương tự để tìm từ và khái niệm bằng mô tả từ năm 2003 tại https://onelook.com/thesaurus/
    Ban đầu đó chỉ là tìm kiếm từ điển đảo thuần túy, nhưng trong 20 năm qua chúng tôi đã thử word embedding, sentence embedding, và gần đây là cả mô hình ngôn ngữ lớn; hiện GPT tạo các ứng viên cho những đầu vào mà hệ thống của chúng tôi không tự trả lời được
    Với tác vụ này, mô hình ngôn ngữ lớn vượt xa các phương pháp cũ, nên tôi có phần giảm động lực cải thiện mảng này của OneLook. Tôi thường thấy người nói rằng tìm kiếm định nghĩa ngược là lý do chính họ dùng ChatGPT

  • Tôi thấy hơi muộn, nhưng text embedding, ít nhất là theo những thứ dùng trong bài này, thường không giỏi lắm trong việc tìm kiếm theo cảm nhận/bầu không khí
    Nhìn chung chúng giống với việc so sánh các từ trùng nhau hoặc tìm nội dung tương tự truy vấn hơn
    Tuy nhiên có một bài báo gần đây thực sự thử giải quyết vấn đề này: “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
    Bài báo có nhiều ví dụ tìm kiếm bằng mô tả trừu tượng như “kiến trúc sư thiết kế một tòa nhà”, “một công ty là một phần của công ty khác”, “cuốn sách có ảnh hưởng đến sự phát triển của một thể loại”
    Embedding đó có vẻ hỗ trợ các kiểu tìm kiếm này tốt hơn nhiều, nên sẽ rất thú vị nếu thử lại tìm kiếm Wikipedia ngoại tuyến trong bài được liên kết bằng loại embedding mới này

  • Hiện trang không hoạt động trong môi trường của tôi, vì model_quantized.onnx không được tải
    Ngay trong lúc nhập, nó đã tải xuống đến 19,2MB với tốc độ khoảng 50KB/s; nếu mỗi khách truy cập đều gây ra việc này thì có thể đang gây ra chuyện khủng khiếp cho hóa đơn băng thông của Lee Butterman

    • Tôi đang phục vụ các tệp tĩnh từ một t2.nano, nên cũng không biết sẽ ra sao

  • Bản thân việc nó làm rất ấn tượng, nhưng chất lượng kết quả tìm kiếm có vẻ không tốt
    Theo kinh nghiệm, tôi biết việc đánh giá thủ công chất lượng kết quả tìm kiếm thực sự rất khó. Nó có thể gần chạm tới những kết quả rất tốt, nhưng vẫn trả về các kết quả khớp tệ hơn rất nhiều

    • Đúng vậy. Chất lượng có lẽ không tốt bằng Similar Website Finder https://explore2.marginalia.nu/ đâu ;)
      Dùng sentence embedding mới hơn có lẽ sẽ cải thiện kết quả, và tôi cần thu thập thêm dữ liệu

  • Công nghệ thì rất ấn tượng, nhưng kết quả thì không
    Tôi tìm “pointy building in Paris” thì ra Tourism in Paris, Bourse de commerce (Paris), Grands Projets of François Mitterrand, List of tallest buildings and structures in the Paris region, List of tourist attractions in Paris, Palais des congrès de Paris, Landmarks in Paris, Palais de la Bourse, Lyon, Outline of Paris, Architecture of Paris
    Hoàn toàn không thấy tòa nhà nhọn nổi tiếng nhất ở Paris
    Với kiểu ứng dụng như thế này, sentence embedding của toàn bộ tài liệu có thể không phải là cách tốt nhất

    • Ít nhất 5 trong số đó có vẻ chứa câu trả lời cần tìm
    • Nếu đang nói đến Eiffel Tower, thì nó không phải là building
      Tôi vừa kiểm tra bài viết: từ “building” xuất hiện 19 lần, nhưng phần lớn là động từ, sau đó là “Chrysler Building”
      Ý tôi là vậy, trừ khi có một tòa nhà nhọn nổi tiếng nào khác mà tôi không nghĩ ra

  • Một phần phép màu của công cụ tìm kiếm nằm ở việc trộn embedding của các trang liên kết đến trang đó, hoặc từ khóa truy xuất thông tin truyền thống, có trọng số theo số lượt nhấp và điểm uy tín
    Nếu không có tín hiệu này, rất nhiều thông tin hữu ích sẽ bị bỏ qua, và kết quả sẽ không còn cảm giác kỳ diệu nữa
    Dù vậy đây vẫn là một demo ấn tượng và thú vị

  • Tôi đã muốn thích nó, nhưng trong các tìm kiếm tôi thử, gần như không có kết quả liên quan
    “The wizard in The Lord of the Rings” không có Gandalf hay Saruman, chỉ có các sách liên quan đến LOTR
    “Protagonist of Scorsese's Taxi Driver” không có Travis Bickle
    “A person that plants trees for a living” thì không hiểu sao gardener không có trong danh sách
    “Curly-haired painter on TV” hoàn toàn không có Bob Ross
    “Unusually shaped modern art museum in Spain” thì Bilbao xuất hiện ở vị trí thứ 4, nhưng các kết quả còn lại không có hình dạng đặc biệt
    “Dog shaped like a sausage” thì có lẽ dachshund nên nằm trong các kết quả hàng đầu

    • Đáng nói là tất cả những kết quả mong muốn ở đây đều có bài viết Wikipedia
      Nếu không có bài viết thì việc bị thiếu sẽ đỡ lạ hơn, nhưng thực tế là tất cả đều tồn tại

  • “Vibes” là cách nói dễ cảm hơn nhiều so với “sentence embeddings”. Có lẽ tôi cũng nên bắt đầu dùng cách nói đó :)

    • Đây là cách diễn đạt hy sinh rất nhiều độ chính xác để đổi lấy “dễ cảm”
      Tác giả bài gốc không giải thích vì sao chọn từ đó, và nó không khớp với bất kỳ cách dùng nào của “vibe” mà tôi biết
      Có lẽ “gist” chưa đủ giống buzzword chăng