Искусственный интеллект для поиска новых лекарственных препаратов

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

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

Требования к решению:

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

Стоит отметить, что работа над проектом велась совместно с командой клиента и научными командами двух университетов, расположенных в Бельгии и Австрии. Это требовало умения работать в многоязыковой среде (C++, Java, Python, R) и интегрировать все части в единый конвейер данных.

В качестве основного языка разработки был выбран Python, поскольку в научных исследованиях чаще всего используются Python-скрипты и файлы Jupyter Notebook. Мы создали библиотеку инструментов, которая позволяет пользователям строить конвейеры в соответствии с потребностями бизнеса. Библиотека имеет интуитивно понятный пользовательский интерфейс, напоминающий известные библиотеки Python, такие как scikit-learn.

В библиотеке 7 основных блоков:

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

Подготовка данных

Инструменты для предварительной обработки данных позволяют пользователям

• конвертировать записи соединения из различных форматов в один стандарт (например, в канонический SMILES),
• стандартизировать представление структуры молекулы,
• удалять тяжелые атомы и др.

Инструменты для фингерпринтирования используют различные методы для трансформации молекулы в числовой тензор.

В соответствии с требованием клиента, описанные выше инструменты позволяют использовать разнообразные алгоритмы как из популярных химических библиотек (Chemaxon, RDKit), так и из новейших научных исследований. Для этого мы:

• адаптировали и внедрили ряд алгоритмов из недавних научных публикаций,
• разработали на Python оболочку для Chemaxon, поскольку эта библиотека написана на Java,
• разработали оболочки для некоторых инструментов на основе RDKit, чтобы компенсировать некоторые недочёты в работе этого пакета.

Машинное обучение

Инструменты для кластеризации используют как классические, так и продвинутые алгоритмы. По запросу клиента мы модифицировали библиотеку Faiss, написанную на C++: мы создали для неё оболочку на Python, которая добавляет метрику сходства Tanimoto — наилучший метод для вычисления схожести фингерпринтов молекул.

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

Для обучения с учителем мы использовали различные алгоритмы, в том числе:

• классическое машинное обучение — SVM, KNN, SGD и т. д.,
• алгоритмы, основанные на дереве поиска решений, такие как XGBoost, NGBoost,
• алгоритмы глубокого обучения на основе Pytorch и
• нейронные сети обмена сообщениями (message-passing neural networks, MPNN). Последние являются мощными инструментами для прогнозирования свойств молекулярных графов. Некоторые из них мы реализовали с нуля на основе научных работ.

Оценка результатов

Для оценки неопределённости мы включили индуктивные конформные предикторы Мондриана и предикторы Venn-ABERS. Кроме того, модели, построенные на основе ансамбля (метамодели) могут вычислять неопределённость данных (алеаторическую) и неопределённость знаний (эпистемическую).

Модели машинного обучения и оценщики неопределённости собраны в единую сущность, что позволяет

• обучать одну модель,
• выполнять k-кратную проверку,
• выполнять вложенную перекрестную проверку (кросс-валидацию).

Кроме того, это позволяет пользователям оптимизировать гиперпараметры двумя методами: поиск по сетке и древесно-структурированная оценка Парзена.

Процедура обучения показала высокую производительность, так как вычисления распределены между всеми ядрами процессора и несколькими графическими процессорами.

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

API

Ко всем конвейерам можно обращаться посредством API. Пользователи (химики, обучающие модель ИИ) могут создать простой YAML-файл, в котором они описывают, какие методы они хотели бы использовать для предварительной обработки соединения и создания фингерпринта, а также какую модель они хотели бы использовать для вывода. Затем они запускают сервер Flask. После этого клиенты (другие химики) могут обращаться размещённой на сервере модели и запрашивать предсказания для тех или иных химических соединений. Высокая производительность позволяет использовать этот API в приложениях реального времени, таких как инструменты рисования молекул, которые генерируют прогнозы по ходу выполнения. Иными словами, химик «рисует» химическое соединение в редакторе и через API сразу получает информацию о его характеристиках, предсказанных моделью ИИ.

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

Продукт хорошо задокументирован и содержит множество примеров файлов Jupyter Notebook для различных задач и ситуаций.

Наше решение позволило обрабатывать большие объёмы данных, представленных в различных форматах, используемых в научном сообществе. При поиске потенциальных лекарственных средств используются различные алгоритмы машинного обучения; сложные вычисления можно распределить на несколько ядер процессора, графических процессоров или кластеров.

Использование Python в качестве основного языка разработки и интуитивно понятный пользовательский интерфейс позволяют пользователям быстро освоить работу с библиотекой. Кроме того, они могут интегрировать отдельные блоки нашей библиотеки в сторонние скрипты и инструменты Python — например, инструменты для активного обучения.

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