Рекомендатель кино или как я писал свое DIY-решение для поиска новых фильмов

Вечер. Пересматриваю «Пятницу 13». Не люблю пересматривать фильмы, даже хорошие. Но выбрать интересное кино из потока новинок сложно. Поэтому мне захотелось написать свой рекомендатор кино. Этим и займусь в выходные. 

В статье покажу, что получилось написать за 2 дня. Писал всё «на коленке» по доступным библиотекам и данным. Получилcя DIY-рецепт. Всё платформозависимое работает в Docker, чтобы повторить и развернуть можно было везде. 

Определимся с деталями проекта

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

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

Основные задачи

Накидаю задач для проекта: 

• найти датасет с информацией о кино. Чем больше данных, тем лучше;

• векторизовать фильмы для сравнения и поиска;

• настроить векторную базу данных;

• написать интерфейс на Flask.

Писать я буду на Python, ведь на нём удобно работать с векторами и датасетами. 

Поиск датасета фильмов

На поиск датасета я потратил 3 часа. Перебрал 10 вариантов. В одних наборах данных было меньше 1000 строк, а другие не содержали важных колонок: описания, актёров, жанров. 

Сначала я нашёл официальные датасеты imdb.com, но они разделены на 4 файла. Их нужно мержить по ID в одну таблицу. У файла с работниками нелинейная структура: много строчек на 1 фильм. Для мержа надо самостоятельно отделить и сгруппировать актёров. Колонки с описаниями в официальном датасете нет.

Далее я искал данные на Kaggle и Github. Остановился на TMDB + IMDB Movies Dataset 2024. Это датасет в CSV-формате на 1 млн. строк. В нём 27 колонок: название, актёры, рекламные слоганы, описания, жанры и другие. Его я распаковал в movies.csv.

Для тестов я решил оставить только фильмы с известными актёрами. Для фильтрации использовал датасет Top 100 Greatest Hollywood Actors of All Time. 

Готовим данные

Перед использованием нужно почистить датасет: удалить строки без данных, отфильтровать фильмы по статусу и сборам. Оставшиеся колонки с NaN — не критичные, поэтому их я заполняю значениями подходящего типа.

$ python -m venv env
$ . ./env/bin/activate
$ pip install pandas

import pandas as pd

movies = pd.read_csv(
    'movies.csv',
    usecols=['id', 'title', 'release_date', 'revenue', 'status',
        'imdb_id', 'original_language', 'original_title', 'overview',
        'tagline', 'genres', 'production_companies',
        'production_countries',
        'spoken_languages', 'cast', 'director', 'writers',
        'imdb_rating', 'imdb_votes'
    ]
)

# Вырежем строки с пустыми колонками
movies.dropna(subset=[
    'title', 'overview', 'genres', 'release_date', 'status',
    'cast', 'director', 'writers', 'imdb_id'], 
    inplace=True
)

# Оставим фильмы, которые уже вышли и что-то заработали
movies = movies[movies['status'] == 'Released']
movies = movies[movies['revenue'] > 0]
movies.drop(['status', 'revenue'], axis=1, inplace=True)
movies.reset_index(drop=True, inplace=True)

# Заполняем нулями пустые рейтинги и отзывы
movies.fillna({'imdb_rating': 0, 'imdb_votes': 0}, inplace=True)

# Заполняем пропущенные слоганы пустотой строкой
movies['tagline'] = movies['tagline'].fillna('')

Теперь загружу топ актёров. Сначала из даты рождения я выделю год. Потом отсортирую всех по нему в порядке убывания и получу имена первых 50 актёров. Преобразую их в set-множество.

actors = pd.read_csv('actors.csv')

# Выделяем год рождения
actors['Year of Birth'] = actors['Date of Birth'].apply(
    lambda d: d.split()[-1]
)

# Сортируем по году рождения и берём первых 50 актёров
actors.sort_values('Year of Birth', ascending=False, inplace=True)
actor_set = set(actors.head(50)['Name'])

