ऑफ़लाइन में लाखों Wikipedia पेजों को ‘vibe’ से खोजें

 (leebutterman.com)
1 पॉइंट द्वारा GN⁺ 2023-09-03 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • बिना सर्वर के, ब्राउज़र के अंदर ऑफ़लाइन रीयल-टाइम embedding search लागू करके English Wikipedia के 60 लाख दस्तावेज़ों को लोकल में इंटरैक्टिव तरीके से एक्सप्लोर किया गया
  • 384-डायमेंशन float32 embeddings का मूल आकार लगभग 9GB है, इसलिए यह ब्राउज़र ऐप के लिए उपयुक्त नहीं था; इसे all-minilm-l6-v2 और page chunk के औसत embedding के आधार पर छोटा किया गया
  • Product Quantization को 8-डायमेंशन subspace इकाइयों पर लागू करके dataset को लगभग 288MB तक घटाया गया, और 96MB में लगभग 20 लाख embeddings रखे जा सकते हैं
  • खोज के दौरान decompression के बिना compressed domain में distance calculation किया जाता है, और ONNX के GatherElements opcode से conversion output को ठीक करके computation लगभग 4 गुना तेज़ बनाया गया
  • आधुनिक डिवाइसों पर MiniLM inference और 100k embeddings की distance calculation कम समय में पूरी हो जाती है, इसलिए React UI 100~300ms के अंतराल पर परिणाम अपडेट कर सकता है

ब्राउज़र के अंदर चलने वाली ऑफ़लाइन vector search

  • लक्ष्य यह दिखाना है कि जटिल server search infrastructure के बिना भी लाखों Wikipedia pages को ब्राउज़र में रीयल-टाइम खोजा जा सकता है
  • मुख्य डिज़ाइन तीन बिंदुओं में सिमटती है
    • vector search केवल Product Quantization और linear scan से भी काफ़ी तेज़ हो सकती है
    • ONNX Runtime का उपयोग करते हुए भी, PyTorch conversion से पैदा हुई inefficiency को opcode स्तर पर सुधारा जा सकता है
    • ब्राउज़र WASM के साथ रीयल-टाइम inference चलाने लायक तेज़ है, और WebGPU आगे की संभावना के रूप में मौजूद है
  • खोज की शर्तें body embedding की vector similarity और title के पहले अक्षर की equality condition, दोनों का साथ में उपयोग करती हैं
    • यह उस database query case के क़रीब है जहाँ कुछ columns vector similarity से और अन्य columns equality condition से खोजे जाते हैं
  • ऑफ़लाइन browser app की सीमाओं के कारण embeddings छोटे होने चाहिए, embedding model open source होना चाहिए, और model size व computation दोनों हल्के होने चाहिए
    • 2023 के मानदंड के अनुसार lightweight model का अर्थ आम तौर पर 100MB से कम होता है

English Wikipedia डेटा और embedding का आकार

  • dataset English Wikipedia है, और इसे all-minilm-l6-v2 मॉडल से embed किया गया है
  • दस्तावेज़ लगभग 60 लाख हैं, और pages को chunks में बाँटने के बाद हर chunk embedding का औसत लेकर page embedding बनाई जाती है
  • pages को पहले लंबे pages से क्रमबद्ध किया जाता है
    • database को progressively load करते समय, ऐसे लंबे pages पहले दिखाने के लिए जिन पर संभवतः लोगों ने ज़्यादा मेहनत की हो
  • embedding dimension 384 है
    • float32 में store करने पर 6M * 384 * 4 = 9GB
    • 96MB में केवल लगभग 64k embeddings ही रखी जा सकती हैं
  • मूल float32 encoding ब्राउज़र-आधारित ऑफ़लाइन खोज के लिए बहुत बड़ी है

