Соединяем физику и лирику. Как я собрал рекомендательную систему для стихов с помощью Flask, sqlite-vec и Hugging Face

Люблю я кодить и стихи —
Вот, в общем, все мои грехи...

А. С. Пушкин

Привет! Я Константин Хабазня, преподаватель программирования и математики, а также автор (что бы это ни значило).

N-нное время назад увлёкся NLP (Natural Language Processing), что вполне логично для писателя, который кодит (или кодера, который пишет).

Почитав интернет и пару вводных книжек, отправился учиться на ДПО в МФТИ. В качестве выпускного проекта придумал себе задачу — создать рекомендательную систему для стихов. Для песен рекомендательные системы есть, для книг — есть, для фильмов — есть, для стихов — нет. Непорядочек. А значит, этому городу нужен герой! Чип и Дейл Обезьянка с Питоном спешат на помощь!

По ходу дальнейшего текста расскажу, как собрал первый прототип системы на небольшом датасете с Kaggle, используя Flask, sqlite-vec, энкодер с Hugging Face и чистый (ну почти) HTML.

0. Как всё работает

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

Берём много стихов (на картинке это Data). С помощью энкодера и библиотеки Sentence-Transformers (BERT-model) превращаем стихи в 512-мерные векторы. Складываем векторы и соответствующие им тексты в векторную базу (Vector base).

Пользователь (User) вводит в систему (App) стихотворение. Модель (BERT-model) на лету превращает его в вектор (стихотворение, а не пользователя). Этот вектор сверяется с векторами базы (Vector base). Система (App) выбирает три наиболее схожих по косинусному расстоянию вектора, берёт соответствующие им тексты (включая исходный, если он есть в базе) и возвращает их пользователю (User) в качестве рекомендаций.

Фух. Вы что-нибудь поняли? Я тоже нет. Давайте разбираться.

1. Датасет

«Content is King» — сказал Билл Гейтс ещё в 1994 году и, кажется, был прав. Датасет — уже полдела, и это не преувеличение.

Данных должно быть МНОГО, они должны быть ЧИСТЫМИ. А ещё их надо ХРАНИТЬ в какой-либо удобной для использования структуре.

Я решил начать с простого. Нашёл на Kaggle датасет:

Классические стихи, небольшой объём — для сборки прототипа годится как нельзя лучше

2. Векторная база

Для машины тексты бесполезны. Душевных порывов она не понимает (пока), и виршами в ресторане под луной (с бокальчиком вина) её не соблазнишь. Тексты надо превратить в векторы, которые будут храниться в векторной базе данных. Вот на векторы машина уже клюнет. Дело пошло.

Методом проб и ошибок остановил свой выбор на библиотеке sqlite-vec. Почему не облачное векторное хранилище?

Во-первых, запрос к сторонним северам удлинил бы время отклика системы. Во-вторых, для первого прототипа не хотелось плодить сущности (в которых ещё нужно разобраться) — чем проще, тем лучше.

Базу собираем в Google Colab. Грузим библиотеки и датасет, инициализируем модель:

!pip install sqlite-vec
import sqlite3
import sqlite_vec
import pandas as pd
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List

model = SentenceTransformer('distiluse-base-multilingual-cased-v1', device='cpu')

df = pd.read_csv('russianPoetryWithTheme.csv')

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

texts = []
for idx in range(16694):
  author = df.iloc[idx]['author']
  name = df.iloc[idx]['name']
  text = df.iloc[idx]['text']
  texts.append(author + '\n' + str(name) + '\n' + text)

def get_normalized_embeddings(texts: List[str]) -> np.ndarray:
    embeddings = model.encode(texts, convert_to_numpy=True)
    return embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

embeddings = get_normalized_embeddings(texts)

Инициализируем базу "vectors.db" и создаём в ней таблицу, в которой будут храниться пары «вектор — текст»:

db = sqlite3.connect("vectors.db")
db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)

db.execute("""
    CREATE VIRTUAL TABLE documents USING vec0(
        embedding float[512],
        text text
    )
""")

Заполняем базу векторами и текстами:

from tqdm.notebook import tqdm
BATCH_SIZE = 1
with db:
    for i in tqdm(range(0, len(texts), BATCH_SIZE)):
        batch = [
            (j+1,
             embeddings[j].tobytes(),
             texts[j])
            for j in range(i, min(i+BATCH_SIZE, len(texts)))
        ]
        db.executemany(
            "INSERT INTO documents(rowid, embedding, text) VALUES (?, ?, ?)",
            batch
        )

Копируем в проект готовый файл "vectors.db".

3. Веб-приложение на Flask

Итак, подошли к самому интересному — пишем сервер на Flask. Выбрал именно этот фреймворк, так как слышал о нём как о золотой середине между простотой и функциональностью — идеально для новичка.

