tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Releases 2 Packages Activity Actions

fix: support batching chromadb

Browse source
This commit is contained in:
Steven Kreitzer 2024-04-09 09:38:40 -05:00
parent 839efa4443
commit 0bae789d39
No known key found for this signature in database
GPG key ID: 76F902141A05C232
1 changed files with 9 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
backend/apps/rag/main.py
View file

@ -15,6 +15,7 @@ from typing import List
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
from chromadb.utils import embedding_functions from chromadb.utils import embedding_functions
from chromadb.utils.batch_utils import create_batches
from langchain_community.document_loaders import ( from langchain_community.document_loaders import (
WebBaseLoader, WebBaseLoader,
@ -331,9 +332,14 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
embedding_function=app.state.sentence_transformer_ef, embedding_function=app.state.sentence_transformer_ef,
) )
collection.add( for batch in create_batches(
documents=texts, metadatas=metadatas, ids=[str(uuid.uuid1()) for _ in texts] api=CHROMA_CLIENT,
) ids=[str(uuid.uuid1()) for _ in texts],
metadatas=metadatas,
documents=texts,
):
collection.add(*batch)
return True return True
except Exception as e: except Exception as e:
log.exception(e) log.exception(e)