Embedding pipeline optimization (#70)

* Fixes to EmbeddingPipeline
* Removed embedding table from Embedding Pipeline
* Added error message when element is not in index

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

Author: eugeniashurko

bluegraph/downstream/data_structures.py

@@ -22,7 +22,9 @@
 import re
 import pickle
 import shutil
+import warnings
 
+from bluegraph.exceptions import BlueGraphException, BlueGraphWarning
 from .similarity import SimilarityProcessor
 
 
@@ -87,6 +89,7 @@ def predict(self, pgframe, predict_elements=None):
         return self.model.predict(data)
 
 
+
 class Preprocessor(ABC):
     """Preprocessor inferface for EmbeddingPipeline."""
 
@@ -118,11 +121,9 @@ def fit_model(self):
 class EmbeddingPipeline(object):
 
     def __init__(self, preprocessor=None, embedder=None,
-                 embedding_table=None,
                  similarity_processor=None):
         self.preprocessor = preprocessor
         self.embedder = embedder
-        self.embedding_table = embedding_table
         self.similarity_processor = similarity_processor
 
     def is_transductive(self):
@@ -138,21 +139,38 @@ def run_fitting(self, data):
             train_data = self.preprocessor.transform(data)
         else:
             train_data = data
-        # Train the embedder
-        self.embedding_table = self.embedder.fit_model(train_data)
+        if not self.embedder:
+            raise EmbeddingPipelineException(
+                "Embedder object is not specified: cannot run fitting")
+        else:
+            # Train the embedder
+            embedding_table = self.embedder.fit_model(train_data)
         # Create a similarity processor
         vectors =\
-            self.embedding_table["embedding"].tolist()
+            embedding_table["embedding"].tolist()
         self.similarity_processor._initialize_model(vectors)
-        self.similarity_processor.add(vectors, self.embedding_table.index)
-        self.similarity_processor.index = self.embedding_table.index
+        self.similarity_processor.add(vectors, embedding_table.index)
+        self.similarity_processor.index = embedding_table.index
 
     def run_prediction(self, data):
         pass
 
+    def generate_embedding_table(self):
+        """Generate embedding table from similarity index."""
+        index = self.similarity_processor.index
+        pairs = [
+            (ind, self.similarity_processor._model.reconstruct(i))
+            for i, ind in enumerate(index)
+        ]
+        return pd.DataFrame(
+            pairs, columns=["@id", "embedding"]).set_index("@id")
+
+
     def retrieve_embeddings(self, indices):
-        if self.embedding_table is not None:
-            return self.embedding_table.loc[indices]["embedding"].tolist()
+        if self.similarity_processor is None:
+            raise EmbeddingPipelineException(
+                "Similarity processor object is None, cannot "
+                "retrieve embedding vectors")
         else:
             return [
                 el.tolist()
@@ -191,12 +209,6 @@ def load(cls, path, embedder_interface=None, embedder_ext="pkl"):
                 embedder = embedder_interface.load(
                     os.path.join(path, "embedder.zip"))
 
-        # Load the embedding table
-        embedding_table = None
-        if os.path.isfile(os.path.join(path, "vectors.pkl")):
-            embedding_table = pd.read_pickle(
-                os.path.join(path, "vectors.pkl"))
-
         # Load the similarity processor
         similarity_processor = SimilarityProcessor.load(
             os.path.join(path, "similarity.pkl"),
@@ -205,7 +217,6 @@ def load(cls, path, embedder_interface=None, embedder_ext="pkl"):
         pipeline = cls(
             preprocessor=encoder,
             embedder=embedder,
-            embedding_table=embedding_table,
             similarity_processor=similarity_processor)
 
         if decompressed:
@@ -214,7 +225,6 @@ def load(cls, path, embedder_interface=None, embedder_ext="pkl"):
         return pipeline
 
     def save(self, path, compress=False):
-
         if not os.path.isdir(path):
             os.mkdir(path)
 
@@ -223,12 +233,12 @@ def save(self, path, compress=False):
             pickle.dump(self.preprocessor, f)
 
         # Save the embedding model
-        self.embedder.save(
-            os.path.join(path, "embedder"), compress=True)
-
-        # Save the embedding table
-        self.embedding_table.to_pickle(
-            os.path.join(path, "vectors.pkl"))
+        if self.embedder:
+            self.embedder.save(
+                os.path.join(path, "embedder"), compress=True)
+        else:
+            with open(os.path.join(path, "embedder.pkl"), "wb") as f:
+                pickle.dump(self.preprocessor, f)
 
         # Save the similarity processor
         if self.similarity_processor is not None:
@@ -239,3 +249,9 @@ def save(self, path, compress=False):
         if compress:
             shutil.make_archive(path, 'zip', path)
             shutil.rmtree(path)
+
+    class EmbeddingPipelineException(BlueGraphException):
+        pass
+
+    class EmbeddingPipelineWarning(BlueGraphWarning):
+        pass

bluegraph/downstream/similarity.py

@@ -20,6 +20,9 @@
 import faiss
 import os
 
+from bluegraph.exceptions import BlueGraphException
+
+
 # This is to avoid a wierd Faiss segmentation fault (TODO: investigate)
 os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
@@ -113,8 +116,14 @@ def _initialize_model(self, initial_vectors=None):
 
     def get_vectors(self, existing_indices):
         if self.index is not None:
-            existing_indices = self.index.get_indexer(existing_indices)
-        x = [self._model.reconstruct(int(i)) for i in existing_indices]
+            int_idices = self.index.get_indexer(existing_indices)
+        try:
+            x = [self._model.reconstruct(int(i)) for i in int_idices]
+        except RuntimeError:
+            raise SimilarityProcessor.SimilarityException(
+                "Cannot retrieve vectors for provided elements {} ".format(
+                    existing_indices) +
+                "make sure all the elements are in the index.")
         return x
 
     def query_existing(self, existing_indices, k=10):
@@ -161,16 +170,16 @@ def get_similar_points(self, vectors=None, vector_indices=None,
             indices = int_index
         return indices, distance
 
-    class TrainException(Exception):
+    class TrainException(BlueGraphException):
         pass
 
-    class SimilarityException(Exception):
+    class SimilarityException(BlueGraphException):
         pass
 
-    class IndexException(Exception):
+    class IndexException(BlueGraphException):
         pass
 
-    class QueryException(Exception):
+    class QueryException(BlueGraphException):
         pass