wehub-resource-sync
647 lines
22 KiB
Python
647 lines
22 KiB
Python
"""MovieLens dataset"""
|
|
import os
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
from torch import LongTensor, Tensor
|
|
|
|
from ..base import dgl_warning
|
|
from ..convert import heterograph
|
|
from .dgl_dataset import DGLDataset
|
|
|
|
from .utils import (
|
|
_get_dgl_url,
|
|
download,
|
|
extract_archive,
|
|
load_graphs,
|
|
load_info,
|
|
save_graphs,
|
|
save_info,
|
|
split_dataset,
|
|
)
|
|
|
|
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"]
|
|
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
HAS_TORCH = False
|
|
else:
|
|
HAS_TORCH = True
|
|
|
|
|
|
def check_pytorch():
|
|
"""Check if PyTorch is the backend."""
|
|
if not HAS_TORCH:
|
|
raise ModuleNotFoundError(
|
|
"MovieLensDataset requires PyTorch to be the backend."
|
|
)
|
|
|
|
|
|
class MovieLensDataset(DGLDataset):
|
|
r"""MovieLens dataset for edge prediction tasks. The raw datasets are extracted from
|
|
`MovieLens <https://grouplens.org/datasets/movielens/>`, introduced by
|
|
`Movielens unplugged: experiences with an occasionally connected recommender system <https://dl.acm.org/doi/10.1145/604045.604094>`.
|
|
|
|
The datasets consist of user ratings for movies and incorporate additional user/movie information in the form of features.
|
|
The nodes represent users and movies, and the edges store ratings that users assign to movies.
|
|
|
|
Statistics:
|
|
|
|
MovieLens-100K (ml-100k)
|
|
|
|
- Users: 943
|
|
- Movies: 1,682
|
|
- Ratings: 100,000 (1, 2, 3, 4, 5)
|
|
|
|
MovieLens-1M (ml-1m)
|
|
|
|
- Users: 6,040
|
|
- Movies: 3,706
|
|
- Ratings: 1,000,209 (1, 2, 3, 4, 5)
|
|
|
|
MovieLens-10M (ml-10m)
|
|
|
|
- Users: 69,878
|
|
- Movies: 10,677
|
|
- Ratings: 10,000,054 (0.5, 1, 1.5, ..., 4.5, 5.0)
|
|
|
|
Parameters
|
|
----------
|
|
name: str
|
|
Dataset name. (:obj:`"ml-100k"`, :obj:`"ml-1m"`, :obj:`"ml-10m"`).
|
|
valid_ratio: int
|
|
Ratio of validation samples out of the whole dataset. Should be in (0.0, 1.0).
|
|
test_ratio: int, optional
|
|
Ratio of testing samples out of the whole dataset. Should be in (0.0, 1.0). And its sum with
|
|
:obj:`valid_ratio` should be in (0.0, 1.0) as well. This parameter is invalid
|
|
when :obj:`name` is :obj:`"ml-100k"`, since its testing samples are pre-specified.
|
|
Default: None
|
|
raw_dir : str, optional
|
|
Raw file directory to download/store the data.
|
|
Default: ~/.dgl/
|
|
force_reload : bool, optional
|
|
Whether to re-download(if the dataset has not been downloaded) and re-process the dataset.
|
|
Default: False
|
|
verbose : bool, optional
|
|
Whether to print progress information. Default: True.
|
|
transform : callable, optional
|
|
A transform that takes in a :class:`~dgl.DGLGraph` object and returns
|
|
a transformed version. The :class:`~dgl.DGLGraph` object will be
|
|
transformed before every access.
|
|
random_state : int, optional
|
|
Random seed used for random dataset split. Default: 0
|
|
|
|
Notes
|
|
-----
|
|
- When :obj:`name` is :obj:`"ml-100k"`, the :obj:`test_ratio` is invalid, and the training ratio is equal to 1-:obj:`valid_ratio`.
|
|
When :obj:`name` is :obj:`"ml-1m"` or :obj:`"ml-10m"`, the :obj:`test_ratio` is valid,
|
|
and the training ratio is equal to 1-:obj:`valid_ratio`-:obj:`test_ratio`.
|
|
- The number of edges is doubled to form an undirected(bidirected) graph structure.
|
|
|
|
Examples
|
|
--------
|
|
>>> from dgl.data import MovieLensDataset
|
|
>>> dataset = MovieLensDataset(name='ml-100k', valid_ratio=0.2)
|
|
>>> g = dataset[0]
|
|
>>> g
|
|
Graph(num_nodes={'movie': 1682, 'user': 943},
|
|
num_edges={('movie', 'movie-user', 'user'): 100000, ('user', 'user-movie', 'movie'): 100000},
|
|
metagraph=[('movie', 'user', 'movie-user'), ('user', 'movie', 'user-movie')])
|
|
|
|
>>> # get ratings of edges in the training graph.
|
|
>>> rate = g.edges['user-movie'].data['rate'] # or rate = g.edges['movie-user'].data['rate']
|
|
>>> rate
|
|
tensor([5., 5., 3., ..., 3., 3., 5.])
|
|
|
|
>>> # get train, valid and test mask of edges
|
|
>>> train_mask = g.edges['user-movie'].data['train_mask']
|
|
>>> valid_mask = g.edges['user-movie'].data['valid_mask']
|
|
>>> test_mask = g.edges['user-movie'].data['test_mask']
|
|
|
|
>>> # get train, valid and test ratings
|
|
>>> train_ratings = rate[train_mask]
|
|
>>> valid_ratings = rate[valid_mask]
|
|
>>> test_ratings = rate[test_mask]
|
|
|
|
>>> # get input features of users
|
|
>>> g.nodes["user"].data["feat"] # or g.nodes["movie"].data["feat"] for movie nodes
|
|
tensor([[0.4800, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
|
|
[1.0600, 1.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
|
|
[0.4600, 0.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
|
|
...,
|
|
[0.4000, 0.0000, 1.0000, ..., 0.0000, 0.0000, 0.0000],
|
|
[0.9600, 1.0000, 0.0000, ..., 0.0000, 0.0000, 0.0000],
|
|
[0.4400, 0.0000, 1.0000, ..., 0.0000, 0.0000, 0.0000]])
|
|
|
|
"""
|
|
|
|
_url = {
|
|
"ml-100k": "dataset/ml-100k.zip",
|
|
"ml-1m": "dataset/ml-1m.zip",
|
|
"ml-10m": "dataset/ml-10m.zip",
|
|
}
|
|
|
|
def __init__(
|
|
self,
|
|
name,
|
|
valid_ratio,
|
|
test_ratio=None,
|
|
raw_dir=None,
|
|
force_reload=None,
|
|
verbose=None,
|
|
transform=None,
|
|
random_state=0,
|
|
):
|
|
check_pytorch()
|
|
assert name in [
|
|
"ml-100k",
|
|
"ml-1m",
|
|
"ml-10m",
|
|
], f"currently movielens does not support {name}"
|
|
|
|
# test regarding valid and test split ratio
|
|
assert (
|
|
valid_ratio > 0.0 and valid_ratio < 1.0
|
|
), f"valid_ratio {valid_ratio} must be in (0.0, 1.0)"
|
|
|
|
if name in ["ml-1m", "ml-10m"]:
|
|
assert (
|
|
test_ratio is not None and test_ratio > 0.0 and test_ratio < 1.0
|
|
), f"test_ratio({test_ratio}) must be set to a value in (0.0, 1.0) when using ml-1m and ml-10m"
|
|
assert (
|
|
test_ratio + valid_ratio > 0.0
|
|
and test_ratio + valid_ratio < 1.0
|
|
), f"test_ratio({test_ratio}) + valid_ratio({valid_ratio}) must be set to (0.0, 1.0) when using ml-1m and ml-10m"
|
|
|
|
if name == "ml-100k" and test_ratio is not None:
|
|
dgl_warning(
|
|
f"test_ratio ({test_ratio}) is not set to None for ml-100k. "
|
|
"Note that dataset split would not be affected by the test_ratio since "
|
|
"testing samples of ml-100k have been pre-specified."
|
|
)
|
|
|
|
self.valid_ratio = valid_ratio
|
|
self.test_ratio = test_ratio
|
|
self.random_state = random_state
|
|
|
|
if name == "ml-100k":
|
|
self.genres = GENRES_ML_100K
|
|
elif name == "ml-1m":
|
|
self.genres = GENRES_ML_1M
|
|
elif name == "ml-10m":
|
|
self.genres = GENRES_ML_10M
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
super(MovieLensDataset, self).__init__(
|
|
name=name,
|
|
url=_get_dgl_url(self._url[name]),
|
|
raw_dir=raw_dir,
|
|
force_reload=force_reload,
|
|
verbose=verbose,
|
|
transform=transform,
|
|
)
|
|
|
|
def check_version(self):
|
|
valid_ratio, test_ratio = load_info(self.version_path)
|
|
if self.valid_ratio == valid_ratio and (
|
|
self.test_ratio == test_ratio if self.name != "ml-100k" else True
|
|
):
|
|
return True
|
|
else:
|
|
if self.name == "ml-100k":
|
|
print(
|
|
f"The current valid ratio ({self.valid_ratio}) "
|
|
"is not the same as the last setting "
|
|
f"(valid: {valid_ratio}). "
|
|
f"MovieLens {self.name} will be re-processed with the new dataset split setting."
|
|
)
|
|
else:
|
|
print(
|
|
f"At least one of current valid ({self.valid_ratio}) and test ({self.test_ratio}) ratio "
|
|
"are not the same as the last setting "
|
|
f"(valid: {valid_ratio}, test: {test_ratio}). "
|
|
f"MovieLens {self.name} will be re-processed with the new dataset split setting."
|
|
)
|
|
return False
|
|
|
|
def download(self):
|
|
zip_file_path = os.path.join(self.raw_dir, self.name + ".zip")
|
|
download(self.url, path=zip_file_path)
|
|
extract_archive(zip_file_path, self.raw_dir, overwrite=True)
|
|
|
|
def process(self):
|
|
print(f"Starting processing {self.name} ...")
|
|
|
|
# 0. loading movie features
|
|
movie_feat = load_info(
|
|
os.path.join(self.raw_path, "movie_feat.pkl")
|
|
).to(torch.float)
|
|
# 1. dataset split: train + (valid + ) test
|
|
if self.name == "ml-100k":
|
|
train_rating_data = self._load_raw_rates(
|
|
os.path.join(self.raw_path, "u1.base"), "\t"
|
|
)
|
|
test_rating_data = self._load_raw_rates(
|
|
os.path.join(self.raw_path, "u1.test"), "\t"
|
|
)
|
|
indices = np.arange(len(train_rating_data))
|
|
train, valid, _ = split_dataset(
|
|
indices,
|
|
[1 - self.valid_ratio, self.valid_ratio, 0.0],
|
|
shuffle=True,
|
|
random_state=self.random_state,
|
|
)
|
|
train_rating_data, valid_rating_data = (
|
|
train_rating_data.iloc[train.indices],
|
|
train_rating_data.iloc[valid.indices],
|
|
)
|
|
all_rating_data = pd.concat(
|
|
[train_rating_data, valid_rating_data, test_rating_data]
|
|
)
|
|
|
|
elif self.name == "ml-1m" or self.name == "ml-10m":
|
|
all_rating_data = self._load_raw_rates(
|
|
os.path.join(self.raw_path, "ratings.dat"), "::"
|
|
)
|
|
indices = np.arange(len(all_rating_data))
|
|
train, valid, test = split_dataset(
|
|
indices,
|
|
[
|
|
1 - self.valid_ratio - self.test_ratio,
|
|
self.valid_ratio,
|
|
self.test_ratio,
|
|
],
|
|
shuffle=True,
|
|
random_state=self.random_state,
|
|
)
|
|
train_rating_data, valid_rating_data, test_rating_data = (
|
|
all_rating_data.iloc[train.indices],
|
|
all_rating_data.iloc[valid.indices],
|
|
all_rating_data.iloc[test.indices],
|
|
)
|
|
|
|
# 2. load user and movie data, and drop those unseen in rating_data
|
|
user_data = self._load_raw_user_data()
|
|
movie_data = self._load_raw_movie_data()
|
|
user_data = self._drop_unseen_nodes(
|
|
data_df=user_data,
|
|
col_name="id",
|
|
reserved_ids_set=set(all_rating_data["user_id"].values),
|
|
)
|
|
movie_data = self._drop_unseen_nodes(
|
|
data_df=movie_data,
|
|
col_name="id",
|
|
reserved_ids_set=set(all_rating_data["movie_id"].values),
|
|
)
|
|
|
|
user_feat = Tensor(self._process_user_feat(user_data))
|
|
|
|
# 3. generate rating pairs
|
|
# Map user/movie to the global id
|
|
self._global_user_id_map = {
|
|
ele: i for i, ele in enumerate(user_data["id"])
|
|
}
|
|
self._global_movie_id_map = {
|
|
ele: i for i, ele in enumerate(movie_data["id"])
|
|
}
|
|
|
|
# pair value is idx rather than id
|
|
u_indices, v_indices, labels = self._generate_pair_value(
|
|
all_rating_data
|
|
)
|
|
all_rating_pairs = (
|
|
LongTensor(u_indices),
|
|
LongTensor(v_indices),
|
|
)
|
|
all_rating_values = Tensor(labels)
|
|
|
|
graph = self.construct_g(
|
|
all_rating_pairs, all_rating_values, user_feat, movie_feat
|
|
)
|
|
self.graph = self.add_masks(
|
|
graph, train_rating_data, valid_rating_data, test_rating_data
|
|
)
|
|
|
|
print(f"End processing {self.name} ...")
|
|
|
|
def construct_g(self, rate_pairs, rate_values, user_feat, movie_feat):
|
|
g = heterograph(
|
|
{
|
|
("user", "user-movie", "movie"): (rate_pairs[0], rate_pairs[1]),
|
|
("movie", "movie-user", "user"): (rate_pairs[1], rate_pairs[0]),
|
|
}
|
|
)
|
|
ndata = {"user": user_feat, "movie": movie_feat}
|
|
edata = {"user-movie": rate_values, "movie-user": rate_values}
|
|
g.ndata["feat"] = ndata
|
|
g.edata["rate"] = edata
|
|
return g
|
|
|
|
def add_masks(
|
|
self, g, train_rating_data, valid_rating_data, test_rating_data
|
|
):
|
|
train_u_indices, train_v_indices, _ = self._generate_pair_value(
|
|
train_rating_data
|
|
)
|
|
valid_u_indices, valid_v_indices, _ = self._generate_pair_value(
|
|
valid_rating_data
|
|
)
|
|
test_u_indices, test_v_indices, _ = self._generate_pair_value(
|
|
test_rating_data
|
|
)
|
|
|
|
# user-movie
|
|
train_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
|
|
train_mask[
|
|
g.edge_ids(train_u_indices, train_v_indices, etype="user-movie")
|
|
] = True
|
|
valid_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
|
|
valid_mask[
|
|
g.edge_ids(valid_u_indices, valid_v_indices, etype="user-movie")
|
|
] = True
|
|
test_mask = torch.zeros((g.num_edges("user-movie"),), dtype=torch.bool)
|
|
test_mask[
|
|
g.edge_ids(test_u_indices, test_v_indices, etype="user-movie")
|
|
] = True
|
|
|
|
g.edges["user-movie"].data["train_mask"] = train_mask
|
|
g.edges["user-movie"].data["valid_mask"] = valid_mask
|
|
g.edges["user-movie"].data["test_mask"] = test_mask
|
|
|
|
# movie-user
|
|
train_mask_rev = torch.zeros(
|
|
(g.num_edges("movie-user"),), dtype=torch.bool
|
|
)
|
|
train_mask_rev[
|
|
g.edge_ids(train_v_indices, train_u_indices, etype="movie-user")
|
|
] = True
|
|
valid_mask_rev = torch.zeros(
|
|
(g.num_edges("movie-user"),), dtype=torch.bool
|
|
)
|
|
valid_mask_rev[
|
|
g.edge_ids(valid_v_indices, valid_u_indices, etype="movie-user")
|
|
] = True
|
|
test_mask_rev = torch.zeros(
|
|
(g.num_edges("movie-user"),), dtype=torch.bool
|
|
)
|
|
test_mask_rev[
|
|
g.edge_ids(test_v_indices, test_u_indices, etype="movie-user")
|
|
] = True
|
|
|
|
g.edges["movie-user"].data["train_mask"] = train_mask_rev
|
|
g.edges["movie-user"].data["valid_mask"] = valid_mask_rev
|
|
g.edges["movie-user"].data["test_mask"] = test_mask_rev
|
|
|
|
return g
|
|
|
|
def has_cache(self):
|
|
if (
|
|
os.path.exists(self.graph_path)
|
|
and os.path.exists(self.version_path)
|
|
and self.check_version()
|
|
):
|
|
return True
|
|
return False
|
|
|
|
def save(self):
|
|
save_graphs(self.graph_path, [self.graph])
|
|
save_info(self.version_path, [self.valid_ratio, self.test_ratio])
|
|
if self.verbose:
|
|
print(f"Done saving data into {self.raw_path}.")
|
|
|
|
def load(self):
|
|
g_list, _ = load_graphs(self.graph_path)
|
|
self.graph = g_list[0]
|
|
|
|
"""
|
|
To avoid the problem each time loading boolean tensor from the disk, boolean values
|
|
would be automatically converted into torch.uint8 types, and a deprecation warning would
|
|
be raised for using torch.uint8
|
|
"""
|
|
for e in self.graph.etypes:
|
|
self.graph.edges[e].data["train_mask"] = (
|
|
self.graph.edges[e].data["train_mask"].to(torch.bool)
|
|
)
|
|
self.graph.edges[e].data["valid_mask"] = (
|
|
self.graph.edges[e].data["valid_mask"].to(torch.bool)
|
|
)
|
|
self.graph.edges[e].data["test_mask"] = (
|
|
self.graph.edges[e].data["test_mask"].to(torch.bool)
|
|
)
|
|
|
|
def __getitem__(self, idx):
|
|
assert (
|
|
idx == 0
|
|
), "This dataset has only one set of training, validation and testing graph"
|
|
if self._transform is None:
|
|
return self.graph
|
|
else:
|
|
return self._transform(self.graph)
|
|
|
|
def __len__(self):
|
|
return 1
|
|
|
|
@property
|
|
def raw_path(self):
|
|
return os.path.join(self.raw_dir, self.name)
|
|
|
|
@property
|
|
def graph_path(self):
|
|
return os.path.join(self.raw_path, self.name + ".bin")
|
|
|
|
@property
|
|
def version_path(self):
|
|
return os.path.join(self.raw_path, self.name + "_version.pkl")
|
|
|
|
def _process_user_feat(self, user_data):
|
|
if self.name == "ml-100k" or self.name == "ml-1m":
|
|
ages = user_data["age"].values.astype(np.float32)
|
|
gender = (user_data["gender"] == "F").values.astype(np.float32)
|
|
all_occupations = set(user_data["occupation"])
|
|
occupation_map = {ele: i for i, ele in enumerate(all_occupations)}
|
|
occupation_one_hot = np.zeros(
|
|
shape=(user_data.shape[0], len(all_occupations)),
|
|
dtype=np.float32,
|
|
)
|
|
occupation_one_hot[
|
|
np.arange(user_data.shape[0]),
|
|
np.array(
|
|
[occupation_map[ele] for ele in user_data["occupation"]]
|
|
),
|
|
] = 1
|
|
user_features = np.concatenate(
|
|
[
|
|
ages.reshape((user_data.shape[0], 1)) / 50.0,
|
|
gender.reshape((user_data.shape[0], 1)),
|
|
occupation_one_hot,
|
|
],
|
|
axis=1,
|
|
)
|
|
elif self.name == "ml-10m":
|
|
user_features = np.zeros(
|
|
shape=(user_data.shape[0], 1), dtype=np.float32
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
return user_features
|
|
|
|
def _load_raw_user_data(self):
|
|
if self.name == "ml-100k":
|
|
user_data = pd.read_csv(
|
|
os.path.join(self.raw_path, "u.user"),
|
|
sep="|",
|
|
header=None,
|
|
names=["id", "age", "gender", "occupation", "zip_code"],
|
|
engine="python",
|
|
)
|
|
elif self.name == "ml-1m":
|
|
user_data = pd.read_csv(
|
|
os.path.join(self.raw_path, "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.raw_path, "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",
|
|
)
|
|
user_data = pd.DataFrame(
|
|
np.unique(rating_info["user_id"].values.astype(np.int32)),
|
|
columns=["id"],
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
return user_data
|
|
|
|
def _load_raw_movie_data(self):
|
|
file_path = os.path.join(self.raw_path, "u.item")
|
|
if self.name == "ml-100k":
|
|
movie_data = pd.read_csv(
|
|
file_path,
|
|
sep="|",
|
|
header=None,
|
|
names=[
|
|
"id",
|
|
"title",
|
|
"release_date",
|
|
"video_release_date",
|
|
"url",
|
|
]
|
|
+ GENRES_ML_100K,
|
|
engine="python",
|
|
encoding="ISO-8859-1",
|
|
)
|
|
elif self.name == "ml-1m" or self.name == "ml-10m":
|
|
file_path = os.path.join(self.raw_path, "movies.dat")
|
|
movie_data = 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(self.genres)}
|
|
genre_map["Children's"] = genre_map["Children"]
|
|
genre_map["Childrens"] = genre_map["Children"]
|
|
movie_genres = np.zeros(
|
|
shape=(movie_data.shape[0], len(self.genres)), dtype=np.float32
|
|
)
|
|
for i, genres in enumerate(movie_data["genres"]):
|
|
for ele in genres.split("|"):
|
|
if ele in genre_map:
|
|
movie_genres[i, genre_map[ele]] = 1.0
|
|
else:
|
|
movie_genres[i, genre_map["unknown"]] = 1.0
|
|
for idx, genre_name in enumerate(self.genres):
|
|
movie_data[genre_name] = movie_genres[:, idx]
|
|
movie_data = movie_data.drop(columns=["genres"])
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
return movie_data
|
|
|
|
def _load_raw_rates(self, file_path, sep):
|
|
rating_data = 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",
|
|
)
|
|
rating_data = rating_data.reset_index(drop=True)
|
|
return rating_data
|
|
|
|
def _drop_unseen_nodes(self, data_df, col_name, reserved_ids_set):
|
|
data_df = data_df[data_df[col_name].isin(reserved_ids_set)]
|
|
data_df.reset_index(drop=True, inplace=True)
|
|
return data_df
|
|
|
|
def _generate_pair_value(self, rating_data):
|
|
rating_pairs = (
|
|
np.array(
|
|
[
|
|
self._global_user_id_map[ele]
|
|
for ele in rating_data["user_id"]
|
|
],
|
|
dtype=np.int32,
|
|
),
|
|
np.array(
|
|
[
|
|
self._global_movie_id_map[ele]
|
|
for ele in rating_data["movie_id"]
|
|
],
|
|
dtype=np.int32,
|
|
),
|
|
)
|
|
rating_values = rating_data["rating"].values.astype(np.float32)
|
|
return rating_pairs[0], rating_pairs[1], rating_values
|
|
|
|
def __repr__(self):
|
|
return (
|
|
f'Dataset("{self.name}", num_graphs={len(self)},'
|
|
+ f" save_path={self.raw_path}), valid_ratio={self.valid_ratio}, test_ratio={self.test_ratio}"
|
|
)
|