Мы с Крисом недавно «с нуля» буквально за пару часов создали механизм поиска для моего блога. Основную часть проделал именно Крис, так как до этого с word2vec я был знаком лишь отдалённо.

Разработанный нами поисковик основывается на векторных представлениях (эмбеддингах) слов. Принцип здесь следующий. Функция получает слово и отображает его в N-мерное пространство (в данном случае N=300), где каждое измерение отражает определённый оттенок смысла. Вот хорошая статья (англ.) о том, как обучить собственную модель word2vec, и её внутреннем устройстве.

Суть работы созданного нами поиска заключается в преобразовании моих статей, а точнее, составляющих их слов, в эмбеддинги, сохраняемые в общем пространстве. Затем при выполнении конкретного поиска текст его запроса преобразуется аналогичным образом и сопоставляется с векторами статей. В результате этого сопоставления, используя метрику косинусного сходства, мы ранжируем статьи по их релевантности запросу.

Уравнение ниже может показаться пугающим, но в нём говорится, что косинусное сходство, представляющее косинус угла между двух векторов cos(theta), определяется в виде скалярного произведения, поделённого на произведение величин каждого вектора. Разберём всё это подробнее.

Косинусное расстояние является, пожалуй, простейшим методом сравнения эмбеддинга запроса с эмбеддингами документов для их ранжирования. Ещё одним интуитивным подходом будет метрика евклидова расстояния, которая вместо угла между векторами измеряет, насколько они удалены друг от друга.

Мы отдаём предпочтение косинусному расстоянию, так как оно соответствует нашему интуитивному представлению, что два вектора имеют одинаковое значение, если пропорции их измерений совпадают. То есть при наличии двух векторов, устремлённых в одном направлении, можно сделать вывод, что они несут одинаковый смысл, даже если один из них очень короткий, а второй очень длинный. (Например, в двух документах, выраженных такими векторами, речь идёт о кошках, просто в одном слово «кошка» встречается гораздо чаще).

Теперь запустим word2vec и векторизуем наши первые слова.

Создание эмбеддингов

Мы взяли готовую базу данных из 10 000 самых популярных эмбеддингов слов, хранящихся в файле формата pickle размером 12 МБ. Вот её фрагмент:

couch  [0.23, 0.05, ..., 0.10]
banana [0.01, 0.80, ..., 0.20]
...

Крис прислал мне её по сети. Если эту базу данных развернуть, то мы получим структуру NumPy: словарь, в котором строки сопоставляются с массивами numpy.float32. Я написал скрипт для преобразования этого файла в простые числа с плавающей запятой и списки, так как хотел сделать всё вручную.

Код загрузки весьма прост — просто используем библиотеку pickle. Здесь присутствуют стандартные риски безопасности, но я доверяю Крису.

import pickle

def load_data(path):
    with open(path, "rb") as f:
        return pickle.load(f)

 word2vec = load_data("word2vec.pkl")

При желании можете вывести word2vec, но вывод получится огромным, имейте в виду. Я познал это на собственном печальном опыте… Предлагаю вывести просто word2vec["cat"]. В выводе вы получите соответствующий эмбеддинг.

Для векторизации слова достаточно найти его в необъятном словаре. Хотя бессмысленного или нетипичного слова может в нём не быть, в случае чего мы вместо ошибки вернём None.

def embed_word(word2vec, word):
    return word2vec.get(word)

Для векторизации нескольких слов мы будем обрабатывать их по отдельности, после чего объединим полученные эмбеддинги попарно. Если какое-то слово нельзя преобразовать в вектор, игнорируем его. Проблема возникает только в том случае, когда мы не можем понять ни одно слово.

def vec_add(a, b):
    return [x + y for x, y in zip(a, b)]

def embed_words(word2vec, words):
    result = [0.0] * len(next(iter(word2vec.values())))
    num_known = 0
    for word in words:
        embedding = word2vec.get(word)
        if embedding is not None:
            result = vec_add(result, embedding)
            num_known += 1
    if not num_known:
        raise SyntaxError(f"Не понимаю ничего из {words}")
    return result

В этом вся суть создания эмбеддингов — выполнение поиска по словарю и сложение векторов.

embed_words([a, b]) == vec_add(embed_word(a), embed_word(b))

А теперь создадим «поисковый индекс», а точнее, эмбеддинги для всех моих статей.

Создание эмбеддингов для всех статей

Векторизация всех статей подразумевает рекурсивный обход каталогов, в ходе которого создаётся словарь сопоставлений путей файлов с эмбеддингами этих файлов.

