Add Embeddings context manager, closes #832

src/python/txtai/embeddings/base.py

@@ -85,6 +85,12 @@ def __init__(self, config=None, models=None, **kwargs):
         # Set initial configuration
         self.configure(config)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
     def score(self, documents):
         """
         Builds a term weighting scoring index. Only used by word vectors models.
@@ -530,6 +536,9 @@ def load(self, path=None, cloud=None, config=None, **kwargs):
             cloud: cloud storage configuration
             config: configuration overrides
             kwargs: additional configuration as keyword args
+
+        Returns:
+            Embeddings
         """
 
         # Load from cloud, if configured
@@ -589,6 +598,8 @@ def load(self, path=None, cloud=None, config=None, **kwargs):
         # Query model
         self.query = self.loadquery()
 
+        return self
+
     def save(self, path, cloud=None, **kwargs):
         """
         Saves an index in a directory at path unless path ends with tar.gz, tar.bz2, tar.xz or zip.

test/python/testembeddings.py

@@ -81,6 +81,22 @@ def testColumns(self):
         uid = embeddings.search("lottery", 1)[0][0]
         self.assertEqual(uid, 4)
 
+    def testContext(self):
+        """
+        Test embeddings context manager
+        """
+
+        # Generate temp file path
+        index = os.path.join(tempfile.gettempdir(), "embeddings.context")
+
+        with Embeddings() as embeddings:
+            embeddings.index(self.data)
+            embeddings.save(index)
+
+        with Embeddings().load(index) as embeddings:
+            uid = embeddings.search(self.data[4], 1)[0][0]
+            self.assertEqual(uid, 4)
+
     def testDefaults(self):
         """
         Test default configuration