chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -0,0 +1,53 @@
+import unittest
+
+import backend as F
+
+import dgl
+from dgl.data.movielens import MovieLensDataset
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Datasets don't need to be tested on GPU.",
+)
+@unittest.skipIf(
+    dgl.backend.backend_name != "pytorch", reason="only supports pytorch"
+)
+def test_movielens():
+    transform = dgl.AddSelfLoop(new_etypes=True)
+
+    movielens = MovieLensDataset(name="ml-100k", valid_ratio=0.2, verbose=True)
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 100000
+    assert (
+        g.nodes["user"].data["feat"].shape[1]
+        == g.nodes["user"].data["feat"].shape[1]
+        == g.nodes["user"].data["feat"].shape[1]
+        == 23
+    )
+    assert (
+        g.nodes["movie"].data["feat"].shape[1]
+        == g.nodes["movie"].data["feat"].shape[1]
+        == g.nodes["movie"].data["feat"].shape[1]
+        == 320
+    )
+
+    movielens = MovieLensDataset(
+        name="ml-100k", valid_ratio=0.2, transform=transform, verbose=True
+    )
+    g1 = movielens[0]
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+    assert g1.num_edges() - g.num_edges() == g.num_nodes()
+
+    movielens = MovieLensDataset(
+        name="ml-1m", valid_ratio=0.2, test_ratio=0.1, verbose=True
+    )
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 1000209
+
+    movielens = MovieLensDataset(
+        name="ml-10m", valid_ratio=0.2, test_ratio=0.1, verbose=True
+    )
+    g = movielens[0]
+    assert g.num_edges("user-movie") == g.num_edges("movie-user") == 10000054