import os
def load_post(pathname):
    with open(pathname, "r") as f:
        contents = f.read()
    return normalize_text(contents).split()
  
def load_posts():
    # Обходим _posts в поиске файлов *.md.
    posts = {}
    for root, dirs, files in os.walk("_posts"):
        for file in files:
            if file.endswith(".md"):
                pathname = os.path.join(root, file)
                posts[pathname] = load_post(pathname)
    return posts

post_embeddings = {pathname: embed_words(word2vec, words)
                   for pathname, words in posts.items()}

with open("post_embeddings.pkl", "wb") as f:
    pickle.dump(post_embeddings, f)

Но здесь мы делаем ещё кое-что: normalize_text. Всё дело в том, что статьи содержат всяческую пунктуацию, заглавные буквы и другие сбивающие с толку элементы. Нам же для нахождения лучшего совпадения нужно рассматривать слова вроде «CoMpIlEr» и «compiler» как одинаковые.

import re

def normalize_text(text):
    return re.sub(r"[^a-zA-Z]", r" ", text).lower()

Всё это мы также проделываем для каждого запроса. Далее создадим простенькую REPL для поиска.

Скромная среда REPL для поиска

Для создания REPL-среды используем встроенный в Python модуль code. Напишем подкласс, определяющий метод runsource. Он будет просто обрабатывать ввод source и возвращать ложное значение (в противном случае ожидать дополнительный ввод).

import code

class SearchRepl(code.InteractiveConsole):
    def init(self, word2vec, post_embeddings):
        super().__init__()
        self.word2vec = word2vec
        self.post_embeddings = post_embeddings

    def runsource(self, source, filename="<input>", symbol="single"):
        for result in self.search(source):
            print(result)

Далее определим функцию search, которая объединит все остальные функции. Вуаля, и наш поиск готов:

class SearchRepl(code.InteractiveConsole):
    # ...
    def search(self, query_text, n=5):
        # Векторизуем запрос.
        words = normalize_text(query_text).split()
        try:
            query_embedding = embed_words(self.word2vec, words)
        except SyntaxError as e:
            print(e)
            return
        # Вычисляем косинусное сходство
        post_ranks = {pathname: vec_cosine_similarity(query_embedding,
                                                      embedding) for pathname,
                      embedding in self.post_embeddings.items()}
        posts_by_rank = sorted(post_ranks.items(),
                               reverse=True,
                               key=lambda entry: entry[1])
        top_n_posts_by_rank = posts_by_rank[:n]
        return [path for path,  in topn_posts_by_rank]

Да, нужно вычислять косинусное сходство. И здесь нам повезло, так как взятый из Wikipedia фрагмент практически один в один переводится в код Python:

import math
 
def vec_norm(v):
    return math.sqrt(sum([x*x for x in v]))

def vec_cosine_similarity(a, b):
    assert len(a) == len(b)
    a_norm = vec_norm(a)
    b_norm = vec_norm(b)
    dot_product = sum([ax*bx for ax, bx in zip(a, b)])
    return dot_product/(a_norm*b_norm)

Наконец, создаём и запускаем REPL.

sys.ps1 = "QUERY. "
sys.ps2 = "...... "

repl = SearchRepl(word2vec, post_embeddings)
repl.interact(banner="", exitmsg="")

Так выглядит процесс взаимодействия:

QUERY. type inference
_posts/2024-10-15-type-inference.md
_posts/2025-03-10-lattice-bitset.md
_posts/2025-02-24-sctp.md
_posts/2022-11-07-inline-caches-in-skybison.md
_posts/2021-01-14-inline-caching.md
QUERY.

Это образец запроса к небольшому датасету (моему блогу). Здесь мы получили неплохой результат, но он не отражает общее качество поиска. Крис говорит, что мне следует выбирать лучшие ответы, потому что «в сфере ИИ все так делают».

Хорошо, с этим разобрались. Но большинство людей, которые ищут что-либо на моём сайте, не используют терминал. Несмотря на то, что мой блог заточен под удобство чтения из текстовых браузеров вроде Lynx, многие читают его из графических. Значит, нужно сделать для поиска фронтенд.

Веб-поиск в миниатюре

К этому моменту мы всё выполняли на моей локальной машине, где меня нисколько не напрягает присутствие файла весов размером 12 МБ. Теперь же мы переходим в веб-среду, и мне не хочется утруждать рядовой браузер неожиданно большим объёмом загрузки. Здесь нужно грамотное решение.

