Harrison/faiss merge (#1843)

Co-authored-by: Ting Su <ting.su.1995@outlook.com>
This commit is contained in:
Harrison Chase 2023-03-20 22:54:08 -07:00 committed by GitHub
parent 6f02286805
commit 1c22657256
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 127 additions and 5 deletions

View File

@ -62,10 +62,6 @@
"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",
@ -200,10 +196,104 @@
"docs[0]"
]
},
{
"cell_type": "markdown",
"id": "57da60d4",
"metadata": {},
"source": [
"## Merging\n",
"You can also merge two FAISS vectorstores"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "6dfd2b78",
"metadata": {},
"outputs": [],
"source": [
"db1 = FAISS.from_texts([\"foo\"], embeddings)\n",
"db2 = FAISS.from_texts([\"bar\"], embeddings)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "29960da7",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'e0b74348-6c93-4893-8764-943139ec1d17': Document(page_content='foo', lookup_str='', metadata={}, lookup_index=0)}"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"db1.docstore._dict"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "83392605",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'bdc50ae3-a1bb-4678-9260-1b0979578f40': Document(page_content='bar', lookup_str='', metadata={}, lookup_index=0)}"
]
},
"execution_count": 9,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"db2.docstore._dict"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "a3fcc1c7",
"metadata": {},
"outputs": [],
"source": [
"db1.merge_from(db2)"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "41c51f89",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"{'e0b74348-6c93-4893-8764-943139ec1d17': Document(page_content='foo', lookup_str='', metadata={}, lookup_index=0),\n",
" 'd5211050-c777-493d-8825-4800e74cfdb6': Document(page_content='bar', lookup_str='', metadata={}, lookup_index=0)}"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"db1.docstore._dict"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bc8b71f7",
"id": "f80b60de",
"metadata": {},
"outputs": [],
"source": []

View File

@ -221,6 +221,38 @@ class FAISS(VectorStore):
docs = self.max_marginal_relevance_search_by_vector(embedding, k, fetch_k)
return docs
def merge_from(self, target: FAISS) -> None:
"""Merge another FAISS object with the current one.
Add the target FAISS to the current one.
Args:
target: FAISS object you wish to merge into the current one
Returns:
None.
"""
if not isinstance(self.docstore, AddableMixin):
raise ValueError("Cannot merge with this type of docstore")
# Numerical index for target docs are incremental on existing ones
starting_len = len(self.index_to_docstore_id)
# Merge two IndexFlatL2
self.index.merge_from(target.index)
# Create new id for docs from target FAISS object
full_info = []
for i in target.index_to_docstore_id:
doc = target.docstore.search(target.index_to_docstore_id[i])
if not isinstance(doc, Document):
raise ValueError("Document should be returned")
full_info.append((starting_len + i, str(uuid.uuid4()), doc))
# Add information to docstore and index_to_docstore_id.
self.docstore.add({_id: doc for _, _id, doc in full_info})
index_to_id = {index: _id for index, _id, _ in full_info}
self.index_to_docstore_id.update(index_to_id)
@classmethod
def from_texts(
cls,