Fix json result processor

Author: MGlauer

chebai/result/base.py

@@ -53,7 +53,7 @@ def _generate_predictions(self, data_path, raw=False, **kwargs):
                 for x in self.dataset._load_dict(data_path)
             ]
         else:
-            data_tuples = torch.load(data_path)
+            data_tuples = [(x.get("raw_features", x["ident"]), x["ident"], x) for x in  torch.load(data_path)]
 
         for raw_features, ident, row in tqdm.tqdm(data_tuples):
             raw_labels = row.get("labels")

chebai/result/prediction_json.py

@@ -16,11 +16,11 @@ def close(self):
             json.dump(self.data, fout)
             del self.data
 
-    def process_prediction(self, proc_id, features, labels, pred, ident):
+    def process_prediction(self, proc_id, raw_features, labels, preds, ident, **kwargs):
         self.data.append(
             dict(
                 ident=ident,
                 labels=labels if labels is not None else None,
-                prediction=pred.tolist(),
+                prediction=preds.tolist(),
             )
         )