diff --git a/src/models/code2vec_features.py b/src/models/code2vec_features.py
index 2d7ad7d..cb1fe2a 100644
--- a/src/models/code2vec_features.py
+++ b/src/models/code2vec_features.py
@@ -14,16 +14,16 @@ def construct(self, value2index, path2index, value2freq, path2freq, path_context
         self._path2index = path2index
         self._value2freq = value2freq
         self._path2freq = path2freq
-
         self._path_contexts = path_contexts
         return self
 
     def _load_tree(self, tree):
-        self.construct(value2index=tree["value2index"],
-                       path2index=tree["path2index"],
-                       value2freq=tree["value2freq"],
-                       path2freq=tree["path2freq"],
-                       path_contexts=tree["path_contexts"])
+        self.construct(
+            value2index=tree["value2index"],
+            path2index={tuple(val[0]): key for (key, val) in tree["index2path_freq"].items()},
+            value2freq=tree["value2freq"],
+            path2freq={tuple(val[0]): val[1] for (_, val) in tree["index2path_freq"].items()},
+            path_contexts=tree["path_contexts"])
 
     @property
     def value2index(self):
@@ -86,9 +86,9 @@ def path2freq_items(self):
 
     def _generate_tree(self):
         return {"value2index": self._value2index,
-                "path2index": self._path2index,
+                "index2path_freq": {val: (key, self._path2freq[key])
+                                    for (key, val) in self._path2index.items()},
                 "value2freq": self._value2freq,
-                "path2freq": self._path2freq,
                 "path_contexts": self._path_contexts}
 
     def dump(self):
@@ -98,9 +98,9 @@ def dump(self):
                "First 10 path -> ID: %s\n" \
                "First 10 value -> frequency: %s\n" \
                "First 10 path -> frequency: %s" % \
-               (len(self._value2index_freq),
-                len(self.path2index_freq),
-                list(islice(self._value2index, 10)),
-                list(islice(self._path2index, 10)),
-                list(islice(self._value2freq, 10)),
-                list(islice(self._path2freq, 10)))
+               (len(self._value2index),
+                len(self._path2index),
+                list(islice(self.value2index_items(), 10)),
+                list(islice(self.path2index_items(), 10)),
+                list(islice(self.value2freq_items(), 10)),
+                list(islice(self.path2freq_items(), 10)))