семантический браузер, использующий глубокое обучение и эластичный поиск для поиска документов COVID

Сегодня мы собираемся создать семантический браузер с использованием глубокого обучения для поиска в более чем 50 тысячах статей о недавней болезни COVID-19.

Весь код находится в моем репо на GitHub. Пока живая версия этой статьи находится здесь

Ключевой идеей является кодирование каждой статьи в векторе, представляющем ее семантическое содержание, а затем поиск с использованием косинусного сходства между запросом и всеми закодированными документами. Это тот же процесс, который используют браузеры изображений (например, Google Images) для поиска похожих изображений.

Итак, наша головоломка состоит из трех частей: данных, отображения документов на векторы и способа поиска.

Большая часть работы основана на этом проекте, в котором я работаю со студентами Университета Триеста (Италия). Живая демонстрация доступна здесь.

Давайте начнем!

Данные

Все начинается с данных. Мы будем использовать этот набор данных от Kaggle. Список из более чем 57 000 научных статей, подготовленных Белым домом и коалицией ведущих исследовательских групп. Фактически, нам нужен только файл metadata.csv, содержащий информацию о статьях и полный текст аннотации. Вам нужно сохранить файл внутри ./dataset.

Давайте взглянем

import pandas as pd
df = pd.read_csv('./dataset/metadata.csv')
df.head(5)

Как видите, информации у нас много. Нас явно интересуют текстовые столбцы. Работа с пандами не идеальна, поэтому давайте создадим Dataset. Это позволит нам позже создать DataLoader для выполнения пакетного кодирования. Если вы не знакомы с экосистемой загрузки данных Pytorch, вы можете узнать больше о здесь

Чтобы создать собственный набор данных, я разделил torch.utils.data.Dataset на подклассы. Набор данных ожидает фрейм данных в качестве входных данных, из которого мы сохранили только интересные столбцы. Затем мы удалили некоторые строки, в которых столбцы abstract и title совпадали с одним из «нежелательных» слов в FILTER_TITLE и FILTER_ABSTRACT соответственно. Это сделано потому, что статьи были списаны автоматически, и многие из них содержат нерелевантные записи вместо заголовка / аннотации.

Набор данных возвращает словарь, поскольку pd.DataFrame не поддерживается в PyTorch. Чтобы дать нашей поисковой системе больше контекста, мы объединяем title и abstract вместе, результат сохраняется в ключе title_abstract.

Теперь мы можем вызвать набор данных и посмотреть, все ли правильно.

ds = CovidPapersDataset.from_path('./dataset/metadata.csv')
ds[0]['title']

Выход:

'Sequence requirements for RNA strand transfer during nidovirus discontinuous subgenomic RNA synthesis'

Встроить

Теперь нам нужен способ создать вектор (встраивание) из каждой точки данных. Мы определяем класс Embedder, который автоматически загружает модель из HuggingFace's transformers с помощью библиотеки предложения_преобразователи.

Выбранная модель - gsarti / biobert-nli или BioBERT, настроенная на SNLI и MultiNLI для создания универсальных вложений предложений. Gabriele Sarti произвела тонкую настройку, код для ее воспроизведения доступен здесь.

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

Под капотом модель сначала токенизирует входную строку в токенах, а затем создает по одному вектору для каждого из них. Итак, если у нас есть N токенов в одной статье, мы получим [N, 768] вектор (обратите внимание, что токен часто соответствует части слова, подробнее о стратегиях токенизации читайте здесь. Таким образом, если две статьи имеют разный размер слова, мы будем иметь два вектора с двумя разными первыми измерениями.Это проблема, поскольку нам нужно сравнить их для поиска.

Чтобы получить фиксированное вложение для каждой бумаги, мы применяем средний пул. Эта методология вычисляет среднее значение каждого слова и выводит вектор тусклых изображений фиксированного размера [1, 768]

Итак, давайте запрограммируем Embedder класс

Мы можем попробовать наш встраивающий модуль на точке данных

embedder = Embedder()
emb = embedder(ds[0]['title_abstract'])
emb[0].shape // (768,)

И вуаля! Мы закодировали одну бумагу.

Поиск

Хорошо, мы знаем, как встраивать каждую статью, но как мы можем искать в данных с помощью запроса? Предполагая, что мы встроили все документы, мы также можем встроить запрос и вычислить косинусное сходство между запросом и всеми вложениями. Затем мы можем показать результаты, отсортированные по расстоянию (баллу). Интуитивно понятно, что чем ближе они находятся в пространстве встраивания к запросу, тем больше у них общего контекстного сходства.

Но как? Во-первых, нам нужен правильный способ управления данными и достаточно быстрого выполнения косинусного подобия. К счастью, на помощь приходит Elastic Search!

Эластичный поиск

Elastic Search - это база данных с одной целью, да вы не ошиблись: поиск. Сначала мы сохраним все вложения в эластичном, а затем будем использовать его API для выполнения поиска. Если ты такой же ленивый, как я, можешь установить эластичный поиск с докером

docker pull docker.elastic.co/elasticsearch/elasticsearch:7.6.2
docker run -p 9200:9200 -p 9300:9300 -e "discovery.type=single-node" docker.elastic.co/elasticsearch/elasticsearch:7.6.2

Идеально. Следующим шагом является сохранение вложений и информации о документах по эластичному поиску. Это очень простой процесс. Нам нужно создать index (новую базу данных), а затем создать по одной записи для каждой статьи.

Чтобы создать index, нам нужно описать для эластичного элемента, что мы хотим сохранить. В нашем случае:

Подробнее о создании индекса можно прочитать в эластичном поиске doc. Последняя запись определяет поле embed как плотный вектор с 768. Это действительно наше вложение. Для удобства я сохранил конфигурацию в файле a.json и создал класс с именем ElasticSearchProvider для обработки процесса сохранения.

Большая часть работы выполняется в create_and_bulk_documents, где мы просто разбираем по одной записи за раз и добавляем два параметра эластичного поиска.

К сожалению, Elastic Search не может сериализовать numpy массивы. Итак, нам нужно создать адаптер для наших данных. Этот класс принимает в качестве входных данных бумажные данные и встраивание и «адаптирует» их для работы в нашем ElasticSearchProvider.

Хорошо, у нас все на месте. Способ представления данных, способ кодирования в векторе и метод сохранения результата. Давайте все завернем и закодируем все бумаги.

Здесь есть две хитрости. Во-первых, мы используем torch.utils.data.DataLoader для создания пакетного итератора. В общем, загрузка данных в модель в пакетном режиме, а не в виде одной точки, повышает производительность (в моем случае x100). Во-вторых, мы заменяем параметр collate_fn в конструкторе DataLoader. Это потому, что по умолчанию Pytorch попытается преобразовать все наши данные в torch.Tensor, но не сможет преобразовать строки. Таким образом мы просто возвращаем массив словарей, вывод из CovidPapersDataset. Итак, batch - это список словарей длиной batch_size. После того, как мы закончили (~ 7 м на 1080ti), мы можем взглянуть на http://localhost:9200/covid/_search?pretty=true&q=*:*.

Если все работает правильно, вы должны увидеть наши данные, отображаемые эластичным поиском.

Сделать запрос

Мы почти закончили. Последний кусок головоломки - это способ поиска в базе данных. Эластичный поиск может выполнять косинусное сходство между одним входным вектором и целевым векторным полем во всех документах. Синтаксис очень прост:

{
    "query": {
        "match_all": {}
    },
    "script": {
        "source":
        "cosineSimilarity(params.query_vector, doc['embed']) + 1.0",
        "params": {
            "query_vector": vector
        }
    }
}

Где vector - наш ввод. Итак, мы создали класс, который принимает вектор в качестве входных данных и показывает все результаты запроса.

Давайте посмотрим на первый результат (я скопировал и вставил аннотацию из первой подходящей статьи)

es_search = ElasticSearcher()
es_search(embedder(['Effect of the virus on pregnant women'])[0].tolist())

По мере того как специалисты общественного здравоохранения реагируют на возникающие инфекции, особое внимание необходимо уделять беременным женщинам и их детям. Беременные женщины могут быть более восприимчивыми к возникающим инфекциям или более серьезно пострадать от них. Влияние новой материнской инфекции на эмбрион или плод предсказать сложно. Некоторые лекарства, рекомендуемые для профилактики или лечения, могут нанести вред эмбриону или плоду. Мы обсуждаем проблемы реагирования на возникающие инфекции среди беременных женщин и предлагаем стратегии по их преодолению.

Это сработало! Я также создал командную строку, в которой пользователь может ввести запрос. Конечный результат:

Больше запросов

Наконец, мы можем задать вопрос и найти интересные статьи. Эмпирически более подробные запросы работают лучше, поскольку они предоставляют больше контекста.

Например, мы можем узнать, Какова эффективность хлорохина при COVID-19. Результаты

Или Как COVID-19 связывается с рецептором ACE2?

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

Выводы

В этом проекте мы создаем семантический браузер для поиска по более чем 50 тысячам документов о COVID-19. Оригинальный проект, над которым я работал со студентами из Университета Триеста, находится здесь. Живая демонстрация доступна здесь

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

Подтверждение

Я хотел бы поблагодарить Gabriele Sarti за помощь в написании этой статьи, Marco Franzon и Tommaso Rodani за поддержку в реализации эластичного поиска.

Спасибо за чтение

Быть в безопасности,

Франческо Саверио Цуппичини