Product Quantization से embedding compression

  • Product Quantization में floating-point values को सीधे store नहीं किया जाता; इसके बजाय कई palettes बनाई जाती हैं और input values को palette indices में बदला जाता है
  • palettes में आम तौर पर 256 या उससे कम values होती हैं, इसलिए हर index अधिकतम 1 byte का हो सकता है
    • implicit palette को -127~127 या 0~255 जैसी number line को shift और scale करके इस्तेमाल किया जाता है
    • explicit palette में 256 numbers को क्रम से store किया जाता है
    • Product Quantization explicit paletteization का उपयोग करता है
  • न्यूनतम configuration में 384 dimensions को 1-byte indices में बदलने पर dataset 6M * 384 = 2.25GB हो जाता है
    • 96MB में लगभग 256k embeddings store की जा सकती हैं
    • यह तरीका भी लक्ष्य की तुलना में अभी inefficent है
  • अगर quantization 2-डायमेंशन points पर की जाए, तो 384-डायमेंशन embedding को 192 indices से व्यक्त किया जा सकता है और 2x space saving मिलती है
  • वास्तविक configuration 8-डायमेंशन points का उपयोग करती है
    • 384 dimensions को 48 indices से व्यक्त किया जाता है
    • palette size लगभग 384KB के स्तर पर रहती है
    • dataset size 6M * 48 = 288M हो जाती है
    • 96MB में लगभग 20 लाख embeddings store की जा सकती हैं
  • किसी भी quantization स्तर पर precision और recall का मूल्यांकन किया जा सकता है

बिना decompress किए सीधे search

  • Product Quantization का फ़ायदा यह है that decompression के बिना भी compressed state में search किया जा सकता है
  • अगर n dimensions को n/2 के 2-डायमेंशन points में group किया जाए, तो हर palette point और query के संबंधित 2-डायमेंशन point के बीच distance पहले से compute की जा सकती है
  • इसके बाद हर embedding, palette indices के ज़रिए distance values को lookup करके जोड़ती है और कुल distance निकाली जाती है
    • यह पूरे n-डायमेंशन point को reconstruct करने के बाद हर dimension की distance निकालने से कम काम करता है
  • वास्तविक 8-डायमेंशन point configuration में distance calculation, uncompressed embeddings की तुलना में 1/8 काम से संभव है

डेटा interchange format के रूप में Arrow का उपयोग

  • data movement और conversion cost बड़ी होने के कारण compressed embeddings को ऐसे format में serialize करना फ़ायदेमंद है जिसे तुरंत उपयोग किया जा सके
  • Arrow इस तरह के उपयोग के लिए उपयुक्त column-oriented format है
    • embeddings और page titles को Arrow tables में store किया जाता है
    • palette indices को 2D array की तरह, और titles को 1D string array की तरह संभाला जाता है
    • parsing या loading से ज़्यादा ध्यान bits को सही जगह copy करने पर रखा जा सकता है
  • Arrow array format केवल 1D data store करता है, इसलिए embedding की 48 dimensions को संभालने के लिए दो schemas का उपयोग किया जाता है
    • metadata schema 1 लाख rows की इकाई में है
    • embedding schema 1 लाख * 48 rows की इकाई में है
    • load के समय embeddings को फिर से reshape किया जाता है
  • safetensors भी यही design principle अपनाता है
  • JSON, 48-element arrays को variable-length ASCII text के रूप में serialize करता है, जिससे loading चरण जटिल हो जाता है
  • Protocol Buffers, 32-bit से छोटे integers के लिए base-128 variable-width integers का उपयोग करता है, और यह integer format अभी compute kernel में अच्छी तरह supported नहीं है

Wikipedia parsing और embedding generation

  • Wikipedia markup parsing के लिए mediawiki parser from hell, यानी mwparserfromhell का उपयोग किया जाता है
  • OLM का Wikipedia dataset किसी विशेष भाषा की Wikipedia का नवीनतम data dump लाकर उसे title और body जैसी rows में parse करता है
    • यह dataset हानिरहित serialized data से ज़्यादा, ऐसे code के क़रीब है जिसे execute करना पड़ता है, इसलिए execution trust की ज़रूरत होती है
  • parsing एक मशीन के सभी cores पर की जाती है
    • आधुनिक मशीनों में दर्जनों cores होते हैं, इसलिए यह वर्तमान 60 लाख-page वाली English Wikipedia जैसे 10 करोड़ से कम दस्तावेज़ वाले सिस्टम के लिए उपयुक्त है
  • embedding model के रूप में all-minilm-l6-v2 का उपयोग किया जाता है
    • यह अंग्रेज़ी के लिए performant और lightweight sentence transformer है
    • यह query और documents को एक ही space में embed करता है
    • 2.2 करोड़ parameters के साथ यह बहुत हल्का model है
  • ब्राउज़र JavaScript execution के लिए ONNX और transformers.js के माध्यम से 8-bit quantized version का भी उपयोग किया जाता है
  • मॉडल 128-token sequences पर train किया गया था, इसलिए इसका context window औसत page length से काफ़ी छोटा है
    • pages को chunks में बाँटकर और हर chunk embedding का औसत लेकर page average embedding बनाई जाती है

