utils.py

# coding=utf-8

import re
from collections import defaultdict
from tqdm import tqdm_notebook as tqdm
import numpy as np
from scipy.sparse import lil_matrix


def load_data(max_films=None):
    movie_id_patt = re.compile('(\d+)\:$')
    user_line_patt = re.compile('(\d+),(\d+),.*$')
    movie_ids = []
    user_ids = set()
    data_train = defaultdict(list)
    data_test = defaultdict(list)
    np.random.seed(20181116)
    for filename in ['./combined_data_1.txt', './combined_data_2.txt', './combined_data_3.txt', './combined_data_4.txt']:
        with open(filename, 'r') as in_file:
            for line in tqdm(in_file):
                match = movie_id_patt.match(line)
                if match:
                    if max_films and len(movie_ids) == max_films:
                        break
                    movie_id = int(match.group(1)) - 1
                    movie_ids.append(movie_id)
                match = user_line_patt.match(line)
                if match:
                    user_id, target = int(match.group(1)), int(match.group(2))
                    if np.random.randint(2) == 0:
                        data_train[movie_id].append((user_id, target))
                    else:
                        data_test[movie_id].append((user_id, target))
                    user_ids.add(user_id)
    movie_ids = np.array(movie_ids)
    user_ids = np.array(sorted(user_ids))
    return movie_ids, user_ids, data_train, data_test


def get_interaction_matrix(data, user_indexes):
    X = lil_matrix((len(data), len(user_indexes)), dtype=np.float32)
    for movie_id, user_ratings in tqdm(data.items()):
        for (user_id, target) in user_ratings:
            X[movie_id, user_indexes[user_id]] = target
    return X


def get_train_test_lil_matrix(max_films=None):
    movie_ids, user_ids, data_train, data_test = load_data(max_films=max_films)
    print(u'Число уникальных фильмов: {}'.format(movie_ids.shape[0]))
    print(u'Число уникальных фильмов в трейне: {}, в тесте: {}'.format(
          len(data_train), len(data_test)))
    print(u'Число уникальных пользователей: {}'.format(user_ids.shape[0]))
    user_indexes = {user_id: index for index, user_id in enumerate(user_ids)}

    X_tr = get_interaction_matrix(data_train, user_indexes)
    print(u'Размер матрицы для обучения: {}'.format(X_tr.shape))
    X_te = get_interaction_matrix(data_test, user_indexes)
    print(u'Размер матрицы для валидации: {}'.format(X_te.shape))
    return X_tr, X_te


def get_train_test(max_films=None):
    X_tr, X_te = get_train_test_lil_matrix(max_films=max_films)
    X_tr = X_tr.tocsr()
    X_te = X_te.tocsr()
    return X_tr, X_te