- Adds support for multiple index backends (Closes #4, Closes #5)

- Adds id map for all index types to support any id type (Closes #6)
- Skip installing Faiss on Windows (Closes #1, Closes #8)
- Default to Faiss when available, Annoy otherwise
- Switch to faiss-cpu (Closes #2, Closes #9)
- Add quantize flag for Faiss (Closes #10)

setup.py

@@ -5,7 +5,7 @@ with open("README.md", "r") as f:
     DESCRIPTION = f.read()
 
 setup(name="txtai",
-      version="1.0.0",
+      version="1.1.0",
       author="NeuML",
       description="AI-powered search engine",
       long_description=DESCRIPTION,
@@ -22,8 +22,10 @@ setup(name="txtai",
       keywords="search embedding machine-learning nlp",
       python_requires=">=3.6",
       install_requires=[
-          "faiss-gpu>=1.6.3",
+          "annoy>=1.16.3",
+          "faiss-cpu>=1.6.3; os_name != 'nt'",
           "fasttext>=0.9.2",
+          "hnswlib>=0.4.0",
           "nltk>=3.5",
           "numpy>=1.18.4",
           "pymagnitude-lite>=0.1.43",

src/python/txtai/ann.py

@@ -0,0 +1,200 @@
+"""
+ANN (Approximate Nearest Neighbors) module
+"""
+
+import numpy as np
+
+# pylint: disable=E0611
+from annoy import AnnoyIndex
+from hnswlib import Index
+
+# Conditionally import Faiss as it's only supported on Linux/macOS
+try:
+    import faiss
+    FAISS = True
+except ImportError:
+    FAISS = False
+
+class ANN(object):
+    """
+    Base class for ANN models.
+    """
+
+    @staticmethod
+    def create(config):
+        """
+        Create an ANN model.
+
+        Args:
+            config: index configuration parameters
+
+        Returns:
+            ANN
+        """
+
+        # ANN model
+        model = None
+        backend = config.get("backend")
+
+        # Default backend if not provided, based on available libraries
+        if not backend:
+            backend = "faiss" if FAISS else "annoy"
+
+        # Create ANN instance
+        if backend == "annoy":
+            model = Annoy(config)
+        elif backend == "hnsw":
+            model = HNSW(config)
+        else:
+            # Raise error if trying to create a Faiss index without Faiss installed
+            if not FAISS:
+                raise ImportError("Faiss library is not installed")
+
+            model = Faiss(config)
+
+        # Store config back
+        config["backend"] = backend
+
+        return model
+
+    def __init__(self, config):
+        """
+        Creates a new ANN model.
+        """
+
+        # ANN index
+        self.model = None
+
+        # Model configuration
+        self.config = config
+
+    def load(self, path):
+        """
+        Loads an ANN model at path.
+        """
+
+    def index(self, embeddings):
+        """
+        Builds an ANN model.
+
+        Args:
+            embeddings: embeddings array
+        """
+
+    def search(self, query, limit):
+        """
+        Searches ANN model for query. Returns topn results.
+
+        Args:
+            query: query vector
+            limit: maximum results
+        """
+
+    def save(self, path):
+        """
+        Saves an ANN model at path.
+        """
+
+class Annoy(ANN):
+    """
+    Builds an ANN model using the Annoy library.
+    """
+
+    def load(self, path):
+        # Load index
+        self.model = AnnoyIndex(self.config["dimensions"], self.config["metric"])
+        self.model.load(path)
+
+    def index(self, embeddings):
+        # Inner product is equal to cosine similarity on normalized vectors
+        self.config["metric"] = "dot"
+
+        # Create index
+        self.model = AnnoyIndex(self.config["dimensions"], self.config["metric"])
+
+        # Add items
+        for x in range(embeddings.shape[0]):
+            self.model.add_item(x, embeddings[x])
+
+        # Build index
+        self.model.build(10)
+
+    def search(self, query, limit):
+        # Run the query
+        ids, scores = self.model.get_nns_by_vector(query, n=limit, include_distances=True)
+
+        # Map results to [(id, score)]
+        return list(zip(ids, scores))
+
+    def save(self, path):
+        # Write index
+        self.model.save(path)
+
+class Faiss(ANN):
+    """
+    Builds an ANN model using the Faiss library.
+    """
+
+    def load(self, path):
+        # Load index
+        self.model = faiss.read_index(path)
+
+    def index(self, embeddings):
+        # Create embeddings index. Inner product is equal to cosine similarity on normalized vectors.
+        if self.config.get("quantize"):
+            params = "IVF100,SQ8" if embeddings.shape[0] >= 5000 else "IDMap,SQ8"
+        else:
+            params = "IVF100,Flat" if embeddings.shape[0] >= 5000 else "IDMap,Flat"
+
+        self.model = faiss.index_factory(embeddings.shape[1], params, faiss.METRIC_INNER_PRODUCT)
+
+        # Train model
+        self.model.train(embeddings)
+        self.model.add_with_ids(embeddings, np.array(range(embeddings.shape[0])))
+
+    def search(self, query, limit):
+        # Run the query
+        self.model.nprobe = 6
+        scores, ids = self.model.search(query.reshape(1, -1), limit)
+
+        # Map results to [(id, score)]
+        return list(zip(ids[0].tolist(), (scores[0]).tolist()))
+
+    def save(self, path):
+        # Write index
+        faiss.write_index(self.model, path)
+
+class HNSW(ANN):
+    """
+    Builds an ANN model using the hnswlib library.
+    """
+
+    def load(self, path):
+        # Load index
+        self.model = Index(dim=self.config["dimensions"], space=self.config["metric"])
+        self.model.load_index(path)
+
+    def index(self, embeddings):
+        # Inner product is equal to cosine similarity on normalized vectors
+        self.config["metric"] = "ip"
+
+        # Create index
+        self.model = Index(dim=self.config["dimensions"], space=self.config["metric"])
+        self.model.init_index(max_elements=embeddings.shape[0])
+
+        # Add items
+        self.model.add_items(embeddings, np.array(range(embeddings.shape[0])))
+
+    def search(self, query, limit):
+        # Run the query
+        ids, distances = self.model.knn_query(query.reshape(1, -1), k=limit)
+
+        # Convert distances to similarity scores
+        scores = [1 - d for d in distances[0]]
+
+        # Map results to [(id, score)]
+        return list(zip(ids[0], scores))
+
+    def save(self, path):
+        # Write index
+        self.model.save_index(path)

src/python/txtai/embeddings.py

@@ -5,13 +5,13 @@ Embeddings module
 import pickle
 import os
 
-import faiss
 import numpy as np
 
 from sklearn.decomposition import TruncatedSVD
 
+from .ann import ANN
 from .scoring import Scoring
-from .vectors import Factory
+from .vectors import Vectors
 
 class Embeddings(object):
     """
@@ -45,20 +45,17 @@ class Embeddings(object):
         self.scoring = Scoring.create(self.config["scoring"]) if self.config and self.config.get("scoring") else None
 
         # Sentence vectors model
-        self.model = self.loadVectors(self.config["path"]) if self.config else None
+        self.model = self.loadVectors() if self.config else None
 
-    def loadVectors(self, path):
+    def loadVectors(self):
         """
-        Loads a word vector model at path.
-
-        Args:
-            path: path to word vector model
+        Loads a word vector model set in config.
 
         Returns:
             vector model
         """
 
-        return Factory.create(self.config.get("method"), path, True if not self.embeddings else False, self.scoring)
+        return Vectors.create(self.config.get("method"), self.config["path"], True if not self.embeddings else False, self.scoring)
 
     def score(self, documents):
         """
@@ -100,14 +97,15 @@ class Embeddings(object):
         # Normalize embeddings
         self.normalize(embeddings)
 
-        # pylint: disable=E1136
-        # Create embeddings index. Inner product is equal to cosine similarity on normalized vectors.
-        params = "IVF100,SQ8" if embeddings.shape[0] >= 5000 else "IDMap,SQ8"
-        self.embeddings = faiss.index_factory(embeddings.shape[1], params, faiss.METRIC_INNER_PRODUCT)
+        # Save configuration
+        self.config["ids"] = ids
+        self.config["dimensions"] = dimensions
 
-        # Train on embeddings model
-        self.embeddings.train(embeddings)
-        self.embeddings.add_with_ids(embeddings, np.array(ids))
+        # Create embeddings index, store backend used
+        self.embeddings = ANN.create(self.config)
+
+        # Build the index
+        self.embeddings.index(embeddings)
 
     def buildLSA(self, embeddings, components):
         """
@@ -205,11 +203,14 @@ class Embeddings(object):
         embedding = self.transform((None, query, None))
 
         # Search embeddings index
-        self.embeddings.nprobe = 6
-        results = self.embeddings.search(embedding.reshape(1, -1), limit)
+        results = self.embeddings.search(embedding, limit)
 
-        # Map results to [(id, score)]
-        return list(zip(results[1][0].tolist(), (results[0][0]).tolist()))
+        # Map ids if id mapping available
+        lookup = self.config.get("ids")
+        if lookup:
+            return [(lookup[i], score) for i, score in results]
+
+        return results
 
     def similarity(self, query, documents):
         """
@@ -249,7 +250,8 @@ class Embeddings(object):
             self.config = pickle.load(handle)
 
         # Sentence embeddings index
-        self.embeddings = faiss.read_index("%s/embeddings" % path)
+        self.embeddings = ANN.create(self.config)
+        self.embeddings.load("%s/embeddings" % path)
 
         # Dimensionality reduction
         if self.config.get("pca"):
@@ -262,7 +264,7 @@ class Embeddings(object):
             self.scoring.load(path)
 
         # Sentence vectors model - transforms text into sentence embeddings
-        self.model = self.loadVectors(self.config["path"])
+        self.model = self.loadVectors()
 
     def save(self, path):
         """
@@ -281,7 +283,7 @@ class Embeddings(object):
                 pickle.dump(self.config, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
             # Write sentence embeddings index
-            faiss.write_index(self.embeddings, "%s/embeddings" % path)
+            self.embeddings.save("%s/embeddings" % path)
 
             # Save dimensionality reduction
             if self.lsa:

src/python/txtai/vectors.py

@@ -52,9 +52,9 @@ def transform(document):
 
     return (document[0], VECTORS.transform(document))
 
-class Factory(object):
+class Vectors(object):
     """
-    Creates a new Vectors model instance.
+    Base class for sentence embeddings/vector models.
     """
 
     @staticmethod
@@ -70,11 +70,6 @@ class Factory(object):
         # Create vector model instance
         return TransformersVectors(path) if method == "transformers" else WordVectors(path, blocking, scoring)
 
-class Vectors(object):
-    """
-    Base class for sentence embeddings/vector models.
-    """
-
     def load(self, path, blocking):
         """
         Loads vector model at path.