Вечер. Пересматриваю «Пятницу 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