From b4914888a781a17000eeb2335afc8276f6d08a58 Mon Sep 17 00:00:00 2001 From: Davit Buniatyan Date: Thu, 6 Apr 2023 12:47:33 -0700 Subject: [PATCH] Deep Lake upgrade to include attribute search, distance metrics, returning scores and MMR (#2455) ### Features include - Metadata based embedding search - Choice of distance metric function (`L2` for Euclidean, `L1` for Nuclear, `max` L-infinity distance, `cos` for cosine similarity, 'dot' for dot product. Defaults to `L2` - Returning scores - Max Marginal Relevance Search - Deleting samples from the dataset ### Notes - Added numerous tests, let me know if you would like to shorten them or make smarter --------- Co-authored-by: Davit Buniatyan --- docs/ecosystem/deeplake.md | 12 +- .../vectorstores/examples/deeplake.ipynb | 237 ++++++----- langchain/vectorstores/deeplake.py | 367 +++++++++++++++--- .../vectorstores/test_deeplake.py | 108 ++++++ 4 files changed, 574 insertions(+), 150 deletions(-) diff --git a/docs/ecosystem/deeplake.md b/docs/ecosystem/deeplake.md index 4f4e2151..13f3be0f 100644 --- a/docs/ecosystem/deeplake.md +++ b/docs/ecosystem/deeplake.md @@ -1,10 +1,14 @@ # Deep Lake - This page covers how to use the Deep Lake ecosystem within LangChain. -It is broken into two parts: installation and setup, and then references to specific Deep Lake wrappers. For more information. -1. Here is [whitepaper](https://www.deeplake.ai/whitepaper) and [academic paper](https://arxiv.org/pdf/2209.10785.pdf) for Deep Lake +## Why Deep Lake? +- More than just a (multi-modal) vector store. You can later use the dataset to fine-tune your own LLM models. +- Not only stores embeddings, but also the original data with automatic version control. +- Truly serverless. Doesn't require another service and can be used with major cloud providers (AWS S3, GCS, etc.) +## More Resources +1. [Ultimate Guide to LangChain & Deep Lake: Build ChatGPT to Answer Questions on Your Financial Data](https://www.activeloop.ai/resources/ultimate-guide-to-lang-chain-deep-lake-build-chat-gpt-to-answer-questions-on-your-financial-data/) +1. Here is [whitepaper](https://www.deeplake.ai/whitepaper) and [academic paper](https://arxiv.org/pdf/2209.10785.pdf) for Deep Lake 2. Here is a set of additional resources available for review: [Deep Lake](https://github.com/activeloopai/deeplake), [Getting Started](https://docs.activeloop.ai/getting-started) and [Tutorials](https://docs.activeloop.ai/hub-tutorials) ## Installation and Setup @@ -14,7 +18,7 @@ It is broken into two parts: installation and setup, and then references to spec ### VectorStore -There exists a wrapper around Deep Lake, a data lake for Deep Learning applications, allowing you to use it as a vectorstore (for now), whether for semantic search or example selection. +There exists a wrapper around Deep Lake, a data lake for Deep Learning applications, allowing you to use it as a vector store (for now), whether for semantic search or example selection. To import this vectorstore: ```python diff --git a/docs/modules/indexes/vectorstores/examples/deeplake.ipynb b/docs/modules/indexes/vectorstores/examples/deeplake.ipynb index 3fc5e32c..3ec489ac 100644 --- a/docs/modules/indexes/vectorstores/examples/deeplake.ipynb +++ b/docs/modules/indexes/vectorstores/examples/deeplake.ipynb @@ -13,7 +13,16 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!python3 -m pip install openai deeplake" + ] + }, + { + "cell_type": "code", + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -25,11 +34,22 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "os.environ['OPENAI_API_KEY'] = 'sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx'" + ] + }, + { + "cell_type": "code", + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from langchain.document_loaders import TextLoader\n", + "\n", "loader = TextLoader('../../../state_of_the_union.txt')\n", "documents = loader.load()\n", "text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)\n", @@ -40,17 +60,9 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Evaluating ingest: 100%|██████████| 41/41 [00:00<00:00\n" - ] - } - ], + "outputs": [], "source": [ "db = DeepLake.from_documents(docs, embeddings)\n", "\n", @@ -60,73 +72,136 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n", - "\n", - "We cannot let this happen. \n", - "\n", - "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n", - "\n", - "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n", - "\n", - "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", - "\n", - "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n" - ] - } - ], + "outputs": [], "source": [ "print(docs[0].page_content)" ] }, { + "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ - "## Deep Lake datasets on cloud or local\n", + "### Retrieval Question/Answering" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from langchain.chains import RetrievalQA\n", + "from langchain.llms import OpenAIChat\n", + "\n", + "qa = RetrievalQA.from_chain_type(llm=OpenAIChat(model='gpt-3.5-turbo'), chain_type='stuff', retriever=db.as_retriever())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "query = 'What did the president say about Ketanji Brown Jackson'\n", + "qa.run(query)" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Attribute based filtering in metadata" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import random\n", + "\n", + "for d in docs:\n", + " d.metadata['year'] = random.randint(2012, 2014)\n", + "\n", + "db = DeepLake.from_documents(docs, embeddings)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "db.similarity_search('What did the president say about Ketanji Brown Jackson', filter={'year': 2013})" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Choosing distance function\n", + "Distance function `L2` for Euclidean, `L1` for Nuclear, `Max` l-infinity distnace, `cos` for cosine similarity, `dot` for dot product " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "db.similarity_search('What did the president say about Ketanji Brown Jackson?', distance_metric='cos')" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Maximal Marginal relevance\n", + "Using maximal marginal relevance" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "db.max_marginal_relevance_search('What did the president say about Ketanji Brown Jackson?')" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Deep Lake datasets on cloud (Activeloop, AWS, GCS, etc.) or local\n", "By default deep lake datasets are stored in memory, in case you want to persist locally or to any object storage you can simply provide path to the dataset. You can retrieve token from [app.activeloop.ai](https://app.activeloop.ai/)" ] }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "/bin/bash: -c: line 0: syntax error near unexpected token `newline'\n", - "/bin/bash: -c: line 0: `activeloop login -t '\n" - ] - } - ], + "outputs": [], "source": [ "!activeloop login -t " ] }, { "cell_type": "code", - "execution_count": 20, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Evaluating ingest: 100%|██████████| 4/4 [00:00<00:00\n" - ] - } - ], + "outputs": [], "source": [ "# Embed and store the texts\n", - "dataset_path = \"hub://{username}/{dataset_name}\" # could be also ./local/path (much faster locally), s3://bucket/path/to/dataset, gcs://, etc.\n", + "dataset_path = \"hub://{username}/{dataset_name}\" # could be also ./local/path (much faster locally), s3://bucket/path/to/dataset, gcs://path/to/dataset, etc.\n", "\n", "embedding = OpenAIEmbeddings()\n", "vectordb = DeepLake.from_documents(documents=docs, embedding=embedding, dataset_path=dataset_path)" @@ -134,27 +209,9 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "In state after state, new laws have been passed, not only to suppress the vote, but to subvert entire elections. \n", - "\n", - "We cannot let this happen. \n", - "\n", - "Tonight. I call on the Senate to: Pass the Freedom to Vote Act. Pass the John Lewis Voting Rights Act. And while you’re at it, pass the Disclose Act so Americans can know who is funding our elections. \n", - "\n", - "Tonight, I’d like to honor someone who has dedicated his life to serve this country: Justice Stephen Breyer—an Army veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court. Justice Breyer, thank you for your service. \n", - "\n", - "One of the most serious constitutional responsibilities a President has is nominating someone to serve on the United States Supreme Court. \n", - "\n", - "And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.\n" - ] - } - ], + "outputs": [], "source": [ "query = \"What did the president say about Ketanji Brown Jackson\"\n", "docs = db.similarity_search(query)\n", @@ -163,43 +220,21 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Dataset(path='./local/path', tensors=['embedding', 'ids', 'metadata', 'text'])\n", - "\n", - " tensor htype shape dtype compression\n", - " ------- ------- ------- ------- ------- \n", - " embedding generic (4, 1536) None None \n", - " ids text (4, 1) str None \n", - " metadata json (4, 1) str None \n", - " text text (4, 1) str None \n" - ] - } - ], + "outputs": [], "source": [ "vectordb.ds.summary()" ] }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ "embeddings = vectordb.ds.embedding.numpy()" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { @@ -218,7 +253,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.1" + "version": "3.10.0" }, "vscode": { "interpreter": { diff --git a/langchain/vectorstores/deeplake.py b/langchain/vectorstores/deeplake.py index f4a7ddf5..11419739 100644 --- a/langchain/vectorstores/deeplake.py +++ b/langchain/vectorstores/deeplake.py @@ -3,40 +3,76 @@ from __future__ import annotations import logging import uuid -from typing import Any, Iterable, List, Optional, Sequence +from functools import partial +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple import numpy as np from langchain.docstore.document import Document from langchain.embeddings.base import Embeddings from langchain.vectorstores.base import VectorStore +from langchain.vectorstores.utils import maximal_marginal_relevance logger = logging.getLogger() +distance_metric_map = { + "l2": lambda a, b: np.linalg.norm(a - b, axis=1, ord=2), + "l1": lambda a, b: np.linalg.norm(a - b, axis=1, ord=1), + "max": lambda a, b: np.linalg.norm(a - b, axis=1, ord=np.inf), + "cos": lambda a, b: np.dot(a, b.T) + / (np.linalg.norm(a) * np.linalg.norm(b, axis=1)), + "dot": lambda a, b: np.dot(a, b.T), +} + + +def vector_search( + query_embedding: np.ndarray, + data_vectors: np.ndarray, + distance_metric: str = "L2", + k: Optional[int] = 4, +) -> Tuple[List, List]: + """Naive search for nearest neighbors + + args: + query_embedding: np.ndarray + data_vectors: np.ndarray + k (int): number of nearest neighbors + distance_metric: distance function 'L2' for Euclidean, 'L1' for Nuclear, 'Max' + l-infinity distnace, 'cos' for cosine similarity, 'dot' for dot product + + returns: + nearest_indices: List, indices of nearest neighbors + """ + # Calculate the distance between the query_vector and all data_vectors + distances = distance_metric_map[distance_metric](query_embedding, data_vectors) + nearest_indices = np.argsort(distances) + nearest_indices = ( + nearest_indices[::-1][:k] if distance_metric in ["cos"] else nearest_indices[:k] + ) + + return nearest_indices.tolist(), distances[nearest_indices].tolist() -def L2_search( - query_embedding: np.ndarray, data_vectors: np.ndarray, k: int = 4 -) -> list: - """naive L2 search for nearest neighbors""" - # Calculate the L2 distance between the query_vector and all data_vectors - distances = np.linalg.norm(data_vectors - query_embedding, axis=1) - # Sort the distances and return the indices of the k nearest vectors - nearest_indices = np.argsort(distances)[:k] - return nearest_indices.tolist() +def dp_filter(x: dict, filter: Dict[str, str]) -> bool: + """Filter helper function for Deep Lake""" + metadata = x["metadata"].data()["value"] + return all(k in metadata and v == metadata[k] for k, v in filter.items()) class DeepLake(VectorStore): """Wrapper around Deep Lake, a data lake for deep learning applications. - It not only stores embeddings, but also the original data and queries with - version control automatically enabled. + We implement naive similarity search and filtering for fast prototyping, + but it can be extended with Tensor Query Language (TQL) for production use cases + over billion rows. - It is more than just a vector store. You can use the dataset to fine-tune - your own LLM models or use it for other downstream tasks. + Why Deep Lake? - We implement naive similiarity search, but it can be extended with Tensor - Query Language (TQL for production use cases) over billion rows. + - Not only stores embeddings, but also the original data with version control. + - Serverless, doesn't require another service and can be used with major + cloud providers (S3, GCS, etc.) + - More than just a multi-modal vector store. You can use the dataset + to fine-tune your own LLM models. To use, you should have the ``deeplake`` python package installed. @@ -80,10 +116,34 @@ class DeepLake(VectorStore): else: self.ds = deeplake.empty(dataset_path, token=token, overwrite=True) with self.ds: - self.ds.create_tensor("text", htype="text") - self.ds.create_tensor("metadata", htype="json") - self.ds.create_tensor("embedding", htype="generic") - self.ds.create_tensor("ids", htype="text") + self.ds.create_tensor( + "text", + htype="text", + create_id_tensor=False, + create_sample_info_tensor=False, + create_shape_tensor=False, + ) + self.ds.create_tensor( + "metadata", + htype="json", + create_id_tensor=False, + create_sample_info_tensor=False, + create_shape_tensor=False, + ) + self.ds.create_tensor( + "embedding", + htype="generic", + create_id_tensor=False, + create_sample_info_tensor=False, + create_shape_tensor=False, + ) + self.ds.create_tensor( + "ids", + htype="text", + create_id_tensor=False, + create_sample_info_tensor=False, + create_shape_tensor=False, + ) self._embedding_function = embedding_function @@ -116,11 +176,9 @@ class DeepLake(VectorStore): embeddings = self._embedding_function.embed_documents(text_list) if metadatas is None: - metadatas_to_use: Sequence[Optional[dict]] = [None] * len(text_list) - else: - metadatas_to_use = metadatas + metadatas = [{}] * len(text_list) - elements = zip(text_list, embeddings, metadatas_to_use, ids) + elements = zip(text_list, embeddings, metadatas, ids) @self._deeplake.compute def ingest(sample_in: list, sample_out: list) -> None: @@ -133,32 +191,210 @@ class DeepLake(VectorStore): sample_out.append(s) ingest().eval(list(elements), self.ds) - self.ds.commit() + self.ds.commit(allow_empty=True) return ids - def similarity_search( - self, query: str, k: int = 4, **kwargs: Any - ) -> List[Document]: - """Return docs most similar to query.""" + def search( + self, + query: Any[str, None] = None, + embedding: Any[float, None] = None, + k: int = 4, + distance_metric: str = "L2", + use_maximal_marginal_relevance: Optional[bool] = False, + fetch_k: Optional[int] = 20, + filter: Optional[Dict[str, str]] = None, + return_score: Optional[bool] = False, + **kwargs: Any, + ) -> Any[List[Document], List[Tuple[Document, float]]]: + """Return docs most similar to query. + + Args: + query: Text to look up documents similar to. + embedding: Embedding function to use. Defaults to None. + k: Number of Documents to return. Defaults to 4. + distance_metric: `L2` for Euclidean, `L1` for Nuclear, + `max` L-infinity distance, `cos` for cosine similarity, + 'dot' for dot product. Defaults to `L2`. + filter: Attribute filter by metadata example {'key': 'value'}. + Defaults to None. + maximal_marginal_relevance: Whether to use maximal marginal relevance. + Defaults to False. + fetch_k: Number of Documents to fetch to pass to MMR algorithm. + Defaults to 20. + return_score: Whether to return the score. Defaults to False. + + Returns: + List of Documents selected by the specified distance metric, + if return_score True, return a tuple of (Document, score) + """ + view = self.ds + + # attribute based filtering + if filter is not None: + view = view.filter(partial(dp_filter, filter=filter)) + + if len(view) == 0: + return [] + if self._embedding_function is None: - self.ds.summary() - ds_view = self.ds.filter(lambda x: query in x["text"].data()["value"]) + view = view.filter(lambda x: query in x["text"].data()["value"]) + scores = [1.0] * len(view) + + if use_maximal_marginal_relevance: + raise ValueError( + "For MMR search, you must specify an embedding function on" + "creation." + ) + else: - query_emb = np.array(self._embedding_function.embed_query(query)) - embeddings = self.ds.embedding.numpy() - indices = L2_search(query_emb, embeddings, k=k) - ds_view = self.ds[indices] + emb = embedding or self._embedding_function.embed_query( + query + ) # type: ignore + query_emb = np.array(emb, dtype=np.float32) + embeddings = view.embedding.numpy() + + k_search = fetch_k if use_maximal_marginal_relevance else k + indices, scores = vector_search( + query_emb, + embeddings, + k=k_search, + distance_metric=distance_metric.lower(), + ) + view = view[indices] + + if use_maximal_marginal_relevance: + indices = maximal_marginal_relevance( + query_emb, embeddings[indices], k=min(k, len(indices)) + ) + view = view[indices] + scores = [scores[i] for i in indices] docs = [ Document( page_content=el["text"].data()["value"], metadata=el["metadata"].data()["value"], ) - for el in ds_view + for el in view ] + + if return_score: + return [(doc, score) for doc, score in zip(docs, scores)] + return docs + def similarity_search( + self, query: str, k: int = 4, **kwargs: Any + ) -> List[Document]: + """Return docs most similar to query. + + Args: + query: text to embed and run the query on. + k: Number of Documents to return. + Defaults to 4. + query: Text to look up documents similar to. + embedding: Embedding function to use. + Defaults to None. + k: Number of Documents to return. + Defaults to 4. + distance_metric: `L2` for Euclidean, `L1` for Nuclear, `max` + L-infinity distance, `cos` for cosine similarity, 'dot' for dot product + Defaults to `L2`. + filter: Attribute filter by metadata example {'key': 'value'}. + Defaults to None. + maximal_marginal_relevance: Whether to use maximal marginal relevance. + Defaults to False. + fetch_k: Number of Documents to fetch to pass to MMR algorithm. + Defaults to 20. + return_score: Whether to return the score. Defaults to False. + + Returns: + List of Documents most similar to the query vector. + """ + return self.search(query=query, k=k, **kwargs) + + def similarity_search_by_vector( + self, embedding: List[float], k: int = 4, **kwargs: Any + ) -> List[Document]: + """Return docs most similar to embedding vector. + + Args: + embedding: Embedding to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + Returns: + List of Documents most similar to the query vector. + """ + return self.search(embedding=embedding, k=k, **kwargs) + + def similarity_search_with_score( + self, + query: str, + distance_metric: str = "L2", + k: int = 4, + filter: Optional[Dict[str, str]] = None, + ) -> List[Tuple[Document, float]]: + """Run similarity search with Deep Lake with distance returned. + + Args: + query (str): Query text to search for. + distance_metric: `L2` for Euclidean, `L1` for Nuclear, `max` L-infinity + distance, `cos` for cosine similarity, 'dot' for dot product. + Defaults to `L2`. + k (int): Number of results to return. Defaults to 4. + filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None. + Returns: + List[Tuple[Document, float]]: List of documents most similar to the query + text with distance in float. + """ + return self.search( + query=query, + k=k, + filter=filter, + return_score=True, + distance_metric=distance_metric, + ) + + def max_marginal_relevance_search_by_vector( + self, embedding: List[float], k: int = 4, fetch_k: int = 20 + ) -> List[Document]: + """Return docs selected using the maximal marginal relevance. + Maximal marginal relevance optimizes for similarity to query AND diversity + among selected documents. + Args: + embedding: Embedding to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + fetch_k: Number of Documents to fetch to pass to MMR algorithm. + Returns: + List of Documents selected by maximal marginal relevance. + """ + return self.search( + embedding=embedding, + k=k, + fetch_k=fetch_k, + use_maximal_marginal_relevance=True, + ) + + def max_marginal_relevance_search( + self, query: str, k: int = 4, fetch_k: int = 20 + ) -> List[Document]: + """Return docs selected using the maximal marginal relevance. + Maximal marginal relevance optimizes for similarity to query AND diversity + among selected documents. + Args: + query: Text to look up documents similar to. + k: Number of Documents to return. Defaults to 4. + fetch_k: Number of Documents to fetch to pass to MMR algorithm. + Returns: + List of Documents selected by maximal marginal relevance. + """ + if self._embedding_function is None: + raise ValueError( + "For MMR search, you must specify an embedding function on" "creation." + ) + return self.search( + query=query, k=k, fetch_k=fetch_k, use_maximal_marginal_relevance=True + ) + @classmethod def from_texts( cls, @@ -171,22 +407,24 @@ class DeepLake(VectorStore): ) -> DeepLake: """Create a Deep Lake dataset from a raw documents. - If a persist_directory is specified, the collection will be persisted there. + If a dataset_path is specified, the dataset will be persisted there. Otherwise, the data will be ephemeral in-memory. Args: path (str, pathlib.Path): - The full path to the dataset. Can be: - - a Deep Lake cloud path of the form ``hub://username/datasetname``. + - Deep Lake cloud path of the form ``hub://username/dataset_name``. To write to Deep Lake cloud datasets, ensure that you are logged in to Deep Lake (use 'activeloop login' from command line) - - an s3 path of the form ``s3://bucketname/path/to/dataset``. - Credentials are required in either the environment or - passed to the creds argument. - - a local file system path of the form ``./path/to/dataset`` or + - AWS S3 path of the form ``s3://bucketname/path/to/dataset``. + Credentials are required in either the environment + - Google Cloud Storage path of the form + ``gcs://bucketname/path/to/dataset``Credentials are required + in either the environment + - Local file system path of the form ``./path/to/dataset`` or ``~/path/to/dataset`` or ``path/to/dataset``. - - a memory path of the form ``mem://path/to/dataset`` which doesn't - save the dataset but keeps it in memory instead. + - In-memory path of the form ``mem://path/to/dataset`` which doesn't + save the dataset, but keeps it in memory instead. Should be used only for testing as it does not persist. documents (List[Document]): List of documents to add. embedding (Optional[Embeddings]): Embedding function. Defaults to None. @@ -203,9 +441,48 @@ class DeepLake(VectorStore): deeplake_dataset.add_texts(texts=texts, metadatas=metadatas, ids=ids) return deeplake_dataset + def delete( + self, + ids: Any[List[str], None] = None, + filter: Any[Dict[str, str], None] = None, + delete_all: Any[bool, None] = None, + ) -> bool: + """Delete the entities in the dataset + + Args: + ids (Optional[List[str]], optional): The document_ids to delete. + Defaults to None. + filter (Optional[Dict[str, str]], optional): The filter to delete by. + Defaults to None. + delete_all (Optional[bool], optional): Whether to drop the dataset. + Defaults to None. + """ + if delete_all: + self.ds.delete() + return True + + view = None + if ids: + view = self.ds.filter(lambda x: x["ids"].data()["value"] in ids) + ids = list(view.sample_indices) + + if filter: + if view is None: + view = self.ds + view = view.filter(partial(dp_filter, filter=filter)) + ids = list(view.sample_indices) + + with self.ds: + for id in sorted(ids)[::-1]: + self.ds.pop(id) + + self.ds.commit(f"deleted {len(ids)} samples", allow_empty=True) + + return True + def delete_dataset(self) -> None: """Delete the collection.""" - self.ds.delete() + self.delete(delete_all=True) def persist(self) -> None: """Persist the collection.""" diff --git a/tests/integration_tests/vectorstores/test_deeplake.py b/tests/integration_tests/vectorstores/test_deeplake.py index a8316f21..2463a4bb 100644 --- a/tests/integration_tests/vectorstores/test_deeplake.py +++ b/tests/integration_tests/vectorstores/test_deeplake.py @@ -1,9 +1,31 @@ """Test Deep Lake functionality.""" +import deeplake +import pytest +from pytest import FixtureRequest + from langchain.docstore.document import Document from langchain.vectorstores import DeepLake from tests.integration_tests.vectorstores.fake_embeddings import FakeEmbeddings +@pytest.fixture +def deeplake_datastore() -> DeepLake: + texts = ["foo", "bar", "baz"] + metadatas = [{"page": str(i)} for i in range(len(texts))] + docsearch = DeepLake.from_texts( + dataset_path="mem://test_path", + texts=texts, + metadatas=metadatas, + embedding=FakeEmbeddings(), + ) + return docsearch + + +@pytest.fixture(params=["L1", "L2", "max", "cos"]) +def distance_metric(request: FixtureRequest) -> str: + return request.param + + def test_deeplake() -> None: """Test end to end construction and search.""" texts = ["foo", "bar", "baz"] @@ -31,6 +53,9 @@ def test_deeplake_with_metadatas() -> None: def test_deeplakewith_persistence() -> None: """Test end to end construction and search, with persistence.""" dataset_path = "./tests/persist_dir" + if deeplake.exists(dataset_path): + deeplake.delete(dataset_path) + texts = ["foo", "bar", "baz"] docsearch = DeepLake.from_texts( dataset_path=dataset_path, @@ -56,3 +81,86 @@ def test_deeplakewith_persistence() -> None: # Persist doesn't need to be called again # Data will be automatically persisted on object deletion # Or on program exit + + +def test_similarity_search(deeplake_datastore: DeepLake, distance_metric: str) -> None: + """Test similarity search.""" + output = deeplake_datastore.similarity_search( + "foo", k=1, distance_metric=distance_metric + ) + assert output == [Document(page_content="foo", metadata={"page": "0"})] + deeplake_datastore.delete_dataset() + + +def test_similarity_search_by_vector( + deeplake_datastore: DeepLake, distance_metric: str +) -> None: + """Test similarity search by vector.""" + embeddings = FakeEmbeddings().embed_documents(["foo", "bar", "baz"]) + output = deeplake_datastore.similarity_search_by_vector( + embeddings[1], k=1, distance_metric=distance_metric + ) + assert output == [Document(page_content="bar", metadata={"page": "1"})] + deeplake_datastore.delete_dataset() + + +def test_similarity_search_with_score( + deeplake_datastore: DeepLake, distance_metric: str +) -> None: + """Test similarity search with score.""" + output, score = deeplake_datastore.similarity_search_with_score( + "foo", k=1, distance_metric=distance_metric + )[0] + assert output == Document(page_content="foo", metadata={"page": "0"}) + if distance_metric == "cos": + assert score == 1.0 + else: + assert score == 0.0 + deeplake_datastore.delete_dataset() + + +def test_similarity_search_with_filter( + deeplake_datastore: DeepLake, distance_metric: str +) -> None: + """Test similarity search.""" + + output = deeplake_datastore.similarity_search( + "foo", k=1, distance_metric=distance_metric, filter={"page": "1"} + ) + assert output == [Document(page_content="bar", metadata={"page": "1"})] + deeplake_datastore.delete_dataset() + + +def test_max_marginal_relevance_search(deeplake_datastore: DeepLake) -> None: + """Test max marginal relevance search by vector.""" + + output = deeplake_datastore.max_marginal_relevance_search("foo", k=1, fetch_k=2) + + assert output == [Document(page_content="foo", metadata={"page": "0"})] + + embeddings = FakeEmbeddings().embed_documents(["foo", "bar", "baz"]) + output = deeplake_datastore.max_marginal_relevance_search_by_vector( + embeddings[0], k=1, fetch_k=2 + ) + + assert output == [Document(page_content="foo", metadata={"page": "0"})] + deeplake_datastore.delete_dataset() + + +def test_delete_dataset_by_ids(deeplake_datastore: DeepLake) -> None: + """Test delete dataset.""" + id = deeplake_datastore.ds.ids.data()["value"][0] + deeplake_datastore.delete(ids=[id]) + assert deeplake_datastore.similarity_search("foo", k=1, filter={"page": "0"}) == [] + assert len(deeplake_datastore.ds) == 2 + + deeplake_datastore.delete_dataset() + + +def test_delete_dataset_by_filter(deeplake_datastore: DeepLake) -> None: + """Test delete dataset.""" + deeplake_datastore.delete(filter={"page": "1"}) + assert deeplake_datastore.similarity_search("bar", k=1, filter={"page": "1"}) == [] + assert len(deeplake_datastore.ds) == 2 + + deeplake_datastore.delete_dataset()