Add a get_embedding_sources composite method.

PiperOrigin-RevId: 704357725
google-research · Dec 9, 2024 · 025f7f4 · 025f7f4
1 parent b156d6a
commit 025f7f4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/hoplite/db/interface.py b/hoplite/db/interface.py
@@ -263,6 +263,12 @@ def get_embeddings(
     embeddings = [self.get_embedding(int(idx)) for idx in embedding_ids]
     return embedding_ids, np.array(embeddings)
 
+  def get_embedding_sources(
+      self, embedding_ids: np.ndarray
+  ) -> tuple[EmbeddingSource, ...]:
+    """Get an array of embedding sources for the indicated IDs."""
+    return tuple(self.get_embedding_source(int(idx)) for idx in embedding_ids)
+
   def random_batched_iterator(
       self,
       ids: np.ndarray,

diff --git a/hoplite/db/tests/hoplite_test.py b/hoplite/db/tests/hoplite_test.py
@@ -82,6 +82,8 @@ def test_graph_db_interface(self, db_type, thread_split):
       embs = db.get_embeddings_by_source(source.dataset_name, 'fake_id', None)
       self.assertEqual(embs.shape[0], 0)
 
+      sources = db.get_embedding_sources(idxes[:3])
+      self.assertLen(sources, 3)
     db.commit()
 
   @parameterized.product(db_type=PERSISTENT_DB_TYPES)