Вечер. Пересматриваю «Пятницу 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 pandasimport 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 torchimport 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
ссылка на оригинал статьи https://habr.com/ru/articles/850686/
Добавить комментарий