Embedding pipeline / similarity processing and embedding service fixes (#71)

* Added model loading from local fs
* Updated to python3.7
* Enabled embedder fetching from local models
* Fixes to EmbeddingPipeline: added bugfixes and prediction to EmbeddingPipeline
* App fixes
* Changed prediction API of the embedder app
* Added service prediction for json and nexus pgframes
* Updated notebooks

Co-authored-by: Ievgeniia Oshurko <eugenia.oshurko@epfl.ch>

Author: eugeniashurko

.dockerignore

@@ -1,8 +1,25 @@
-tests/*
 examples/*
+cord19kg/examples/*
+tests/*
 docs/*
-cord19kg/*
-.pytest_cache/*
 build/*
 dist/*
-bluegraph.egg-info/*
+.pytest_cache/*
+.tox/*
+.eggs/*
+*.egg-info/
+*.ipynb_checkpoints/
+*.DS_Store
+bluegraph.egg-info/*
+bluegraph/version.py
+tests/data/*
+__MACOSX/
+build/
+dist/
+.pytest_cache
+__pycache__
+*/__pycache__
+attri2vec_test_model.zip
+neo4j_sage_embedder.zip
+stellar_sage_embedder.zip
+services/embedder/examples/*

.gitignore

@@ -3,7 +3,6 @@ build/
 dist/
 .pytest_cache
 __pycache__
-kganalytics/notebooks/data/gephi*
 *.ipynb_checkpoints/
 docs/_build/
 *.DS_Store
@@ -16,17 +15,24 @@ examples/data/attri2vec_test_model.zip
 examples/data/nasa.json
 examples/data/nasa_comention.json
 cord19kg/examples/data/Glucose_risk_3000_papers.csv
+cord19kg/examples/data/Glucose_risk_3000_paper_meta_data.csv
 cord19kg/examples/data/NCIT_ontology_linking.csv
 cord19kg/examples/data/NCIT_ontology_linking_3000_papers.csv
 cord19kg/examples/data/CORD_19_v47_occurrence_top_10000.json
 cord19kg/examples/data/_MACOSX/*
 cord19kg/examples/config/forge-config.yml
 cord19kg/examples/models/neuroshapes/
 cord19kg/examples/data/output_graphs/*
-cord19kg/examples/data/Glucose_risk_3000_paper_meta_data.csv
+cord19kg/examples/data/*.zip
+attri2vec_test_model.zip
+neo4j_sage_emedder.zip
+stellar_sage_emedder.zip
 __MACOSX/
-*.zip
+# *.zip
 .tox/
 .eggs
 bluegraph/version.py
 tests/data/*
+test_sim_proc.faiss
+test_sim_proc.pkl
+.coverage

bluegraph/backends/neo4j/embed/embedders.py

@@ -206,7 +206,7 @@ def _fit_inductive_embedder(self, train_graph):
         train_graph.execute(train_query)
         return model_id
 
-    def _predict_embeddings(self, graph,
+    def _predict_embeddings(self, graph, nodes=None,
                             write=False, write_property=None):
         node_edge_selector = graph.get_projection_query(
             self.graph_configs["edge_weight"],

bluegraph/backends/neo4j/io.py

@@ -487,6 +487,8 @@ def subgraph(self, nodes_to_include=None, edges_to_include=None,
         """Get a node/edge induced subgraph."""
         return Neo4jGraphView(
             self.driver, self.node_label, self.edge_label,
+            nodes_to_include=nodes_to_include,
+            edges_to_include=edges_to_include,
             nodes_to_exclude=nodes_to_exclude,
             edges_to_exclude=edges_to_exclude,
             directed=self.directed)
@@ -507,11 +509,16 @@ class Neo4jGraphView(object):
     TODO: make methods public
     """
     def __init__(self, driver, node_label,
-                 edge_label, nodes_to_exclude=None,
-                 edges_to_exclude=None, directed=True):
+                 edge_label,
+                 nodes_to_include=None, edges_to_include=None,
+                 nodes_to_exclude=None, edges_to_exclude=None,
+                 directed=True):
+        """Initialize an instance of Neo4jGraphView."""
         self.driver = driver
         self.node_label = node_label
         self.edge_label = edge_label
+        self.nodes_to_include = nodes_to_include
+        self.edges_to_include = edges_to_include
         self.nodes_to_exclude = nodes_to_exclude if nodes_to_exclude else []
         self.edges_to_exclude = edges_to_exclude if edges_to_exclude else []
         self.directed = directed

bluegraph/backends/stellargraph/embed/embedders.py

@@ -220,12 +220,14 @@ def _fit_inductive_embedder(self, train_graph):
         embedding_model = Model(inputs=x_inp_src, outputs=x_out_src)
         return embedding_model
 
-    def _predict_embeddings(self, graph):
+    def _predict_embeddings(self, graph, nodes=None):
+        if nodes is None:
+            nodes = graph.nodes()
         node_generator = _dispatch_generator(
-            graph, self.model_name, self.params).flow(graph.nodes())
+            graph, self.model_name, self.params).flow(nodes)
 
         node_embeddings = self._embedding_model.predict(node_generator)
-        res = dict(zip(graph.nodes(), node_embeddings.tolist()))
+        res = dict(zip(nodes, node_embeddings.tolist()))
         return res
 
     @staticmethod

bluegraph/core/embed/embedders.py

@@ -61,7 +61,7 @@ def _fit_inductive_embedder(self, train_graph):
         pass
 
     @abstractmethod
-    def _predict_embeddings(self, graph):
+    def _predict_embeddings(self, graph, nodes=None):
         pass
 
     @staticmethod
@@ -159,16 +159,17 @@ def fit_model(self, pgframe):
             embeddings = embeddings.set_index("@id")
         return embeddings
 
-    def predict_embeddings(self, pgframe):
+    def predict_embeddings(self, pgframe, nodes=None):
         """Predict embeddings of out-sample elements."""
+        if nodes is None:
+            nodes = pgframe.nodes()
         if self._embedding_model is None:
             raise ElementEmbedder.PredictionException(
                 "Embedder does not have a predictive model")
 
-        input_graph = self._generate_graph(
-            pgframe, self.graph_configs)
+        input_graph = self._generate_graph(pgframe, self.graph_configs)
 
-        node_embeddings = self._predict_embeddings(input_graph)
+        node_embeddings = self._predict_embeddings(input_graph, nodes=nodes)
         node_embeddings = pd.DataFrame(
             node_embeddings.items(), columns=["@id", "embedding"])
         node_embeddings = node_embeddings.set_index("@id")