File size: 550 Bytes
a3682cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import pickle
import pandas as pd

from src.graph.dataset_builder import build_graph_dataset


def main():
    print("Loading dataset...")
    df = pd.read_csv("data/processed/transactions.csv")
    users = pd.read_csv("data/processed/users.csv")

    print("Building graph dataset...")
    graph_data = build_graph_dataset(df, users)

    os.makedirs("data/graph", exist_ok=True)

    with open("data/graph/graph.pkl", "wb") as f:
        pickle.dump(graph_data, f)

    print("Graph dataset saved")


if __name__ == "__main__":
    main()