wehub-resource-sync
790 lines
27 KiB
Python
790 lines
27 KiB
Python
"""MovieLens dataset"""
|
|
import os
|
|
import re
|
|
|
|
import dgl
|
|
import numpy as np
|
|
import pandas as pd
|
|
import scipy.sparse as sp
|
|
import torch as th
|
|
from dgl.data.utils import download, extract_archive, get_download_dir
|
|
from utils import to_etype_name
|
|
|
|
_urls = {
|
|
"ml-100k": "http://files.grouplens.org/datasets/movielens/ml-100k.zip",
|
|
"ml-1m": "http://files.grouplens.org/datasets/movielens/ml-1m.zip",
|
|
"ml-10m": "http://files.grouplens.org/datasets/movielens/ml-10m.zip",
|
|
}
|
|
|
|
READ_DATASET_PATH = get_download_dir()
|
|
GENRES_ML_100K = [
|
|
"unknown",
|
|
"Action",
|
|
"Adventure",
|
|
"Animation",
|
|
"Children",
|
|
"Comedy",
|
|
"Crime",
|
|
"Documentary",
|
|
"Drama",
|
|
"Fantasy",
|
|
"Film-Noir",
|
|
"Horror",
|
|
"Musical",
|
|
"Mystery",
|
|
"Romance",
|
|
"Sci-Fi",
|
|
"Thriller",
|
|
"War",
|
|
"Western",
|
|
]
|
|
GENRES_ML_1M = GENRES_ML_100K[1:]
|
|
GENRES_ML_10M = GENRES_ML_100K + ["IMAX"]
|
|
|
|
|
|
class MovieLens(object):
|
|
"""MovieLens dataset used by GCMC model
|
|
|
|
TODO(minjie): make this dataset more general
|
|
|
|
The dataset stores MovieLens ratings in two types of graphs. The encoder graph
|
|
contains rating value information in the form of edge types. The decoder graph
|
|
stores plain user-movie pairs in the form of a bipartite graph with no rating
|
|
information. All graphs have two types of nodes: "user" and "movie".
|
|
|
|
The training, validation and test set can be summarized as follows:
|
|
|
|
training_enc_graph : training user-movie pairs + rating info
|
|
training_dec_graph : training user-movie pairs
|
|
valid_enc_graph : training user-movie pairs + rating info
|
|
valid_dec_graph : validation user-movie pairs
|
|
test_enc_graph : training user-movie pairs + validation user-movie pairs + rating info
|
|
test_dec_graph : test user-movie pairs
|
|
|
|
Attributes
|
|
----------
|
|
train_enc_graph : dgl.DGLGraph
|
|
Encoder graph for training.
|
|
train_dec_graph : dgl.DGLGraph
|
|
Decoder graph for training.
|
|
train_labels : torch.Tensor
|
|
The categorical label of each user-movie pair
|
|
train_truths : torch.Tensor
|
|
The actual rating values of each user-movie pair
|
|
valid_enc_graph : dgl.DGLGraph
|
|
Encoder graph for validation.
|
|
valid_dec_graph : dgl.DGLGraph
|
|
Decoder graph for validation.
|
|
valid_labels : torch.Tensor
|
|
The categorical label of each user-movie pair
|
|
valid_truths : torch.Tensor
|
|
The actual rating values of each user-movie pair
|
|
test_enc_graph : dgl.DGLGraph
|
|
Encoder graph for test.
|
|
test_dec_graph : dgl.DGLGraph
|
|
Decoder graph for test.
|
|
test_labels : torch.Tensor
|
|
The categorical label of each user-movie pair
|
|
test_truths : torch.Tensor
|
|
The actual rating values of each user-movie pair
|
|
user_feature : torch.Tensor
|
|
User feature tensor. If None, representing an identity matrix.
|
|
movie_feature : torch.Tensor
|
|
Movie feature tensor. If None, representing an identity matrix.
|
|
possible_rating_values : np.ndarray
|
|
Available rating values in the dataset
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
Dataset name. Could be "ml-100k", "ml-1m", "ml-10m"
|
|
device : torch.device
|
|
Device context
|
|
mix_cpu_gpu : boo, optional
|
|
If true, the ``user_feature`` attribute is stored in CPU
|
|
use_one_hot_fea : bool, optional
|
|
If true, the ``user_feature`` attribute is None, representing an one-hot identity
|
|
matrix. (Default: False)
|
|
symm : bool, optional
|
|
If true, the use symmetric normalize constant. Otherwise, use left normalize
|
|
constant. (Default: True)
|
|
test_ratio : float, optional
|
|
Ratio of test data
|
|
valid_ratio : float, optional
|
|
Ratio of validation data
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name,
|
|
device,
|
|
mix_cpu_gpu=False,
|
|
use_one_hot_fea=False,
|
|
symm=True,
|
|
test_ratio=0.1,
|
|
valid_ratio=0.1,
|
|
):
|
|
self._name = name
|
|
self._device = device
|
|
self._symm = symm
|
|
self._test_ratio = test_ratio
|
|
self._valid_ratio = valid_ratio
|
|
# download and extract
|
|
download_dir = get_download_dir()
|
|
zip_file_path = "{}/{}.zip".format(download_dir, name)
|
|
download(_urls[name], path=zip_file_path)
|
|
extract_archive(zip_file_path, "{}/{}".format(download_dir, name))
|
|
if name == "ml-10m":
|
|
root_folder = "ml-10M100K"
|
|
else:
|
|
root_folder = name
|
|
self._dir = os.path.join(download_dir, name, root_folder)
|
|
print("Starting processing {} ...".format(self._name))
|
|
self._load_raw_user_info()
|
|
self._load_raw_movie_info()
|
|
print("......")
|
|
if self._name == "ml-100k":
|
|
self.all_train_rating_info = self._load_raw_rates(
|
|
os.path.join(self._dir, "u1.base"), "\t"
|
|
)
|
|
self.test_rating_info = self._load_raw_rates(
|
|
os.path.join(self._dir, "u1.test"), "\t"
|
|
)
|
|
self.all_rating_info = pd.concat(
|
|
[self.all_train_rating_info, self.test_rating_info]
|
|
)
|
|
elif self._name == "ml-1m" or self._name == "ml-10m":
|
|
self.all_rating_info = self._load_raw_rates(
|
|
os.path.join(self._dir, "ratings.dat"), "::"
|
|
)
|
|
num_test = int(
|
|
np.ceil(self.all_rating_info.shape[0] * self._test_ratio)
|
|
)
|
|
shuffled_idx = np.random.permutation(self.all_rating_info.shape[0])
|
|
self.test_rating_info = self.all_rating_info.iloc[
|
|
shuffled_idx[:num_test]
|
|
]
|
|
self.all_train_rating_info = self.all_rating_info.iloc[
|
|
shuffled_idx[num_test:]
|
|
]
|
|
else:
|
|
raise NotImplementedError
|
|
print("......")
|
|
num_valid = int(
|
|
np.ceil(self.all_train_rating_info.shape[0] * self._valid_ratio)
|
|
)
|
|
shuffled_idx = np.random.permutation(
|
|
self.all_train_rating_info.shape[0]
|
|
)
|
|
self.valid_rating_info = self.all_train_rating_info.iloc[
|
|
shuffled_idx[:num_valid]
|
|
]
|
|
self.train_rating_info = self.all_train_rating_info.iloc[
|
|
shuffled_idx[num_valid:]
|
|
]
|
|
self.possible_rating_values = np.unique(
|
|
self.train_rating_info["rating"].values
|
|
)
|
|
|
|
print("All rating pairs : {}".format(self.all_rating_info.shape[0]))
|
|
print(
|
|
"\tAll train rating pairs : {}".format(
|
|
self.all_train_rating_info.shape[0]
|
|
)
|
|
)
|
|
print(
|
|
"\t\tTrain rating pairs : {}".format(
|
|
self.train_rating_info.shape[0]
|
|
)
|
|
)
|
|
print(
|
|
"\t\tValid rating pairs : {}".format(
|
|
self.valid_rating_info.shape[0]
|
|
)
|
|
)
|
|
print(
|
|
"\tTest rating pairs : {}".format(self.test_rating_info.shape[0])
|
|
)
|
|
|
|
self.user_info = self._drop_unseen_nodes(
|
|
orign_info=self.user_info,
|
|
cmp_col_name="id",
|
|
reserved_ids_set=set(self.all_rating_info["user_id"].values),
|
|
label="user",
|
|
)
|
|
self.movie_info = self._drop_unseen_nodes(
|
|
orign_info=self.movie_info,
|
|
cmp_col_name="id",
|
|
reserved_ids_set=set(self.all_rating_info["movie_id"].values),
|
|
label="movie",
|
|
)
|
|
|
|
# Map user/movie to the global id
|
|
self.global_user_id_map = {
|
|
ele: i for i, ele in enumerate(self.user_info["id"])
|
|
}
|
|
self.global_movie_id_map = {
|
|
ele: i for i, ele in enumerate(self.movie_info["id"])
|
|
}
|
|
print(
|
|
"Total user number = {}, movie number = {}".format(
|
|
len(self.global_user_id_map), len(self.global_movie_id_map)
|
|
)
|
|
)
|
|
self._num_user = len(self.global_user_id_map)
|
|
self._num_movie = len(self.global_movie_id_map)
|
|
|
|
### Generate features
|
|
if use_one_hot_fea:
|
|
self.user_feature = None
|
|
self.movie_feature = None
|
|
else:
|
|
# if mix_cpu_gpu, we put features in CPU
|
|
if mix_cpu_gpu:
|
|
self.user_feature = th.FloatTensor(self._process_user_fea())
|
|
self.movie_feature = th.FloatTensor(self._process_movie_fea())
|
|
else:
|
|
self.user_feature = th.FloatTensor(self._process_user_fea()).to(
|
|
self._device
|
|
)
|
|
self.movie_feature = th.FloatTensor(
|
|
self._process_movie_fea()
|
|
).to(self._device)
|
|
if self.user_feature is None:
|
|
self.user_feature_shape = (self.num_user, self.num_user)
|
|
self.movie_feature_shape = (self.num_movie, self.num_movie)
|
|
else:
|
|
self.user_feature_shape = self.user_feature.shape
|
|
self.movie_feature_shape = self.movie_feature.shape
|
|
info_line = "Feature dim: "
|
|
info_line += "\nuser: {}".format(self.user_feature_shape)
|
|
info_line += "\nmovie: {}".format(self.movie_feature_shape)
|
|
print(info_line)
|
|
|
|
(
|
|
all_train_rating_pairs,
|
|
all_train_rating_values,
|
|
) = self._generate_pair_value(self.all_train_rating_info)
|
|
train_rating_pairs, train_rating_values = self._generate_pair_value(
|
|
self.train_rating_info
|
|
)
|
|
valid_rating_pairs, valid_rating_values = self._generate_pair_value(
|
|
self.valid_rating_info
|
|
)
|
|
test_rating_pairs, test_rating_values = self._generate_pair_value(
|
|
self.test_rating_info
|
|
)
|
|
|
|
def _make_labels(ratings):
|
|
labels = th.LongTensor(
|
|
np.searchsorted(self.possible_rating_values, ratings)
|
|
).to(device)
|
|
return labels
|
|
|
|
self.train_enc_graph = self._generate_enc_graph(
|
|
train_rating_pairs, train_rating_values, add_support=True
|
|
)
|
|
self.train_dec_graph = self._generate_dec_graph(train_rating_pairs)
|
|
self.train_labels = _make_labels(train_rating_values)
|
|
self.train_truths = th.FloatTensor(train_rating_values).to(device)
|
|
|
|
self.valid_enc_graph = self.train_enc_graph
|
|
self.valid_dec_graph = self._generate_dec_graph(valid_rating_pairs)
|
|
self.valid_labels = _make_labels(valid_rating_values)
|
|
self.valid_truths = th.FloatTensor(valid_rating_values).to(device)
|
|
|
|
self.test_enc_graph = self._generate_enc_graph(
|
|
all_train_rating_pairs, all_train_rating_values, add_support=True
|
|
)
|
|
self.test_dec_graph = self._generate_dec_graph(test_rating_pairs)
|
|
self.test_labels = _make_labels(test_rating_values)
|
|
self.test_truths = th.FloatTensor(test_rating_values).to(device)
|
|
|
|
def _npairs(graph):
|
|
rst = 0
|
|
for r in self.possible_rating_values:
|
|
r = to_etype_name(r)
|
|
rst += graph.num_edges(str(r))
|
|
return rst
|
|
|
|
print(
|
|
"Train enc graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.train_enc_graph.num_nodes("user"),
|
|
self.train_enc_graph.num_nodes("movie"),
|
|
_npairs(self.train_enc_graph),
|
|
)
|
|
)
|
|
print(
|
|
"Train dec graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.train_dec_graph.num_nodes("user"),
|
|
self.train_dec_graph.num_nodes("movie"),
|
|
self.train_dec_graph.num_edges(),
|
|
)
|
|
)
|
|
print(
|
|
"Valid enc graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.valid_enc_graph.num_nodes("user"),
|
|
self.valid_enc_graph.num_nodes("movie"),
|
|
_npairs(self.valid_enc_graph),
|
|
)
|
|
)
|
|
print(
|
|
"Valid dec graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.valid_dec_graph.num_nodes("user"),
|
|
self.valid_dec_graph.num_nodes("movie"),
|
|
self.valid_dec_graph.num_edges(),
|
|
)
|
|
)
|
|
print(
|
|
"Test enc graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.test_enc_graph.num_nodes("user"),
|
|
self.test_enc_graph.num_nodes("movie"),
|
|
_npairs(self.test_enc_graph),
|
|
)
|
|
)
|
|
print(
|
|
"Test dec graph: \t#user:{}\t#movie:{}\t#pairs:{}".format(
|
|
self.test_dec_graph.num_nodes("user"),
|
|
self.test_dec_graph.num_nodes("movie"),
|
|
self.test_dec_graph.num_edges(),
|
|
)
|
|
)
|
|
|
|
def _generate_pair_value(self, rating_info):
|
|
rating_pairs = (
|
|
np.array(
|
|
[
|
|
self.global_user_id_map[ele]
|
|
for ele in rating_info["user_id"]
|
|
],
|
|
dtype=np.int64,
|
|
),
|
|
np.array(
|
|
[
|
|
self.global_movie_id_map[ele]
|
|
for ele in rating_info["movie_id"]
|
|
],
|
|
dtype=np.int64,
|
|
),
|
|
)
|
|
rating_values = rating_info["rating"].values.astype(np.float32)
|
|
return rating_pairs, rating_values
|
|
|
|
def _generate_enc_graph(
|
|
self, rating_pairs, rating_values, add_support=False
|
|
):
|
|
user_movie_R = np.zeros(
|
|
(self._num_user, self._num_movie), dtype=np.float32
|
|
)
|
|
user_movie_R[rating_pairs] = rating_values
|
|
|
|
data_dict = dict()
|
|
num_nodes_dict = {"user": self._num_user, "movie": self._num_movie}
|
|
rating_row, rating_col = rating_pairs
|
|
for rating in self.possible_rating_values:
|
|
ridx = np.where(rating_values == rating)
|
|
rrow = rating_row[ridx]
|
|
rcol = rating_col[ridx]
|
|
rating = to_etype_name(rating)
|
|
data_dict.update(
|
|
{
|
|
("user", str(rating), "movie"): (rrow, rcol),
|
|
("movie", "rev-%s" % str(rating), "user"): (rcol, rrow),
|
|
}
|
|
)
|
|
graph = dgl.heterograph(data_dict, num_nodes_dict=num_nodes_dict)
|
|
|
|
# sanity check
|
|
assert (
|
|
len(rating_pairs[0])
|
|
== sum([graph.num_edges(et) for et in graph.etypes]) // 2
|
|
)
|
|
|
|
if add_support:
|
|
|
|
def _calc_norm(x):
|
|
x = x.numpy().astype("float32")
|
|
x[x == 0.0] = np.inf
|
|
x = th.FloatTensor(1.0 / np.sqrt(x))
|
|
return x.unsqueeze(1)
|
|
|
|
user_ci = []
|
|
user_cj = []
|
|
movie_ci = []
|
|
movie_cj = []
|
|
for r in self.possible_rating_values:
|
|
r = to_etype_name(r)
|
|
user_ci.append(graph["rev-%s" % r].in_degrees())
|
|
movie_ci.append(graph[r].in_degrees())
|
|
if self._symm:
|
|
user_cj.append(graph[r].out_degrees())
|
|
movie_cj.append(graph["rev-%s" % r].out_degrees())
|
|
else:
|
|
user_cj.append(th.zeros((self.num_user,)))
|
|
movie_cj.append(th.zeros((self.num_movie,)))
|
|
user_ci = _calc_norm(sum(user_ci))
|
|
movie_ci = _calc_norm(sum(movie_ci))
|
|
if self._symm:
|
|
user_cj = _calc_norm(sum(user_cj))
|
|
movie_cj = _calc_norm(sum(movie_cj))
|
|
else:
|
|
user_cj = th.ones(
|
|
self.num_user,
|
|
)
|
|
movie_cj = th.ones(
|
|
self.num_movie,
|
|
)
|
|
graph.nodes["user"].data.update({"ci": user_ci, "cj": user_cj})
|
|
graph.nodes["movie"].data.update({"ci": movie_ci, "cj": movie_cj})
|
|
|
|
return graph
|
|
|
|
def _generate_dec_graph(self, rating_pairs):
|
|
ones = np.ones_like(rating_pairs[0])
|
|
user_movie_ratings_coo = sp.coo_matrix(
|
|
(ones, rating_pairs),
|
|
shape=(self.num_user, self.num_movie),
|
|
dtype=np.float32,
|
|
)
|
|
g = dgl.bipartite_from_scipy(
|
|
user_movie_ratings_coo, utype="_U", etype="_E", vtype="_V"
|
|
)
|
|
return dgl.heterograph(
|
|
{("user", "rate", "movie"): g.edges()},
|
|
num_nodes_dict={"user": self.num_user, "movie": self.num_movie},
|
|
)
|
|
|
|
@property
|
|
def num_links(self):
|
|
return self.possible_rating_values.size
|
|
|
|
@property
|
|
def num_user(self):
|
|
return self._num_user
|
|
|
|
@property
|
|
def num_movie(self):
|
|
return self._num_movie
|
|
|
|
def _drop_unseen_nodes(
|
|
self, orign_info, cmp_col_name, reserved_ids_set, label
|
|
):
|
|
# print(" -----------------")
|
|
# print("{}: {}(reserved) v.s. {}(from info)".format(label, len(reserved_ids_set),
|
|
# len(set(orign_info[cmp_col_name].values))))
|
|
if reserved_ids_set != set(orign_info[cmp_col_name].values):
|
|
pd_rating_ids = pd.DataFrame(
|
|
list(reserved_ids_set), columns=["id_graph"]
|
|
)
|
|
# print("\torign_info: ({}, {})".format(orign_info.shape[0], orign_info.shape[1]))
|
|
data_info = orign_info.merge(
|
|
pd_rating_ids,
|
|
left_on=cmp_col_name,
|
|
right_on="id_graph",
|
|
how="outer",
|
|
)
|
|
data_info = data_info.dropna(subset=[cmp_col_name, "id_graph"])
|
|
data_info = data_info.drop(columns=["id_graph"])
|
|
data_info = data_info.reset_index(drop=True)
|
|
# print("\tAfter dropping, data shape: ({}, {})".format(data_info.shape[0], data_info.shape[1]))
|
|
return data_info
|
|
else:
|
|
orign_info = orign_info.reset_index(drop=True)
|
|
return orign_info
|
|
|
|
def _load_raw_rates(self, file_path, sep):
|
|
"""In MovieLens, the rates have the following format
|
|
|
|
ml-100k
|
|
user id \t movie id \t rating \t timestamp
|
|
|
|
ml-1m/10m
|
|
UserID::MovieID::Rating::Timestamp
|
|
|
|
timestamp is unix timestamp and can be converted by pd.to_datetime(X, unit='s')
|
|
|
|
Parameters
|
|
----------
|
|
file_path : str
|
|
|
|
Returns
|
|
-------
|
|
rating_info : pd.DataFrame
|
|
"""
|
|
rating_info = pd.read_csv(
|
|
file_path,
|
|
sep=sep,
|
|
header=None,
|
|
names=["user_id", "movie_id", "rating", "timestamp"],
|
|
dtype={
|
|
"user_id": np.int32,
|
|
"movie_id": np.int32,
|
|
"ratings": np.float32,
|
|
"timestamp": np.int64,
|
|
},
|
|
engine="python",
|
|
)
|
|
return rating_info
|
|
|
|
def _load_raw_user_info(self):
|
|
"""In MovieLens, the user attributes file have the following formats:
|
|
|
|
ml-100k:
|
|
user id | age | gender | occupation | zip code
|
|
|
|
ml-1m:
|
|
UserID::Gender::Age::Occupation::Zip-code
|
|
|
|
For ml-10m, there is no user information. We read the user id from the rating file.
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
|
|
Returns
|
|
-------
|
|
user_info : pd.DataFrame
|
|
"""
|
|
if self._name == "ml-100k":
|
|
self.user_info = pd.read_csv(
|
|
os.path.join(self._dir, "u.user"),
|
|
sep="|",
|
|
header=None,
|
|
names=["id", "age", "gender", "occupation", "zip_code"],
|
|
engine="python",
|
|
)
|
|
elif self._name == "ml-1m":
|
|
self.user_info = pd.read_csv(
|
|
os.path.join(self._dir, "users.dat"),
|
|
sep="::",
|
|
header=None,
|
|
names=["id", "gender", "age", "occupation", "zip_code"],
|
|
engine="python",
|
|
)
|
|
elif self._name == "ml-10m":
|
|
rating_info = pd.read_csv(
|
|
os.path.join(self._dir, "ratings.dat"),
|
|
sep="::",
|
|
header=None,
|
|
names=["user_id", "movie_id", "rating", "timestamp"],
|
|
dtype={
|
|
"user_id": np.int32,
|
|
"movie_id": np.int32,
|
|
"ratings": np.float32,
|
|
"timestamp": np.int64,
|
|
},
|
|
engine="python",
|
|
)
|
|
self.user_info = pd.DataFrame(
|
|
np.unique(rating_info["user_id"].values.astype(np.int32)),
|
|
columns=["id"],
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def _process_user_fea(self):
|
|
"""
|
|
|
|
Parameters
|
|
----------
|
|
user_info : pd.DataFrame
|
|
name : str
|
|
For ml-100k and ml-1m, the column name is ['id', 'gender', 'age', 'occupation', 'zip_code'].
|
|
We take the age, gender, and the one-hot encoding of the occupation as the user features.
|
|
For ml-10m, there is no user feature and we set the feature to be a single zero.
|
|
|
|
Returns
|
|
-------
|
|
user_features : np.ndarray
|
|
|
|
"""
|
|
if self._name == "ml-100k" or self._name == "ml-1m":
|
|
ages = self.user_info["age"].values.astype(np.float32)
|
|
gender = (self.user_info["gender"] == "F").values.astype(np.float32)
|
|
all_occupations = set(self.user_info["occupation"])
|
|
occupation_map = {ele: i for i, ele in enumerate(all_occupations)}
|
|
occupation_one_hot = np.zeros(
|
|
shape=(self.user_info.shape[0], len(all_occupations)),
|
|
dtype=np.float32,
|
|
)
|
|
occupation_one_hot[
|
|
np.arange(self.user_info.shape[0]),
|
|
np.array(
|
|
[
|
|
occupation_map[ele]
|
|
for ele in self.user_info["occupation"]
|
|
]
|
|
),
|
|
] = 1
|
|
user_features = np.concatenate(
|
|
[
|
|
ages.reshape((self.user_info.shape[0], 1)) / 50.0,
|
|
gender.reshape((self.user_info.shape[0], 1)),
|
|
occupation_one_hot,
|
|
],
|
|
axis=1,
|
|
)
|
|
elif self._name == "ml-10m":
|
|
user_features = np.zeros(
|
|
shape=(self.user_info.shape[0], 1), dtype=np.float32
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
return user_features
|
|
|
|
def _load_raw_movie_info(self):
|
|
"""In MovieLens, the movie attributes may have the following formats:
|
|
|
|
In ml_100k:
|
|
|
|
movie id | movie title | release date | video release date | IMDb URL | [genres]
|
|
|
|
In ml_1m, ml_10m:
|
|
|
|
MovieID::Title (Release Year)::Genres
|
|
|
|
Also, Genres are separated by |, e.g., Adventure|Animation|Children|Comedy|Fantasy
|
|
|
|
Parameters
|
|
----------
|
|
name : str
|
|
|
|
Returns
|
|
-------
|
|
movie_info : pd.DataFrame
|
|
For ml-100k, the column name is ['id', 'title', 'release_date', 'video_release_date', 'url'] + [GENRES (19)]]
|
|
For ml-1m and ml-10m, the column name is ['id', 'title'] + [GENRES (18/20)]]
|
|
"""
|
|
if self._name == "ml-100k":
|
|
GENRES = GENRES_ML_100K
|
|
elif self._name == "ml-1m":
|
|
GENRES = GENRES_ML_1M
|
|
elif self._name == "ml-10m":
|
|
GENRES = GENRES_ML_10M
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
if self._name == "ml-100k":
|
|
file_path = os.path.join(self._dir, "u.item")
|
|
self.movie_info = pd.read_csv(
|
|
file_path,
|
|
sep="|",
|
|
header=None,
|
|
names=[
|
|
"id",
|
|
"title",
|
|
"release_date",
|
|
"video_release_date",
|
|
"url",
|
|
]
|
|
+ GENRES,
|
|
encoding="iso-8859-1",
|
|
)
|
|
elif self._name == "ml-1m" or self._name == "ml-10m":
|
|
file_path = os.path.join(self._dir, "movies.dat")
|
|
movie_info = pd.read_csv(
|
|
file_path,
|
|
sep="::",
|
|
header=None,
|
|
names=["id", "title", "genres"],
|
|
encoding="iso-8859-1",
|
|
engine="python",
|
|
)
|
|
genre_map = {ele: i for i, ele in enumerate(GENRES)}
|
|
genre_map["Children's"] = genre_map["Children"]
|
|
genre_map["Childrens"] = genre_map["Children"]
|
|
movie_genres = np.zeros(
|
|
shape=(movie_info.shape[0], len(GENRES)), dtype=np.float32
|
|
)
|
|
for i, genres in enumerate(movie_info["genres"]):
|
|
for ele in genres.split("|"):
|
|
if ele in genre_map:
|
|
movie_genres[i, genre_map[ele]] = 1.0
|
|
else:
|
|
print(
|
|
"genres not found, filled with unknown: {}".format(
|
|
genres
|
|
)
|
|
)
|
|
movie_genres[i, genre_map["unknown"]] = 1.0
|
|
for idx, genre_name in enumerate(GENRES):
|
|
assert idx == genre_map[genre_name]
|
|
movie_info[genre_name] = movie_genres[:, idx]
|
|
self.movie_info = movie_info.drop(columns=["genres"])
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
def _process_movie_fea(self):
|
|
"""
|
|
|
|
Parameters
|
|
----------
|
|
movie_info : pd.DataFrame
|
|
name : str
|
|
|
|
Returns
|
|
-------
|
|
movie_features : np.ndarray
|
|
Generate movie features by concatenating embedding and the year
|
|
|
|
"""
|
|
import torchtext
|
|
from torchtext.data.utils import get_tokenizer
|
|
|
|
if self._name == "ml-100k":
|
|
GENRES = GENRES_ML_100K
|
|
elif self._name == "ml-1m":
|
|
GENRES = GENRES_ML_1M
|
|
elif self._name == "ml-10m":
|
|
GENRES = GENRES_ML_10M
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# Old torchtext-legacy API commented below
|
|
# TEXT = torchtext.legacy.data.Field(tokenize='spacy', tokenizer_language='en_core_web_sm')
|
|
tokenizer = get_tokenizer(
|
|
"spacy", language="en_core_web_sm"
|
|
) # new API (torchtext 0.9+)
|
|
embedding = torchtext.vocab.GloVe(name="840B", dim=300)
|
|
|
|
title_embedding = np.zeros(
|
|
shape=(self.movie_info.shape[0], 300), dtype=np.float32
|
|
)
|
|
release_years = np.zeros(
|
|
shape=(self.movie_info.shape[0], 1), dtype=np.float32
|
|
)
|
|
p = re.compile(r"(.+)\s*\((\d+)\)")
|
|
for i, title in enumerate(self.movie_info["title"]):
|
|
match_res = p.match(title)
|
|
if match_res is None:
|
|
print(
|
|
"{} cannot be matched, index={}, name={}".format(
|
|
title, i, self._name
|
|
)
|
|
)
|
|
title_context, year = title, 1950
|
|
else:
|
|
title_context, year = match_res.groups()
|
|
# We use average of glove
|
|
# Upgraded torchtext API: TEXT.tokenize(title_context) --> tokenizer(title_context)
|
|
title_embedding[i, :] = (
|
|
embedding.get_vecs_by_tokens(tokenizer(title_context))
|
|
.numpy()
|
|
.mean(axis=0)
|
|
)
|
|
release_years[i] = float(year)
|
|
movie_features = np.concatenate(
|
|
(
|
|
title_embedding,
|
|
(release_years - 1950.0) / 100.0,
|
|
self.movie_info[GENRES],
|
|
),
|
|
axis=1,
|
|
)
|
|
return movie_features
|
|
|
|
|
|
if __name__ == "__main__":
|
|
MovieLens("ml-100k", device=th.device("cpu"), symm=True)
|