pq.js और linear scan-आधारित facet search

  • जब स्थानीय रूप से लाखों documents खोजने हों, तो जटिल indexing system हमेशा ज़रूरी नहीं होता
  • लक्ष्य किसी दिए गए point के सबसे क़रीब कुछ top results, जैसे top 12 distances, प्राप्त करना है
  • जब 1 करोड़ distance values, 1 करोड़ facet columns और एक facet value हो, तो condition match होने पर 0 और match न होने पर Infinity जोड़कर top-k निकालने के तरीके से filtering की जाती है
    • आधुनिक फ़ोन पर top 100 को 10ms से कम में पाया जा सकता है
    • इसका implementation filtered-topk और pq.js में है
  • Product Quantization-compressed embeddings की distance calculation PyTorch में लिखी जा सकती है
    • subspaceCount palettes होती हैं, और हर palette में codewordCount संख्या के subspaceDim-डायमेंशन points होते हैं
  • PyTorch से export किए गए ONNX model की indexing कुछ awkward रूप ले लेती है
    • ONNX का GatherElements opcode आवश्यक काम सीधे करता है
    • ONNX-modifier जैसे tools से exported ONNX model के dataflow graph nodes जोड़े या हटाए जा सकते हैं
  • indexing के कई चरणों को सही single opcode से बदलने पर distance calculation लगभग 4 गुना तेज़ हो जाती है

streaming computation और interactive UI

  • distance calculation के target embeddings की संख्या fixed नहीं है
    • query और embeddings के एक हिस्से के बीच distances को streaming में compute किया जा सकता है
    • पर्याप्त समय बीतने पर अब तक निकली distances पर top-k चलाकर search results अपडेट किए जा सकते हैं
  • क्योंकि सभी embeddings लोकल में हैं, search infrastructure तक की latency लगभग 0ms है
  • UI responsiveness का लक्ष्य interaction के बाद 100~300ms के भीतर परिणाम दिखाना है
    • MiniLM को 100ms से कम में चलाया जा सकता है
    • edge devices पर कुछ मामलों में यह 15ms से भी कम हो सकता है
    • 100k embeddings की distance calculation लगभग 10ms में की जा सकती है
  • अगर dataset को लंबे लेखों से छोटे लेखों की ओर sort किया जाए, तो अंतिम top search results का अधिकांश हिस्सा शुरुआती streaming results में जल्दी दिखने लगता है
    • क्योंकि लंबे लेखों में अक्सर लिखते समय ज़्यादा ध्यान और मेहनत लगी होती है, और वे कई stub articles की तुलना में search query से बेहतर मेल खा सकते हैं
  • distances को globally compute करने के बाद facet filter और top-k लागू किए जाते हैं, इसलिए एक ही query पर facet value या search result count बदलने पर केवल filtering फिर से चलानी होती है
    • यह filtering 10ms से कम समय लेती है, इसलिए तुरंत महसूस होती है

