Zero-shot классификация текста с Hugging Face и Qdrant
13 июля 2022 г.
Введение
В этой статье я покажу, как можно использовать NLI-модели (Natural Language Inference) от Hugging Face совместно с Qdrant для классификации текстов — без необходимости обучать модель с нуля.
Возможно, у вас есть большой массив данных для обучения классификатора, но не хватает вычислительных ресурсов или времени. Именно здесь на помощь приходит метрическое обучение (metric learning)!
Zero-shot классификация означает, что модели способны распределять данные по категориям, на которых они не обучались. Как правило, этот подход применяется там, где данных для создания классификатора недостаточно. Однако сегодня модели становятся всё больше, и их обучение требует всё больше времени.
Именно поэтому я решил создать этот туториал: он позволит вам построить классификатор за считанные минуты, без обучения, получив при этом хорошую базовую модель!
Датасет
Я буду использовать новостной датасет, содержащий 120 тысяч записей с четырьмя метками: World (Мир), Sports (Спорт), Business (Бизнес) и Science/Tech (Наука/Технологии). Датасет доступен здесь. Напомню, что он используется исключительно в учебных целях. Не стесняйтесь экспериментировать с другими, более объёмными датасетами!
Для загрузки датасета мы воспользуемся трансформерами Hugging Face.
Модель
Я буду использовать модель sentence-transformers/stsb-distilbert-base. При необходимости вы можете заменить её на другую — например, если вам нужна поддержка иного языка или другая NLI-модель.
Итак, у нас есть модель и данные. Мы воспользуемся NLI-моделью для извлечения признаков из текста, а затем на основе полученных векторов найдём наиболее близкие.
Время выполнения кода зависит от размера датасета и объёма памяти. Подберите batch_size в соответствии с доступной памятью для ускорения процесса.
Qdrant
Если вы уже знакомы с Qdrant, этот раздел можно пропустить.
Наши векторы готовы! Теперь нужно загрузить их в Qdrant перед поиском. Qdrant — это по сути база данных для векторного поиска. Можно представить Qdrant как Redis для векторов: вместо ключа вы передаёте вектор, а в ответ получаете ближайшие векторы и связанные с ними данные (payload).
Если Docker ещё не установлен на вашем компьютере, следуйте этой инструкции. Qdrant — проект с открытым исходным кодом, и вы можете напрямую скачать его Docker-образ.
docker pull qdrant/qdrant
После загрузки образа запустите контейнер Qdrant.
Загрузка векторов в Qdrant
Всё готово! Протестируем нашу модель.
Классификация
Чтобы выполнить поиск по входному тексту, сначала необходимо получить его векторное представление. После этого можно искать по вектору с помощью Qdrant.
В результате вы увидите наиболее близкие предложения и их метки. Параметр limit можно изменить, чтобы получить больше или меньше результатов.
Как видите, было найдено 5 ближайших строк, и возвращён их payload. Метки соответствуют значению 3, что является меткой для категории Tech/Science. Попробуйте разные входные данные и посмотрите на результаты!