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 <d@activeloop.ai>
doc
Davit Buniatyan 1 year ago committed by GitHub
parent 2ffb90b161
commit b4914888a7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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

@ -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 youre at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
"\n",
"Tonight, Id 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 nations top legal minds, who will continue Justice Breyers 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 <token>'\n"
]
}
],
"outputs": [],
"source": [
"!activeloop login -t <token>"
]
},
{
"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 youre at it, pass the Disclose Act so Americans can know who is funding our elections. \n",
"\n",
"Tonight, Id 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 nations top legal minds, who will continue Justice Breyers 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": {

@ -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."""

@ -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()

Loading…
Cancel
Save