К счастью, мы с Крисом оба читали эту крутую статью (англ.) о хостинге базы данных SQLite на GitHub Pages. В ней автор рассказывает, как он:

• скомпилировал SQLite в Wasm, чтобы она выполнялась на клиенте,

• собрал виртуальную файловую систему, чтобы она могла считывать файлы БД из сети,

• реализовал продуманное получение страниц, используя существующие индексы SQLite,

• создал дополнительное ПО для получения только небольших фрагментов базы данных с помощью HTTP-запросов Range.

Да, это очень круто, но SQLite вопреки своей миниатюрности для нашего проекта всё же великовата. Мы хотим создать всё с нуля и, к счастью, можем эмулировать все её основные принципы.

Мы можем упорядочить word2vec и разбить его на два файла. Один файл будет содержать просто эмбеддинги слов, без имён. Во втором будет храниться индекс, где каждое слово будет соотноситься с адресом смещения начального байта своих весов и их длиной (предполагается, что при передаче по сети данные о начале и длине весов займут меньший объём, нежели о начале и конце).

# vecs.jsonl
[0.23, 0.05, ..., 0.10]
[0.01, 0.80, ..., 0.20]
...

# index.json
{"couch": [0, 20], "banana": [20, 30], ...}

Хорошо здесь то, что index.json намного меньше блоб-объекта word2vec и весит всего 244 КБ. Поскольку меняться он будет нечасто (как часто вообще меняется word2vec?), я не вижу ничего страшного в том, что пользователям потребуется скачать весь индекс. То же касается post_embeddings.json, который весит всего 388 КБ. Эти файлы можно даже кэшировать. При этом они автоматически сжимаются (и разжимаются) сервером и браузером до размера 84 КБ и 140 КБ соответственно. Они станут ещё меньше, если выбрать двоичный формат, но в текущей статье мы этот нюанс опустим.

Затем можно создать HTTP-запросы Range к серверу и скачивать только те части весов, которые нам нужны. Можно даже связать все диапазоны (ranges) в один запрос (составной диапазон). К сожалению, GitHub Pages не поддерживает составной вариант, поэтому диапазон каждого слова мы будем скачивать в отдельном запросе.

Вот соответствующий JS-код, где (короткие и очень знакомые) функции получения векторов опущены:

(async function() {
  // Скачиваем данные.
  async function get_index() {
    const req = await fetch("index.json");
    return req.json();
  }
  async function get_post_embeddings() {
    const req = await fetch("post_embeddings.json");
    return req.json();
  }
  const index = new Map(Object.entries(await get_index()));
  const post_embeddings = new Map(Object.entries(await get_post_embeddings()));
  // Добавляем обработчик поиска.
  search.addEventListener("input", debounce(async function(value) {
    const query = search.value;
    // TODO(max): нормализация запроса
    const words = query.split(/\s+/);
    if (words.length === 0) {
      // Без слов.
      return;
    }
    const requests = words.reduce((acc, word) => {
      const entry = index.get(word);
      if (entry === undefined) {
        // Недопустимое слово; пропускаем.
        return acc;
      }
      const [start, length] = entry;
      const end = start+length-1;
      acc.push(fetch("vecs.jsonl", {
        headers: new Headers({
          "Range": bytes=${start}-${end},
        }),
      }));
      return acc;
    }, []);
    if (requests.length === 0) {
      // Допустимых слов нет :(
      search_results.innerHTML = "No results :(";
      return;
    }
    const responses = await Promise.all(requests);
    const embeddings = await Promise.all(responses.map(r => r.json()));
    const query_embedding = embeddings.reduce((acc, e) => vec_add(acc, e));
    const post_ranks = {};
    for (const [path, embedding] of post_embeddings) {
      post_ranks[path] = vec_cosine_similarity(embedding, query_embedding);
    }
    const sorted_ranks = Object.entries(post_ranks).sort(function(a, b) {
      // Уменьшаем.
      return b[1]-a[1];
    });
    // Занятный факт: HTML-элементы с атрибутом ‘id’ доступны в виде глобальных объектов JS под тем же именем.
    search_results.innerHTML = "";
    for (let i = 0; i < 5; i++) {
      search_results.innerHTML += <li>${sorted_ranks[i][0]}</li>;
    }
  }));
})();

Зацените действующую страницу поиска. А лучше вдобавок откройте в консоли браузера вкладку Network и проанализируйте загрузку. Прекрасно, если скачивается всего пара фрагментов эмбеддингов по 4 КБ.

Итак, насколько хорошо работает наша технология поиска? Давайте попробуем реализовать механизм объективной оценки.