Создам отдельный датафрейм для отфильтрованных фильмов. В него запишу все фильмы, у которых список актёров пересекается с set-множеством из ТОПа.

def actors_intersect(actors: str):
    """Проверяем пересечение актёров."""
    actors = set(actors.split(', '))
    return bool(actors.intersection(actor_set))

# Отберём фильмы для теста
movies['to_test'] = movies['cast'].apply(actors_intersect)
df = movies[movies['to_test']]
df.reset_index(drop=True, inplace=True)

В тестовый набор попал 1981 фильм.

Объединим франшизы

Некоторые фильмы выпускаются в рамках франшиз. Например, «Мстители» и «Звёздные войны». Такие фильмы должны попадать в рекомендации вместе. Картины одной серии выпускают с похожими названиями. Связать их можно кластеризацией. 

Для кластеризации нужны векторы. Чтобы не тратить лишнего времени, нужен быстрый алгоритм векторизации. Я взял TfidfVectorizer из пакета sklearn. 

Для TfIdf я задал английский словарь стоп-слов, чтобы исключить служебные части речи. Также в параметрах я ограничил вхождения токена min_df и max_df.

Кластеризовать я буду через DBSCAN. Он борется с шумом и может работать с «плотными» облаками точек. А главное — для него не надо заранее знать количество кластеров. 

from sklearn.cluster import DBSCAN
from sklearn.feature_extraction.text import TfidfVectorizer

tfidf = TfidfVectorizer(stop_words='english', min_df=2, max_df=0.2)
matrix = tfidf.fit_transform(df['title'])
dbscan = DBSCAN(eps=0.9, min_samples=2)
df['title_cl'] = dbscan.fit_predict(matrix).astype('str')

Выделяем именованные сущности

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

Выделять сущности я буду через spaCy. Для работы spaCy скачаем англоязычную модель en_core_web_sm, которую заранее обучили на текстах из интернета. 

$ pip install spacy
$ python -m spacy download en_core_web_sm

Теперь выделим сущности из колонок title, tagline, overview. 

import spacy
nlp = spacy.load('en_core_web_sm')

def extract_ents(row):
    """Соберём вместе три колонки и выделим сущности."""
    text = row['title'] + '. ' + row['tagline'] + '. ' + row['overview']
    ents = nlp(text).ents
    return ', '.join(set([ent.text.lower() for ent in ents]))

# В колонку ents запишем список сущностей
df['ents'] = df.apply(extract_ents, axis=1)

Категориальные данные

В датафрейме есть категориальные данные — данные с ограниченным числом значений. Это жанры, режиссёры, сценаристы, именованные сущности и кластеры названий. Их я также использую для рекомендаций. 

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

Для векторизации категориальных данных я снова использую векторизатор TF-IDF. Обычно его не применяют на категориальных данных. Но он гибкий и у него много параметров. Можно контролировать размер выходного вектора, отсеять признаки по частоте. Алгоритм IDF даст редким признакам больше веса, чем часто встречающимся. 

Фильмы из одной франшизы с похожим набором сущностей должны быть выше в списке рекомендаций. На них я выделю больше значений в результирующем векторе. А сценаристам и режиссёрам значений дам меньше. Для векторизатора задам разные настройки min_df в зависимости от категориального параметра. А длину вектора вычислю динамически по количеству признаков: 

• для сценаристов минимум 4 вхождения, длина вектора не более 10% от словаря признаков;

• для режиссёров минимум 2 вхождения и длина не более 20%;

• франшизы и сущности без ограничений.

