Spaces:
Runtime error
Runtime error
| """ | |
| Python Backend API to chat with private data | |
| 08/14/2023 | |
| D.M. Theekshana Samaradiwakara | |
| """ | |
| import os | |
| from dotenv import load_dotenv | |
| import glob | |
| import torch | |
| import pickle | |
| import io | |
| from langchain.vectorstores import Chroma | |
| from langchain.vectorstores import FAISS | |
| from langchain.embeddings import HuggingFaceEmbeddings | |
| from chromadb.config import Settings | |
| load_dotenv() | |
| import streamlit as st | |
| embeddings_model_name = os.environ.get("EMBEDDINGS_MODEL_NAME") | |
| embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name) | |
| def does_chroma_vectorstore_exist(persist_directory: str) -> bool: | |
| # Checks if vectorstore exists | |
| if os.path.exists(os.path.join(persist_directory, 'index')): | |
| if os.path.exists(os.path.join(persist_directory, 'chroma-collections.parquet')) and os.path.exists(os.path.join(persist_directory, 'chroma-embeddings.parquet')): | |
| list_index_files = glob.glob(os.path.join(persist_directory, 'index/*.bin')) | |
| list_index_files += glob.glob(os.path.join(persist_directory, 'index/*.pkl')) | |
| # At least 3 documents are needed in a working vectorstore | |
| if len(list_index_files) > 3: | |
| return True | |
| return False | |
| def load_store(directory: str) -> Chroma: | |
| index_path = "data/{0}".format(directory) | |
| # index_exists = os.path.exists(index_path) | |
| index_exists = does_chroma_vectorstore_exist(index_path) | |
| if index_exists: | |
| try: | |
| CHROMA_SETTINGS = Settings( | |
| chroma_db_impl='duckdb+parquet', | |
| persist_directory=index_path, | |
| anonymized_telemetry=False | |
| ) | |
| # return Chroma.load(index_path) | |
| vectorstore= Chroma( | |
| persist_directory=index_path, | |
| embedding_function=embeddings, | |
| client_settings=CHROMA_SETTINGS | |
| ) | |
| # with open("vectorstore.pkl", "wb") as f: | |
| # pickle.dump(vectorstore, f) | |
| return vectorstore | |
| except Exception as e: | |
| raise Exception(f"Error loading vector store: {e} ") | |
| else: | |
| # raise exception if model_type is not supported | |
| raise Exception(f"A vector store in directory {directory} is not created. Please choose a valid one") | |
| class CPU_Unpickler(pickle.Unpickler): | |
| def find_class(self, module, name): | |
| if module == 'torch.storage' and name == '_load_from_bytes': | |
| return lambda b: torch.load(io.BytesIO(b), map_location='cpu') | |
| else: | |
| return super().find_class(module, name) | |
| def create_db(document_splits,persist_directory): | |
| return Chroma.from_documents( | |
| documents=document_splits, | |
| embedding=embeddings, | |
| persist_directory=persist_directory | |
| ) | |
| def save_files(persist_directory, document_splits): | |
| print(f"Saving document splits...") | |
| if does_chroma_vectorstore_exist(persist_directory): | |
| print(f"Updating esisting vector store. May take some minutes...") | |
| #update function | |
| db = Chroma( | |
| persist_directory=index_path, | |
| embedding_function=embeddings, | |
| ) | |
| db.aadd_documents(document_splits) | |
| else: | |
| print(f"Creating new vector store. May take some minutes...") | |
| index_path = "data/{0}".format(persist_directory) | |
| db = create_db(document_splits,index_path) | |
| db.persist() | |