Deepmemory ux improvements

deeplake/core/vectorstore/dataset_handlers/client_side_dataset_handler.py

@@ -168,6 +168,7 @@ def search(
         return_tensors: List[str],
         return_view: bool,
         deep_memory: bool,
+        return_tql: bool,
     ) -> Union[Dict, Dataset]:
         feature_report_path(
             path=self.bugout_reporting_path,
@@ -243,7 +244,7 @@ def search(
             embedding_tensor=embedding_tensor,
             return_tensors=return_tensors,
             return_view=return_view,
-            deep_memory=deep_memory,
+            return_tql=return_tql,
             token=self.token,
             org_id=self.org_id,
         )

deeplake/core/vectorstore/deep_memory/deep_memory.py

@@ -1,4 +1,5 @@
 import logging
+import pathlib
 import uuid
 from collections import defaultdict
 from pydantic import BaseModel, ValidationError
@@ -8,13 +9,13 @@
 import numpy as np
 
 import deeplake
-from deeplake.enterprise.dataloader import indra_available
 from deeplake.util.exceptions import (
+    DeepMemoryWaitingListError,
     DeepMemoryWaitingListError,
     IncorrectRelevanceTypeError,
     IncorrectQueriesTypeError,
 )
-from deeplake.util.remove_cache import get_base_storage
+from deeplake.util.path import convert_pathlib_to_string_if_needed
 from deeplake.constants import (
     DEFAULT_QUERIES_VECTORSTORE_TENSORS,
     DEFAULT_MEMORY_CACHE_SIZE,
@@ -30,7 +31,15 @@
     feature_report_path,
 )
 from deeplake.util.path import get_path_type
-from deeplake.util.version_control import load_meta
+
+
+def access_control(func):
+    def wrapper(self, *args, **kwargs):
+        if self.client is None:
+            raise DeepMemoryWaitingListError()
+        return func(self, *args, **kwargs)
+
+    return wrapper
 
 
 def use_deep_memory(func):
@@ -46,15 +55,6 @@
     return wrapper
 
 
-def access_control(func):
-    def wrapper(self, *args, **kwargs):
-        if self.client is None:
-            raise DeepMemoryWaitingListError()
-        return func(self, *args, **kwargs)
-
-    return wrapper
-
-
 class Relevance(BaseModel):
     data: List[List[Tuple[str, int]]]
 
@@ -78,7 +78,8 @@
 class DeepMemory:
     def __init__(
         self,
-        dataset_or_path: Union[Dataset, str],
+        dataset: Dataset,
+        path: Union[str, pathlib.Path],
         logger: logging.Logger,
         embedding_function: Optional[Any] = None,
         token: Optional[str] = None,
@@ -87,24 +88,18 @@
         """Based Deep Memory class to train and evaluate models on DeepMemory managed service.
 
         Args:
-            dataset_or_path (Union[Dataset, str]): deeplake dataset object or path.
+            dataset (Dataset): deeplake dataset object or path.
+            path (Union[str, pathlib.Path]): Path to the dataset.
             logger (logging.Logger): Logger object.
             embedding_function (Optional[Any], optional): Embedding funtion class used to convert queries/documents to embeddings. Defaults to None.
             token (Optional[str], optional): API token for the DeepMemory managed service. Defaults to None.
             creds (Optional[Dict[str, Any]], optional): Credentials to access the dataset. Defaults to None.
 
         Raises:
             ImportError: if indra is not installed
-            ValueError: if incorrect type is specified for `dataset_or_path`
         """
-        if isinstance(dataset_or_path, Dataset):
-            self.path = dataset_or_path.path
-        elif isinstance(dataset_or_path, str):
-            self.path = dataset_or_path
-        else:
-            raise ValueError(
-                "dataset_or_path should be a Dataset object or a string path"
-            )
+        self.dataset = dataset
+        self.path = convert_pathlib_to_string_if_needed(path)
 
         feature_report_path(
             path=self.path,
@@ -143,7 +138,8 @@
             relevance (List[List[Tuple[str, int]]]): List of relevant documents for each query with their respective relevance score.
                 The outer list corresponds to the queries and the inner list corresponds to the doc_id, relevence_score pair for each query.
                 doc_id is the document id in the corpus dataset. It is stored in the `id` tensor of the corpus dataset.
-                relevence_score is the relevance score of the document for the query. The range is between 0 and 1, where 0 stands for not relevant and 1 stands for relevant.
+                relevence_score is the relevance score of the document for the query. The value is either 0 and 1, where 0 stands for not relevant (unknown relevance)
+                and 1 stands for relevant. Currently, only values of 1 contribute to the training, and there is no reason to provide examples with relevance of 0.
             embedding_function (Optional[Callable[[str], np.ndarray]], optional): Embedding funtion used to convert queries to embeddings. Defaults to None.
             token (str, optional): API token for the DeepMemory managed service. Defaults to None.
 
@@ -178,7 +174,7 @@
             )
 
         if embedding_function is None and self.embedding_function is not None:
-            embedding_function = self.embedding_function.embed_documents
+            embedding_function = self.embedding_function
 
         runtime = None
         if get_path_type(corpus_path) == "hub":
@@ -484,10 +480,8 @@
         if embedding is not None:
             query_embs = embedding
         else:
-            if self.embedding_function is not None:
-                embedding_function = (
-                    embedding_function or self.embedding_function.embed_documents
-                )
+            if self.embedding_function is not None and embedding_function is None:
+                embedding_function = self.embedding_function
 
             if embedding_function is None:
                 raise ValueError(
@@ -554,6 +548,46 @@
         self.queries_dataset.commit()
         return recalls
 
+    @access_control
+    def get_model(self):
+        """Get the name of the model currently being used by DeepMemory managed service."""
+        return self.dataset.embedding.info["deepmemory"]["model.npy"]["job_id"]
+
+    @access_control
+    def set_model(self, model_name: str):
+        """Set model.npy to use `model_name` instead of default model
+        Args:
+            model_name (str): name of the model to use
+        """
+
+        if "npy" not in model_name:
+            model_name += ".npy"
+
+        # verify model_name
+        self._verify_model_name(model_name)
+
+        # set model.npy to use `model_name` instead of default model
+        self._set_model_npy(model_name)
+
+    def _verify_model_name(self, model_name: str):
+        if model_name not in self.dataset.embedding.info["deepmemory"]:
+            raise ValueError(
+                "Invalid model name. Please choose from the following models: "
+                + ", ".join(self.dataset.embedding.info["deepmemory"].keys())
+            )
+
+    def _set_model_npy(self, model_name: str):
+        # get new model.npy
+        new_model_npy = self.dataset.embedding.info["deepmemory"][model_name]
+
+        # get old deepmemory dictionary and update it:
+        old_deepmemory = self.dataset.embedding.info["deepmemory"]
+        new_deepmemory = old_deepmemory.copy()
+        new_deepmemory.update({"model.npy": new_model_npy})
+
+        # assign new deepmemory dictionary to the dataset:
+        self.dataset.embedding.info["deepmemory"] = new_deepmemory
+
     def _get_dm_client(self):
         path = self.path
         path_type = get_path_type(path)

deeplake/core/vectorstore/deep_memory/test_deepmemory.py

@@ -3,6 +3,7 @@
 import pytest
 import sys
 from time import sleep
+from unittest.mock import MagicMock
 
 import deeplake
 from deeplake import VectorStore
@@ -40,9 +41,9 @@
     assert db.deep_memory is not None
 
 
-def embedding_fn(texts):
+def embedding_fn(texts, embedding_dim=1536):
     return [
-        np.random.uniform(low=-10, high=10, size=(1536)).astype(np.float32)
+        np.random.uniform(low=-10, high=10, size=(embedding_dim)).astype(np.float32)
         for _ in range(len(texts))
     ]
 
@@ -432,7 +433,7 @@
         path=corpus,
         runtime={"tensor_db": True},
         token=hub_cloud_dev_token,
-        embedding_function=DummyEmbedder,
+        embedding_function=embedding_fn,
     )
     recall = db.deep_memory.evaluate(
         queries=queries,
@@ -544,7 +545,7 @@
         token=hub_cloud_dev_token,
     )
 
     output = db.search(
         embedding=query_embedding, deep_memory=True, return_tensors=["id"]
     )
 
@@ -584,7 +585,10 @@
 @requires_libdeeplake
 def test_unsupported_deepmemory_users(local_ds):
     dm = DeepMemory(
-        dataset_or_path=local_ds, logger=logger, embedding_function=DummyEmbedder
+        path=local_ds,
+        dataset=None,
+        logger=logger,
+        embedding_function=DummyEmbedder,
     )
     with pytest.raises(DeepMemoryWaitingListError):
         dm.train(
@@ -660,3 +664,121 @@
             queries=queries,
             relevance="relevances",
         )
+
+
+def test_deepmemory_v2_set_model_should_set_model_for_all_subsequent_loads(
+    local_dmv2_dataset,
+    hub_cloud_dev_token,
+):
+    # Setiing model should set model for all subsequent loads
+    db = VectorStore(path=local_dmv2_dataset, token=hub_cloud_dev_token)
+    assert db.deep_memory.get_model() == "655f86e8ab93e7fc5067a3ac_2"
+
+    # ensure after setting model, get model returns specified model
+    db.deep_memory.set_model("655f86e8ab93e7fc5067a3ac_1")
+
+    assert (
+        db.dataset.embedding.info["deepmemory"]["model.npy"]["job_id"]
+        == "655f86e8ab93e7fc5067a3ac_1"
+    )
+    assert db.deep_memory.get_model() == "655f86e8ab93e7fc5067a3ac_1"
+
+    # ensure after setting model, reloading the dataset returns the same model
+    db = VectorStore(path=local_dmv2_dataset, token=hub_cloud_dev_token)
+    assert db.deep_memory.get_model() == "655f86e8ab93e7fc5067a3ac_1"
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows")
+def test_deepmemory_search_should_contain_correct_answer(
+    corpus_query_relevances_copy,
+    testing_relevance_query_deepmemory,
+    hub_cloud_dev_token,
+):
+    corpus, _, _, _ = corpus_query_relevances_copy
+    relevance, query_embedding = testing_relevance_query_deepmemory
+
+    db = VectorStore(
+        path=corpus,
+        token=hub_cloud_dev_token,
+    )
+
+    output = db.search(
+        embedding=query_embedding, deep_memory=True, return_tensors=["id"]
+    )
+    assert len(output["id"]) == 4
+    assert relevance in output["id"]
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows")
+def test_deeplake_search_should_not_contain_correct_answer(
+    corpus_query_relevances_copy,
+    testing_relevance_query_deepmemory,
+    hub_cloud_dev_token,
+):
+    corpus, _, _, _ = corpus_query_relevances_copy
+    relevance, query_embedding = testing_relevance_query_deepmemory
+
+    db = VectorStore(
+        path=corpus,
+        token=hub_cloud_dev_token,
+    )
+    output = db.search(embedding=query_embedding)
+    assert len(output["id"]) == 4
+    assert relevance not in output["id"]
+
+
+@pytest.mark.slow
+@pytest.mark.flaky(reruns=3)
+@pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows")
+def test_deepmemory_train_with_embedding_function_specified_in_constructor_should_not_throw_any_exception(
+    deepmemory_small_dataset_copy,
+    hub_cloud_dev_token,
+):
+    corpus, queries, relevances, _ = deepmemory_small_dataset_copy
+
+    db = VectorStore(
+        path=corpus,
+        runtime={"tensor_db": True},
+        token=hub_cloud_dev_token,
+        embedding_function=embedding_fn,
+    )
+
+    job_id = db.deep_memory.train(
+        queries=queries,
+        relevance=relevances,
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.flaky(reruns=3)
+@pytest.mark.skipif(sys.platform == "win32", reason="Does not run on Windows")
+def test_deepmemory_evaluate_with_embedding_function_specified_in_constructor_should_not_throw_any_exception(
+    corpus_query_pair_path,
+    hub_cloud_dev_token,
+):
+    corpus, queries = corpus_query_pair_path
+
+    db = VectorStore(
+        path=corpus,
+        runtime={"tensor_db": True},
+        token=hub_cloud_dev_token,
+        embedding_function=embedding_fn,
+    )
+
+    queries_vs = VectorStore(
+        path=queries,
+        runtime={"tensor_db": True},
+        token=hub_cloud_dev_token,
+        embedding_function=embedding_fn,
+    )
+
+    queries = queries_vs.dataset[:10].text.data()["value"]
+    relevance = queries_vs.dataset[:10].metadata.data()["value"]
+    relevance = [rel["relevance"] for rel in relevance]
+
+    recall = db.deep_memory.evaluate(
+        queries=queries,
+        relevance=relevance,
+    )

deeplake/core/vectorstore/deeplake_vectorstore.py

@@ -8,6 +8,8 @@
 from deeplake.core.dataset import Dataset
 from deeplake.core.vectorstore.dataset_handlers import get_dataset_handler
 from deeplake.core.vectorstore.deep_memory import DeepMemory
+from deeplake.core.vectorstore.dataset_handlers import get_dataset_handler
+from deeplake.core.vectorstore.deep_memory import DeepMemory
 from deeplake.constants import (
     DEFAULT_VECTORSTORE_TENSORS,
     MAX_BYTES_PER_MINUTE,
@@ -131,7 +133,8 @@ def __init__(
         )
 
         self.deep_memory = DeepMemory(
-            dataset_or_path=self.dataset_handler.path,
+            dataset=self.dataset_handler.dataset,
+            path=self.dataset_handler.path,
             token=self.dataset_handler.token,
             logger=logger,
             embedding_function=embedding_function,
@@ -240,6 +243,7 @@ def search(
         return_tensors: Optional[List[str]] = None,
         return_view: bool = False,
         deep_memory: bool = False,
+        return_tql: bool = False,
     ) -> Union[Dict, Dataset]:
         """VectorStore search method that combines embedding search, metadata search, and custom TQL search.
 
@@ -290,6 +294,7 @@ def search(
             return_view (bool): Return a Deep Lake dataset view that satisfied the search parameters, instead of a dictionary with data. Defaults to False. If ``True`` return_tensors is set to "*" beucase data is lazy-loaded and there is no cost to including all tensors in the view.
             deep_memory (bool): Whether to use the Deep Memory model for improving search results. Defaults to False if deep_memory is not specified in the Vector Store initialization.
                 If True, the distance metric is set to "deepmemory_distance", which represents the metric with which the model was trained. The search is performed using the Deep Memory model. If False, the distance metric is set to "COS" or whatever distance metric user specifies.
+            return_tql (bool): Whether to return the TQL query string used for the search. Defaults to False.
 
         ..
             # noqa: DAR101
@@ -317,6 +322,7 @@ def search(
             embedding_tensor=embedding_tensor,
             return_tensors=return_tensors,
             return_view=return_view,
+            return_tql=return_tql,
             deep_memory=deep_memory,
         )