Next

Author: orazve

examples/kge-distmult-nations.py

@@ -212,7 +212,7 @@ def inspect_graph(G):
     df = gds.kge.model.predict(
         G_train,
         model_name=model_name,
-        top_k=10,
+        top_k=3,
         node_ids=[
             gds.find_node_id(["Entity"], {"text": "brazil"}),
             gds.find_node_id(["Entity"], {"text": "uk"}),

graphdatascience/model/kge_runner.py

@@ -137,10 +137,12 @@ def _stream_results(self, user_name: str, model_name: str, job_id: str) -> DataF
         )
         res.raise_for_status()
 
-        with open("res.json", mode="wb+") as f:
+        res_file_name = f'res_{job_id}.json'
+        with open(res_file_name, mode="wb+") as f:
             f.write(res.content)
 
-        df = pd.read_json("res.json", orient="records", lines=True)
+        df = pd.read_json(res_file_name, orient="records", lines=True)
+        os.remove(res_file_name)
         return df
 
     def _start_job(self, config: Dict[str, Any]) -> str: