Migrate from Google Colab and modularize the code

[aholanda]

As we need to implement walks in graph partitions, the code complexity tends to increase. Then we need to modularize the code to deal with the complexity that partitions will impose to the model.