Оценка

Давайте оценим наш механизм поиска и посмотрим, как часто он возвращает статьи из топа результатов поиска, когда мы используем в запросе собственные ключевые слова.

Начнём со сбора и анализа датасета пар (document, query). Мы изначально исказим этот анализ, собрав датасет самостоятельно. Но я надеюсь, что это поможет нам лучше понять качество поиска. Запросом в данном случае будет просто несколько выражений, которые, на наш взгляд, должны вести к успешной выдаче документа.

sample_documents = {
  "_posts/2024-10-27-on-the-universal-relation.md": "database relation universal tuple function",
  "_posts/2024-08-25-precedence-printing.md": "operator precedence pretty print parenthesis",
  "_posts/2019-03-11-understanding-the-100-prisoners-problem.md": "probability strategy game visualization simulation",
  # ...
}

Теперь, когда датасет собран, реализуем метрику оценки точности на основе top-k результатов. Эта метрика отражает процент случаев, когда документ появляется в верхних k результатах выдачи при соответствующем запросе.

def compute_top_k_accuracy(
    # Сопоставление статьи с образцом поискового запроса (уже нормализовано).
    # Смотрите sample_documents выше.
    eval_set: dict[str, str],
    max_n_keywords: int,
    max_top_k: int,
    n_query_samples: int,
) -> list[list[float]]:
    counts = [[0] * max_top_k for  in range(maxn_keywords)]
    for n_keywords in range(1, max_n_keywords + 1):
        for post_id, keywords_str in eval_set.items():
            for  in range(nquery_samples):
                # Построение поискового запроса путём выборки ключевых слов.
                keywords = keywords_str.split(" ")
                sampled_keywords = random.choices(keywords, k=n_keywords)
                query = " ".join(sampled_keywords)

                # Определяем ранг целевого поста в результатах поиска.
                ids = search(query, n=max_top_k)
                rank = safe_index(ids, post_id)

                # Инкрементируем позицию документа.
                if rank is not None and rank < max_top_k:
                    counts[n_keywords - 1][rank] += 1
    accuracies = [[0.0] * max_top_k for  in range(maxn_keywords)]
    for i in range(max_n_keywords):
        for j in range(max_top_k):
            # Делим на количество образцов для получения среднего,
            # затем делим на размер оцениваемого набора для получения точности по всем постам.
            accuracies[i][j] = counts[i][j] / n_query_samples / len(eval_set)

            # Аккумулируем показатели точности, так как если полученная позиция равна i,
            # значит документ также был успешно извлечён во всех позициях j > i .
            if j > 0:
                accuracies[i][j] += accuracies[i][j - 1]

     return accuracies

Ниже показан график, отражающий точность top-k для различных значений k. Обратите внимание, что при увеличении k точность возрастает — когда это значение становится близко количеству документов, точность также приближается к 100%. Кроме того, с ростом количества ключевых слов точность также увеличивается, и их линии начинают сходиться, что говорит об убывающей отдаче при добавлении каждого нового слова.

Действительно ли эти мегабайты эмбеддингов слов действительно как-то улучшают наш поиск? Нужно провести сравнение с базовыми показателями. Возможно, при вычислении этих показателей для ранжирования документов суммируется количество всех ключевых слов в каждом из них. Но эту проверку мы оставим вам, дорогие читатели, так как наше время на этот проект истекло ?.

Также будет интересно увидеть, насколько увеличение word2vec способствует увеличению точности. При выборке top-k результатов встречается много ошибок (Не могу понять ничего из ['prank', ...]). Эти неизвестные слова из поиска исключаются. Более обширная база данных word2vec (содержащая >10 000 слов) может включать эти менее распространённые слова, а значит, обеспечивать более эффективный поиск.

Подытожим

Вы вполне можете создать простой механизм поиска «с нуля», обойдясь всего сотней строк кода. Ознакомьтесь с полноценной программой в search.py, которая включает дополнительные фрагменты для оценки и построения графика.

Идеи на будущее

Можно выйти за пределы простой оценки косинусного сходства. Представим, что все наши документы из сферы компьютерных технологий, но лишь в одном из них обсуждаются компиляторы (какая жалость). Если запрос будет содержать слово «компьютер», то это не особо поможет сузить область поиска, и в наших эмбеддингах такое слово станет, скорее, шумом. Для уменьшения шумности мы используем технику под названием TF-IDF (частота термина и обратная частота документа), которая позволит исключить из анализа наиболее типичные слова и сосредоточить внимание на более уникальных для каждого документа терминах.