Помним: чем проще — тем лучше, тем быстрее разберёмся с кодом, тем меньше швов придётся латать, тем быстрее соберём работающий прототип и получим дофаминчик. Друзья скажут: «ты гений!» Работодатели завалят баксами! (Работодатели, ау, отзовитесь! Перед вами... гений... бэкенда же... вы что... разве... не понимаете... гений... бэкенда...😢)

По Flask нашёл хороший вводный туториал, который мне очень понравился и закрыл базу по созданию простенького сервера. Большего пока и не нужно.

Собираем звездолёт.

Вот так выглядит "main.py".

from flask import Flask, render_template, request
from translation_pipe import translation
from get_vec_db import search_similar_texts

app = Flask(__name__, template_folder='templates')

@app.route("/", methods=['GET'])
def text_form():
    return render_template('index.html')

@app.route("/poe3", methods=['GET', 'POST'])
def similarity_form():
    if request.method == 'GET':
        return render_template('poe3.html')
    elif request.method == 'POST':
        text = request.form.get('textarea')
        results = search_similar_texts(text, top_k=3)
        text_match = []
        for row in results:
            row_list = row[1].split('\n')
            text_match.append([row[2], row_list]) #f"Доля схожести: {row[2]:.4f}, Текст: {row[1]}...")
        return render_template('validation.html', text_match=text_match)

@app.route("/validation", methods=['GET', 'POST'])
def validation_form():
    if request.method == 'GET':
        return render_template('poe3.html')
    elif request.method == 'POST':
        try:
            score = float(request.form.get('validarea'))
        except Exception:
            return render_template('poe3.html', text_valid=['Не сработало. Попробуй оценить снова!'])

        with open('validation.txt', 'r', encoding='utf-8') as f:
            validate_list = list(map(float, f.readlines()))
        validate_list[0] += 3.0
        validate_list[1] += score
        validate_list[2] = (validate_list[1]/validate_list[0])*100
        to_file = str(validate_list[0]) + '\n' + str(validate_list[1]) + '\n' + str(validate_list[2])
        with open('validation.txt', 'w', encoding='utf-8') as f:
            f.write(to_file)
        return render_template('congrats.html', text_valid=['Спасибо за оценку!'])

@app.route("/score", methods=['GET'])
def get_score():
    with open('validation.txt', 'r', encoding='utf-8') as f:
        score_list = list(map(float, f.readlines()))
    return render_template('score.html', score_list=score_list)

if __name__ == "__main__":
    port = 8080
    app.run(debug=True, host='0.0.0.0', port=port)

Он умеет запускать сервер и назначать функции-обработчики любого запросов из браузера.

"/" — дефолтный GET, который показывает стартовую страницу сайта (не имеющую отношения к проекту).

"/poe3" — вызывает страницу с окошком и кнопкой для ввода текста пользователем, а также обрабатывает стихотворение, введённое пользователем: оно отправляется в функцию "search_similar_texts()", которая возвращает три схожих текста. Эту рекомендацию мы закидываем на страницу "validation.html", куда и редиректим пользователя. Он видит выдачу, а также запрос оценить рекомендацию.

"main.py" главный узел системы. Функция "search_similar_texts()" запакована в отдельный модуль "get_vec_db.py", который выглядит так:

import sqlite3
import sqlite_vec
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List

model = SentenceTransformer('distiluse-base-multilingual-cased-v1', use_auth_token=False)
db = sqlite3.connect('vectors.db', check_same_thread=False)
db.enable_load_extension(True)
sqlite_vec.load(db)
db.enable_load_extension(False)


def search_similar_texts(query_text: str, top_k: int = 5) -> List[tuple]:
    query_embedding = model.encode(query_text, convert_to_numpy=True)
    query_embedding = query_embedding / np.linalg.norm(query_embedding)

    rows = db.execute(
        "SELECT rowid, text, embedding FROM documents"
    ).fetchall()

    res = []

    for r in rows:
        stored_vec = np.frombuffer(r[2], dtype=np.float32)
        similarity = np.dot(query_embedding, stored_vec)
        res.append((r[0], r[1], float(similarity)))

    return sorted(res, key=lambda x: x[2], reverse=True)[:top_k]

Вернёмся к "main.py". А точнее к обработчику запроса "/validation", где мы просим пользователя оценить релевантность ответа системы после выдачи рекомендаций. Здесь остановимся подробнее.

4. Валидация

Хьюстон, у нас проблема: как валидировать релевантность выдачи, если машина равнодушна к виршам под луной и лирика ей до лампочки? Что значит, что система выдала «хорошую» рекомендацию? Каковы критерии? Мы не можем измерить «схожесть» двух стихов машинно. Это работа для человека — вопрос вкусов.

Итак, человеческая валидация. Каждый пользователь может выставить системе максимум 3 балла. Значит, максимальное количество баллов = 3*(количество пользователей, провалидировавших систему).

Однако злые уважаемые пользователи ставят «реальные» баллы: 1, 2, 3. Соответственно, общее «реальное» количество баллов меньше «максимального». Кто в школе учил доли и проценты поймёт, что будет дальше: делим «реальное» количество баллов на «максимальное», умножаем на 100 и получаем процент точности рекомендательной системы.

