Skip to content

Commit

Permalink
Update chromadb.py
Browse files Browse the repository at this point in the history
FIX for SuperboogaV2 when trying to unload databases "Clear Data" not working crashing textgen

There is an issue with textgen when trying to unload a superboogav2 database. Creating a database works, but unloading does not.

Issue created here:
oobabooga#6181
  • Loading branch information
RandomInternetPreson authored Jul 1, 2024
1 parent 31125eb commit aa036a5
Showing 1 changed file with 16 additions and 9 deletions.
25 changes: 16 additions & 9 deletions extensions/superboogav2/chromadb.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,11 @@ def should_merge(s1, s2, s1_start, s2_start):

class ChromaCollector():
def __init__(self):
name = ''.join(random.choice('ab') for _ in range(10))

name = 'static_test_name' # Use a static name for debugging
self.name = name
self.chroma_client = chromadb.Client(Settings(anonymized_telemetry=False))
self.collection = self.chroma_client.create_collection(name=name, embedding_function=embedder)
logger.info(f'Collection {self.name} created successfully.')

self.ids = []
self.id_to_info = {}
Expand Down Expand Up @@ -213,20 +213,24 @@ def _filter_outliers_by_median_distance(self, infos: list[Info], significant_lev
return filtered_infos

def _merge_infos(self, infos: list[Info]):
if not infos:
return [] # Return an empty list if there are no infos to merge

merged_infos = []
current_info = infos[0]

for next_info in infos[1:]:
merged = current_info.merge_with(next_info)
if merged is not None:
current_info = merged
else:
merged_infos.append(current_info)
current_info = next_info

merged_infos.append(current_info)
return merged_infos


# Main function for retrieving chunks by distance. It performs merging, time weighing, and mean filtering.

def _get_documents_ids_distances(self, search_strings: list[str], n_results: int):
Expand Down Expand Up @@ -335,13 +339,16 @@ def delete(self, ids_to_delete: list[str], where: dict):

def clear(self):
with self.lock:
self.chroma_client.reset()
try:
logger.info(f"Attempting to clear collection {self.name}. Current collections: {self.chroma_client.list_collections()}")
self.chroma_client.delete_collection(name=self.name)
logger.info(f"Collection {self.name} cleared successfully.")
self.collection = self.chroma_client.create_collection(name=self.name, embedding_function=embedder)
except Exception as e:
logger.error(f"Failed to clear collection: {e}. Available collections: {self.chroma_client.list_collections()}")


self.ids = []
self.chroma_client.delete_collection(name=self.name)
self.collection = self.chroma_client.create_collection(name=self.name, embedding_function=embedder)

logger.info('Successfully cleared all records and reset chromaDB.')


def make_collector():
Expand Down

0 comments on commit aa036a5

Please sign in to comment.