cat_tfidf_args = {
    'tokenizer': lambda cats: [c.strip() for c in cats.split(',')],
    'token_pattern': None,
}
categories_cols = [
    ('writers', 4, 0.1),
    ('director', 2, 0.2),
    ('title_cl', 1, 1),
    ('ents', 1, 1),
]
for col, min_df, coeff in categories_cols:
    tfidf = TfidfVectorizer(min_df=min_df, **cat_tfidf_args)
    # Определим число признаков и ограничим длину вектора
    tfidf.fit(df[col])
    tfidf.max_features = int(len(tfidf.vocabulary_) * coeff)
    # Запишем вектор в датафрейм
    df[col + '_vec'] = tfidf.fit_transform(df[col]).toarray().tolist()

В тестовом датасете всего 19 уникальных жанров, но они идут наборами с разной длиной. На таких данных TF-IDF сработает неэффективно. Поэтому я применяю MultiLabelBinarizer. Он вернёт для каждой строчки вектор из 19 значений.

from sklearn.preprocessing import MultiLabelBinarizer

mlb = MultiLabelBinarizer()
# На вход MultiLabelBinarizer нужно передать список
# Поэтому строку с жанрами разделим по запятой
genres = df['genres'].apply(
    lambda row: [g.strip() for g in row.split(',')]
)
df['genres_vec'] = mlb.fit_transform(genres).tolist()

Векторизуем описания

Чтобы связывать фильмы по описаниям, нужно векторизовать колонку overview. Для этого я использую модель ROBERTA — переученный BERT от Google с оптимизацией. Поэтому она векторизует лучше оригинала. На Хабре есть статья о различиях BERT'а и ROBERTA. Я возьму transformers от hugging face для работы с моделью. Для transformers нужен бэкэнд, поэтому поставлю ещё и torch.

Перед запуском модели надо посчитать входные данные. Это делает токенизатор. Я указал ему параметр return_tensors на “pt”, чтобы конечные тензоры были в формате PyTorch. Через truncation и max_length я ограничил входные данные до 512 токенов. 

Я запущу модель в контексте no_grad. Это отключит расчёт градиента обратного распространения. Его используют во время обучения моделей, чтобы вычислять ошибку и править веса. Но сейчас я только запускаю модель, поэтому градиент мне не нужен.

$ pip install transformers torch

import torch
from transformers import RobertaModel, RobertaTokenizer

tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
model = RobertaModel.from_pretrained("roberta-base")

def get_embed_text(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512
    )
    with torch.no_grad():
        out = model(**inputs)
    return out.last_hidden_state.mean(axis=1).squeeze().detach().numpy()

df["overview_vec"] = df["overview"].apply(get_embed_text)

У меня видеокарта RTX 3050 и процессор Ryzen 3200G, поэтому для обработки текстов двух тысяч фильмов нужно 2-3 минуты. Обработка того же объёма текста только на процессоре занимает 7-10 минут. 

Объединим векторы

Я векторизовал колонки датафрейма по отдельности. Теперь объединю их векторы через numpy. Размерность конечного вектора выведу на экран.

Датафрейм с векторами запишу в формате pickle в embedded.pkl: 

import numpy as np

def concatenate(row, col_names):
    """Объединим векторы из колонок в один вектор фильма."""
    embedding = np.concatenate(row[col_names].values)
    embedding = np.concatenate((embedding, row['genres_vec']))
    embedding = np.concatenate((embedding, row['overview_vec']))
    return embedding

cat_col_names = [col + '_vec' for col, _, _ in categories_cols]
df.loc[:,'embedding'] = df.apply(lambda x: concatenate(x, cat_col_names), axis=1)
print('Embedding shape:', df['embedding'][0].shape)
df.to_pickle('embedded.pkl')

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

Векторный Postgres

Для хранения векторов я выбрал привычный PostgreSQL. Базу выбирал по критериям:

• база должна работать и с векторами и с обычными данными. Например, для хранения «сырых» полей: названия, описания и рейтинга;

• нужен поиск фильмов по названию. 

Чтобы работать с векторами в постгрисе, потребуется расширение pgvector. Оно добавляет тип данных vector и реализует поиск ближайших векторов.

