SageRec/python/graph_sampler.pyi at main · sathariels/SageRec · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
"""Type stub for the pybind11 graph_sampler extension.

The compiled module is produced by the CMake target `graph_sampler`.
Users occupy [0, num_users); movies occupy [num_users, num_nodes).
Construction takes local (user_id, movie_id) pairs, not global movie IDs.
parse_movielens_100k accepts in-memory tab-separated u.data text only.
"""

from collections.abc import Sequence

class GraphError(ValueError):
    """Invalid graph construction, parse request, or sample request."""

class MovieLens100kInteraction:
    @property
    def user_id(self) -> int: ...
    @property
    def movie_id(self) -> int: ...
    @property
    def rating(self) -> int: ...
    @property
    def timestamp(self) -> int: ...

class MovieLens100kRatings:
    @property
    def num_users(self) -> int: ...
    @property
    def num_movies(self) -> int: ...
    @property
    def interactions(self) -> list[MovieLens100kInteraction]: ...
    @property
    def user_source_ids(self) -> list[int]: ...
    @property
    def movie_source_ids(self) -> list[int]: ...
    def local_pairs(self) -> list[tuple[int, int]]: ...

def parse_movielens_100k(text: str) -> MovieLens100kRatings: ...

class BipartiteCSR:
    def __init__(
        self,
        num_users: int,
        num_movies: int,
        interactions: Sequence[tuple[int, int]],
    ) -> None: ...
    @property
    def num_users(self) -> int: ...
    @property
    def num_movies(self) -> int: ...
    @property
    def num_nodes(self) -> int: ...
    def degree(self, node_id: int) -> int: ...
    def sample_neighbors(self, node_id: int, k: int, seed: int) -> list[int]: ...
    def offsets(self) -> list[int]: ...
    def neighbors(self) -> list[int]: ...