पुन: उपयोग योग्य pq.js components

  • पूरे Wikipedia search app की कई library functions को reusable pq.js components में ले जाया जा सकता है
  • अभी कई ONNX shapes पहले से fixed हैं
  • अगर अलग-अलग quantization levels और विभिन्न embedding dimensions को support किया जाए, तो इसे और व्यापक रूप से reuse किया जा सकता है

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2023-09-03
Hacker News की राय

  • निश्चित रूप से दिलचस्प है, लेकिन दर्शन और मनोविज्ञान के कुछ शब्दों को समझाकर देखा तो जिन entries को खोज रहा था, वे सभी सिर्फ लगभग 20वें स्थान के आसपास ही दिखीं
    ज्यादा प्रसिद्ध लेकिन कम सटीक entries ऊपर भरी हुई थीं; उदाहरण के लिए किसी खास psychotherapy modality को define करने के लिए कुछ भी टाइप करूँ, “psychotherapy” हमेशा नंबर 1 था
    इसके उलट, मैंने ChatGPT से एक बार ऐसा संकीर्ण subfield खोजा था जिसका नाम याद नहीं आ रहा था, और उसने हर बार सही बताया
    description से चीजों के नाम खोजने वाली AI service का idea अच्छा है, लेकिन इसे Wikipedia या Wikipedia article titles तक सीमित रखना सही तरीका है या नहीं, पता नहीं; general-purpose large language models पहले से ही यह काफी अच्छी तरह करते दिखते हैं
    फिर भी, proof of concept के तौर पर और यह देखते हुए कि यह browser में locally चल सकता है, यह सचमुच शानदार है

    • लक्ष्य database engine दिखाना था, और यह demo करना था कि browser द्वारा सब कुछ download कर लेने के बाद यह पूरी तरह isolated environment में भी चल सकता है
      tune करने के लिए कई parameters दिखते हैं। जैसे article का सिर्फ पहला paragraph इस्तेमाल करना है या पूरा article, किसी खास article के करीब के दायरे में search करना है या नहीं—लेकिन इन्हें अभी छुआ नहीं गया है
      Wikipedia demo dataset के रूप में बढ़िया है, और दूसरे datasets भी जोड़कर देखना चाहूँगा। उदाहरण के लिए, जैसे iPhoto में “mountain” search करने पर पहाड़ वाली photos आती हैं, वैसे CLIP जैसे multimodal model से कई datasets search करना मजेदार होगा
    • मैंने इसे इस तरह इस्तेमाल किया है कि idea और goal समझाने पर यह field का नाम या खोजने लायक keywords लौटा दे
      large language models सबसे अच्छे fuzzy search engine जैसे लगते हैं, और traditional search engines से काफी अलग लेकिन complementary तरीके से काम करते हैं

  • concept पसंद आया, लेकिन results अच्छे नहीं आए
    “weird looking monkey” टाइप करके proboscis monkey या golden snub-nosed monkey जैसे results की उम्मीद थी, लेकिन “Pet monkey”, “List of individual monkeys”, “Ethnoprimatology”, “Monkey” जैसे articles ही आए
    वही query Google में डालने पर बिल्कुल उम्मीद के मुताबिक result मिला, इसलिए थोड़ा अफसोस हुआ; मैं ऐसे अजीब दिखने वाले बंदरों के बारे में जानना चाहता था जिनके बारे में मुझे पता नहीं था

    • एक साल से ज्यादा पुराना off-the-shelf sentence-transformer model ज्यों का त्यों इस्तेमाल किया गया है
      demo का focus embedding database दिखाने पर था, लेकिन embeddings खुद भी कुछ हद तक useful हैं
      page पर लोग क्या खोजते हैं और क्या नहीं मिल पाता, इसका कोई analytics data नहीं रखा गया, इसलिए search results सुधारने की तैयारी नहीं है
    • Wikipedia editors और guidelines आम तौर पर subjective adjectives पसंद नहीं करते, और query का “weird looking” Wikipedia articles में टालने लायक expression जैसा लगता है
      इसलिए इस knowledge corpus से अच्छे results लाना मुश्किल हो जाता है

  • implementation सचमुच अच्छी है, और offline यह कर पाना शानदार है। हालांकि embedding quality अभी कम लगती है
    एक मददगार trick यह हो सकती है कि पूरे Wikipedia article के बजाय definition जैसी sentences या आम तौर पर पहली sentence/पहला paragraph ही embed किया जाए। अभी कौन-सा हिस्सा इस्तेमाल हो रहा है, यह मुझे ठीक से नहीं पता
    मेरी site OneLook भी 2003 से https://onelook.com/thesaurus/ पर description से words और concepts खोजने की मिलती-जुलती सुविधा देती रही है
    शुरुआत में यह pure reverse dictionary search था, लेकिन पिछले 20 सालों में मैंने word embeddings, sentence embeddings और हाल में large language models तक आजमाए हैं; आजकल जिन inputs का जवाब system खुद नहीं दे पाता, उनके लिए GPT candidates generate करता है
    इस काम में large language models पुराने तरीकों से कहीं बेहतर हैं, इसलिए OneLook के इस हिस्से को सुधारने की मेरी motivation थोड़ी कम हुई है। अक्सर लोग कहते दिखते हैं कि reverse-definition search ही ChatGPT इस्तेमाल करने की उनकी मुख्य वजह है

  • थोड़ा देर से देखा, लेकिन text embeddings, कम से कम इस post में इस्तेमाल की गई चीजों के हिसाब से, आम तौर पर vibe से search करने में बहुत अच्छे नहीं होते
    वे ज्यादातर overlapping words compare करने या query जैसी content खोजने के करीब होते हैं
    हालांकि, इस problem को सचमुच tackle करने वाला एक हाल का paper है: “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
    paper में “buildings design करने वाला architect”, “किसी दूसरी company का हिस्सा company”, “genre के development पर असर डालने वाली book” जैसी abstract descriptions से search करने के कई examples हैं
    लगता है वह embedding ऐसी searches को कहीं बेहतर support करती है, इसलिए linked post वाली offline Wikipedia search को इस नए type की embedding से फिर करके देखना interesting होगा

  • अभी page मेरे environment में काम नहीं कर रहा है; model_quantized.onnx load नहीं हो रहा
    मैं type कर रहा था तब भी यह 19.2MB तक download हो चुका था और speed लगभग 50KB/s थी; अगर हर visitor के लिए ऐसा हो रहा है तो यह Lee Butterman के bandwidth bill के साथ भयानक चीजें कर रहा होगा

    • एक t2.nano से static files serve कर रहा हूँ, इसलिए पता नहीं क्या होगा

  • काम अपने-आप में बहुत impressive है, लेकिन search result quality अच्छी नहीं लगती
    experience से पता है कि search result quality को manually evaluate करना सचमुच मुश्किल होता है। यह बहुत अच्छे result के बेहद करीब होते हुए भी इससे कहीं खराब match लौटा सकता है

    • सही। quality शायद Similar Website Finder https://explore2.marginalia.nu/ जितनी अच्छी नहीं होगी ;)
      ज्यादा recent sentence embeddings इस्तेमाल करने से results बेहतर हो सकते हैं, और मुझे और data इकट्ठा करना होगा

  • तकनीक बहुत प्रभावशाली है, लेकिन नतीजे वैसे नहीं थे
    “pointy building in Paris” खोजने पर 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 आए
    पेरिस की सबसे मशहूर नुकीली इमारत बिल्कुल नहीं दिखी
    ऐसे application के लिए पूरे document के sentence embeddings शायद सबसे अच्छा तरीका न हों

    • उनमें से कम-से-कम 5 में तो शायद वह जवाब मौजूद होगा जिसे खोजा जा रहा है
    • अगर आप Eiffel Tower की बात कर रहे हैं, तो वह building नहीं है
      अभी document देखा; “building” शब्द 19 बार आता है, लेकिन ज़्यादातर जगह वह क्रिया है, और उसके बाद “Chrysler Building” है
      मतलब यह है कि जब तक कोई और मशहूर नुकीली इमारत नहीं है जो मेरे दिमाग में नहीं आ रही

  • search engine का कुछ जादू उन pages के embeddings, जो उस page को link करते हैं, या पारंपरिक information retrieval keywords को click count और authority score से weight करके मिलाने में होता है
    इस signal के बिना बहुत-सी उपयोगी जानकारी अनदेखी हो जाती है, और results जादुई नहीं लगते
    फिर भी यह प्रभावशाली और दिलचस्प demo है

  • मैं इसे पसंद करना चाहता था, लेकिन जिन searches को मैंने आज़माया उनमें relevant results बहुत कम निकले
    “The wizard in The Lord of the Rings” में Gandalf या Saruman नहीं थे, सिर्फ LOTR से जुड़ी किताबें आईं
    “Protagonist of Scorsese's Taxi Driver” में Travis Bickle नहीं था
    “A person that plants trees for a living” में न जाने क्यों gardener list में नहीं था
    “Curly-haired painter on TV” में Bob Ross बिल्कुल नहीं था
    “Unusually shaped modern art museum in Spain” में Bilbao चौथे स्थान पर तो आता है, लेकिन बाकी unusual shape वाले नहीं थे
    “Dog shaped like a sausage” के लिए dachshund top results में होना चाहिए लगता है

    • यह बात नोट करने लायक है कि यहाँ जिन results की उम्मीद थी, उन सभी के Wikipedia articles मौजूद हैं
      अगर articles नहीं होते तो उनका छूटना कम अजीब लगता, लेकिन असल में वे सब मौजूद हैं

  • “Vibes”, “sentence embeddings” की तुलना में कहीं ज़्यादा महसूस होने वाला expression है। शायद मुझे भी यह expression इस्तेमाल करना शुरू कर देना चाहिए :)

    • “महसूस होने वाला” होने के बजाय यह precision की बड़ी कुर्बानी देने वाला expression है
      मूल पोस्ट लिखने वाले ने यह नहीं समझाया कि उसने वह शब्द क्यों चुना, और “vibe” के जिन usages को मैं जानता हूँ उनमें से किसी से भी यह मेल नहीं खाता
      शायद “gist” पर्याप्त buzzword जैसा नहीं था