Важно: в pgvector размерность векторов должна быть менее 16 тысяч. 

Далее я настроил контейнер для базы с Docker и Docker-compose. Вот содержимое файлов: 

# psql.Dockerfile
FROM postgres:16-alpine3.20

RUN apk update && apk add --no-cache postgresql16-plpython3
RUN apk update; 
    apk add --no-cache --virtual .vector-deps 
      postgresql16-dev 
      git 
      build-base 
      clang15 
      llvm15-dev 
      llvm15; 
    git clone https://github.com/pgvector/pgvector.git /build/pgvector; 
    cd /build/pgvector; 
    make; 
    make install; 
    apk del .vector-deps

COPY docker-entrypoint-initdb.d/* /docker-entrypoint-initdb.d/

# docker-compose.yml
version: '3'

services:
    postgres:
        build:
            dockerfile: psql.Dockerfile
            context: .
        ports:
            - 5432:5432
        environment:
            - POSTGRES_USER=user
            - POSTGRES_PASSWORD=password
            - POSTGRES_DB=db
        volumes:
            - pgdata:/var/lib/postgresql

volumes:
    pgdata:

Создал инициализирующий SQL-скрипт для базы. В нём загрузил pgvector:

$ mkdir docker-entrypoint-initdb.d/
$ touch docker-entrypoint-initdb.d/init.sql

-- docker-entrypoint-initdb.d/init.sql
CREATE EXTENSION IF NOT EXISTS vector;

Проверяю работу: запускаю контейнеры через docker-compose и смотрю логи. 

$ docker-compose up -d
$ docker-compose logs

Теперь нужно создать таблицу, чтобы сохранить в ней векторы фильмов. Я добавлю колонки для названия, рейтинга, описания и IMDb ID. Рекомендатор отобразит их в ленте. Для вектора я задам поле embedding. В pgvector размерность вектора нужно знать заранее. У меня получилась размерность вектора 6399 — эту информацию выдал скрипт обработки датафрейма с фильмами. 

Добавлю создание таблицы в файл docker-entrypoint-initdb.d/init.sql.

-- docker-entrypoint-initdb.d/init.sql
...

CREATE TABLE movies (
    tconst VARCHAR(16) PRIMARY KEY NOT NULL UNIQUE,
    title VARCHAR(64) NOT NULL,
    title_desc VARCHAR(4096) NOT NULL,
    avg_vote NUMERIC NOT NULL DEFAULT 0.0,
    embedding vector(6399)
);

После изменения init.sql нужно пересобрать Docker-образ и перезапустить базу. Это займёт не больше минуты, ведь Docker кэширует сборки.

$ docker-compose down && docker-compose build && docker-compose up -d

Векторы фильмов нужно загрузить в базу. Для этого я написал скрипт в отдельном файле, который берёт данные из embedded.pkl. Работать с базой я буду через библиотеку psycopg. А чтобы она работала с векторами, нужна библиотека pgvector-python.

$ pip install psycopg pgvector

# filldb.py
import asyncio

from pgvector.psycopg import register_vector_async
import pandas as pd
import psycopg

df = pd.read_pickle("embedded.pkl")

async def fill_db():
    async with await psycopg.AsyncConnection.connect(
        'postgresql://user:password@localhost:5432/db'
    ) as conn:
        await register_vector_async(conn)
        async with conn.cursor() as cur:
            for _, row in df.iterrows():
                await cur.execute(
                    """
                    INSERT INTO movies (
                        tconst,
                        title,
                        title_desc,
                        avg_vote,
                        embedding
                    ) VALUES (%s, %s, %s, %s, %s)
                    """,
                    (
                        row['imdb_id'],
                        row['title'],
                        row['overview'],
                        row['imdb_rating'],
                        row["embedding"],
                    ),
                )

asyncio.run(fill_db())

После запуска все векторы запишутся в базу и можно будет искать похожие фильмы.

$ python filldb.py

Интерфейс на Flask

Чтобы удобно пользоваться Рекомендатором, я сделал веб-интерфейс на Flask и Jinja.

$ pip install flask

Интерфейс состоит из одной страницы с полем ввода для названия фильма. Лента рекомендаций появляется ниже после отправки формы. Запрос и номер страницы передаю в GET-параметрах. На странице отображается по 20 фильмов. Для формы поиска я сделал подсказки: 20 случайных названий, которые вытащил из базы данных. Вместо постеров прикрутил случайные фотографии с собаками. Так выдача Рекомендатора смотрится веселее. 

Код Flask приложения и Jinja шаблон привожу ниже. Вот файл app.py:

# app.py
import psycopg
from flask import Flask, render_template, request
from pgvector.psycopg import register_vector

app = Flask(__name__)

@app.route('/')
def main():
    query = request.args.get('q')
    page = max(0, request.args.get('p', 0, type=int))
    with psycopg.connect(
        'postgres://user:password@localhost:5432/db'
    ) as conn:
        register_vector(conn)
        with conn.cursor() as cur:
            hints = cur.execute(
                'SELECT title FROM movies ORDER BY random() LIMIT 20;'
            )
            if query is not None:
                query = query.strip()
                queryset = cur.execute(
                    """
                    WITH selected_movie AS (
                        SELECT *
                        FROM movies
                        WHERE LOWER(title) = LOWER(%s)
                        LIMIT 1
                    )
                    SELECT
                        m2.*,
                        (SELECT COUNT(*) FROM movies) AS total_count,
                        selected_movie.embedding <-> m2.embedding AS euclidean_distance
                    FROM
                        movies m2,
                        selected_movie
                    ORDER BY
                        euclidean_distance ASC
                    LIMIT 20 OFFSET %s;
                    """,
                    (query, 20 * page),
                )
                result = queryset.fetchall()
                num = result[0][5] if result else 0
                return render_template(
                    'search.html',
                    query=query,
                    result=result,
                    page=page,
                    num=num,
                    hints=hints,
                )
            return render_template('search.html', hints=hints)

Вот файл шаблона страницы на Jinja2 templates/search.html: 

<!DOCTYPE html>
<html lang="ru">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    {% if query %}
        <title>{{ query }} - Рекомендатель кино ({{ num }})</title>
    {% else %}
        <title>Рекомендатель кино</title>
    {% endif %}
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/water.css@2/out/light.min.css">
    <style>
        * {
            box-sizing: border-box;
        }
        body {
            max-width: 960px;
        }
        .app {
            margin-top: 30%;
        }
        .app>h1 {
            font-weight: normal;
            font-size: 4.5rem;
            margin-bottom: 1.5rem;
            text-align: center;
        }
        .app>h1>a,
        .app>h1>a:hover,
        .app>h1>a:active,
        .app>h1>a:focus,
        .app>h1>a:visited {
            color: #46178f;
            text-decoration: none;
        }
        #search-box {
            display: block;
            width: 100%;
            max-width: 700px;
            margin: 0 auto;
            padding: 1.25em;
            border-radius: 8px;
            background-color: hsl(0, 0%, 96%);
            border: none !important;
            outline: none !important;
            font-family: sans-serif;
        }
        #search-box:focus {
            box-shadow: 0px 0px 15px -2px #46178f !important;
        }
        .query-result {
            margin-top: 3em;
            width: 100%;
            max-width: 100%;
            overflow-x: auto;
        }
        .query-result>table {
            width: 100%;
        }
        .query-result td {
            padding-top: 1.5em;
            padding-bottom: 1.5em;
        }
        .query-result td.rating {
            vertical-align: middle;
            text-align: center;
            font-size: 1.5em;
        }
        .query-result th.special {
            text-align: center;
            width: 15%;
        }
        .query-result tr:nth-child(2) {
            background-color: #46178f22;
        }
        .pagination {
            font-size: x-large;
            text-align: center;
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="app">
            <h1><a href="/">Рекомендатель</a></h1>
            <form action="/" method="get">
                <input id="search-box" name="q" type="text" value="{{ query }}">
            </form>
        </div>
        {% if query %}
            {% if result %}
                <div class="query-result">
                    <table>
                        <tr>
                            <th class="special">Рейтинг</th>
                            <th class="special">Постер</th>
                            <th>Описание</th>
                        </tr>
                        {% for movie in result %}
                            <tr>
                                <td class="rating">{{ movie[3] }}</td>
                                <td>
                                    <img loading="lazy" decoding="async" src="https://placedog.net/149/209?id={{ loop.index }}" width="149" height="209" alt="">
                                </td>
                                <td>
                                    <b>{{ movie[1] }}</b>
                                    <p>{{ movie[2] }}</p>
                                    <span>
                                        <a href="/?q={{ movie[1]|urlencode }}">Искать похожие</a>
                                        |
                                        <a href="https://imdb.com/title/{{ movie[0] }}" target="_blank">Страничка на IMDb</a>
                                    </span>
                                </td>
                            </tr>
                        {% endfor %}
                    </table>
                </div>
                <p class="pagination">
                    {% if page and page > 0 %}
                        <a href="/?q={{ query }}&p={{ page - 1 }}">{{ page }}</a>
                    {% endif %}
                
                    {{ page + 1 }}
                
                    {% if (result|length) == 20 %}
                        <a href="/?q={{ query }}&p={{ page + 1 }}">{{ page + 2 }}</a>
                    {% endif %}
                </p>
            {% else %}
                <p>Результатов нет...</p>
            {% endif %}
        {% endif %}
    </div>

    <script type="text/javascript">
        let searchBox = document.getElementById("search-box");
        searchBox.addEventListener("keydown", event => {
            if (event.key != "Enter") return;
            let value = event.srcElement.value;
            if (value.length == 0) {
                event.preventDefault();
                return;
            }
        });
        const examples = [
            {% for hint in hints %}
                "{{ hint[0]|safe }}",
            {% endfor %}
        ].map((example) => example += "...");
        let exampleId = 0;
        let letterId = 0;
        let reversed = false;

        function getRandomInt(max) {
            return Math.floor(Math.random() * max);
        }

        function typewriteExample() {
            if (reversed) {
                setTimeout(typewriteExample, 100 - getRandomInt(25));
                if (letterId-- > 0) {
                    searchBox.placeholder = searchBox.placeholder.slice(0, -1);
                    return;
                }
                reversed = false;
                if (++exampleId >= examples.length) {
                    exampleId = 0;
                }
            } else {
                setTimeout(typewriteExample, 150 + (getRandomInt(150) - 75));
                if (letterId < examples[exampleId].length) {
                    searchBox.placeholder += examples[exampleId].charAt(letterId++);
                    return;
                }
                reversed = true;
            }
        }
        if (examples.length > 0) {
            typewriteExample();
        }
</script>
</body>
</html>

Интерфейс можно запустить и проверить так:

$ flask run

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

Что можно доработать

Рекомендатор кино работает: находить новые фильмы стало проще. Но пока это только MVP.

Для полноты можно:

• Сделать постеры для фильмов. Собаки красивые, но хочется релевантности. 

• Сортировать актёров по рейтингу, выделить главные роли, чтобы кодировать их как категориальные признаки.

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

• Сделать fuzzy search по названию. 

• Сделать поиск по актёрам, жанру, режиссёру, сценаристам. 

• Сделать фасеты для страницы результатов, чтобы фильтровать выдачу. 

Автор статьи: Дмитрий Сидоров

НЛО прилетело и оставило здесь промокод для читателей нашего блога:
-15% на заказ любого VDS (кроме тарифа Прогрев) — HABRFIRSTVDS