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

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

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

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

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

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

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

найти датасет с информацией о кино. Чем больше данных, тем лучше;
векторизовать фильмы для сравнения и поиска;
настроить векторную базу данных;
написать интерфейс на 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

Комментарии (1)

mrCOTOHA
15.10.2024 09:28
#27418112
DIW

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