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

Merge pull request #1476 from buroa/dev

Browse source 
fix: support batching chromadb
This commit is contained in:
Timothy Jaeryang Baek 2024-04-09 22:53:51 -07:00 committed by GitHub
commit 2f6e683163
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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)