Процент... провалидировавшие... делим... какое небо сегодня красивое...

Именно этим и занимается обработчик запроса "/validation". В нём вся описанная выше математика работает в связке с txt-файлом, в котором мы храним текущие показатели баллов. Открываем файл, обновляем баллы, записываем обновлённые баллы в файл, закрываем файл. Приводить его не буду ввиду тривиальности, а также потому что его копия уже давно работает в деплое с текущими показателями баллов.

Их, кстати, можно увидеть по запросу "/score", который тоже обрабатывается нашим "main.py". Обработчик этого запроса открывает txt-файл с баллами и рендерит отчёт на отдельной странице.

Раскидал ссылку на систему по своим соцсетям. На момент написания статьи score такой. Систему оценили 134 человека

Кстати, о страницах.

5. HTML + Jinja2 = ❤️

Jinja2 – серый кардинал происходящего. До настоящего момента он управлял нашим фронтендом, а мы и не замечали. Его вотчина — HTML-шаблоны, в которые можно встраивать Python-код. Именно благодаря такому функционалу обработчики запросов из "main.py" могут отдавать в HTML-страницу структуры данных Python, а те будут корректно с ними работать.

Я понимаю, вы уже устали. Приведу для примера код страницы валидации "score.html", скрин которой мы только что видели:

{% extends "base.html" %}
{% block title %}Страница валидации{% endblock %}

{% block content %}
    <h1>Страница валидации</h1>

    <div>Максимальная оценка: {{ score_list[0] }}</div>
    <br>
    <div>Фактическая оценка: {{ score_list[1] }}</div>
    <br>
    <div>Точность в процентах: {{ score_list[2] }} %</div>
    <br>

{% endblock %}

Смотрите, что происходит. С помощью хитроумного синтаксиса я могу обрабатывать в этом HTML-коде питоновский список — и всё работает.

Ещё один маленький рывочек. Строчка "{% extends "base.html" %}" означает, что шаблон страницы валидации наследует базовому шаблону "base.html", чтобы каждый раз руками не прописывать одни и те же постоянные элементы (например, ссылки на мои соцсети, на которые вы сразу же подпишитесь, потому что вам очень понравилась моя статья, и вообще, я смешной, ведь правда?😊)

<!DOCTYPE html>
<html lang="en, ru">
<head>
    <meta charset="UTF-8">
    <title>{% block title %}Default Title{% endblock %}</title>
</head>
<body>
    <p></p>
    {% block content %}{% endblock %}
    <br>
    <br>
    <a href="https://t.me/habaznyaclub">откровения технаря о поэзии 📖</a>
    <br>
    <br>
    <a href="https://t.me/habaznya_math">откровения поэта о коде, AI, NLP 🛠</a>
</body>
</html>

Ну всё-всё, не плачьте, мы уже закончили, всё позади. Я же сказал: как комарик укусит. И ничего в итоге сложного — всё просто и понятно, да?

Усталь

6. Бонус-трек: деплой на сервер

Есть вещи, которые выходят за рамки нашей истории. Например, работа Flask и Jinja2, синтаксис sqlite-vec и особенности энкодеров Hugging Face. Детали деплоя — одна из таких вещей. В сети куча материалов — гуглите, чатботьте, пишите в поддержку своего хостера.

У моего есть такая классная фича — подгружать деплой из git-репозитория. Делаешь новый push, сервер подхватывает его на лету, деплоит обновления сам. Звучит прекрасно, но, чтобы настроить такой деплой, разобраться с зависимостями и так далее, у меня ушло едва ли не больше времени, чем на сам проект.

В общем, в этом месте вы вспомните все нецензурные слова, которые знаете. Удачи!

7. Наконец-то вывод

Итак, разобравшись с азами фреймворков Flask и Jinja2, почитав документацию sqlite-vec, научившись пользоваться энкодерами из коробки (fine-tuning не делали), потратив определённое время и пару метров нервов на репозиторий, зависимости, деплой, мы собрали собственную рекомендательную систему для стихов. Точность у неё небольшая — 67%, но для первого приближения неплохо.

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

Я ввожу первые две строчки стихотворения, которое точно содержится в базе, но не получаю его в выдаче. Или получаю, но не первым. Что это значит? Вектор двух строчек может не вобрать в себя семантику всего стихотворения, поэтому в действительности будет далёк от вектора целого стихотворения.

Речь идёт именно о рекомендательной системе. Целое-к-целому. Как в «Яндекс Музыке». А не поиск по фрагменту, как в «Шазаме».

---

С вами были Обезьянка и Питон. Если было полезно или, по крайней мере, смешно, поделитесь статьей с друзьями, накидайте реакций и комментариев, а может быть даже подпишитесь на мои откровения технаря о литературе и/или литератора о коде.

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

Клац-клац