Spaces:
Sleeping
Sleeping
| from google import genai | |
| from google.genai import types | |
| import langchain_chroma as chromaDB | |
| import config | |
| class GooglePalmEmbeddings: | |
| def __init__(self): | |
| self.client = genai.Client(api_key=config.API_KEY) | |
| def embed_documents(self, documents): | |
| embeddings = (self.client.models.embed_content( | |
| model="models/text-embedding-004", | |
| contents=documents, | |
| config=types.EmbedContentConfig(task_type='retrieval_document'))) | |
| return [e.values for e in embeddings.embeddings] | |
| def embed_query(self, text): | |
| embeddings = self.client.models.embed_content( | |
| model="models/text-embedding-004", | |
| contents=text, | |
| config=types.EmbedContentConfig(task_type='retrieval_query')) | |
| return embeddings.embeddings[0].values | |
| class ChromaDB: | |
| def __init__(self, embeddingFunction): | |
| print(config.session) | |
| self.vector_db = chromaDB.Chroma(persist_directory=f'{str(config.session)}/vectorData', | |
| embedding_function=embeddingFunction) | |
| def addEmbeddings_to_Chroma(self, chunks): | |
| doc_name = chunks[0].metadata.get("source") | |
| existing_items = self.vector_db.get(include=[]) | |
| existing_ids = set(existing_items["ids"]) | |
| print(f"Number of existing documents in DB: {len(existing_ids)}") | |
| print("Adding embeddings to ChromaDB collection for : ", doc_name) | |
| new_chunks = [] | |
| for chunk in chunks: | |
| if chunk.metadata["id"] not in existing_ids: | |
| new_chunks.append(chunk) | |
| if len(new_chunks) != 0: | |
| chunk_ids = [chunk.metadata["id"] for chunk in new_chunks] | |
| self.vector_db.add_documents(new_chunks, ids=chunk_ids) | |
| print(f"Added {len(new_chunks)} new documents.") | |
| print("Added embeddings to ChromaDB.") | |
| else: | |
| print("No new documents to add.") | |
| def observeDB(self): | |
| print("Document IDs of embeddings in database: ", self.vector_db.get(include=[])["ids"]) | |
| def deleteEmbeddings(self, filename): | |
| ids = self.vector_db.get(include=[])["ids"] | |
| to_delete = [id for id in ids if filename in id] | |
| self.vector_db.delete(to_delete) | |
| print(f"Deleted {filename} embeddings from ChromaDB.") | |
| # ======================= Test ============================= # | |
| # from processDocument import load_file_and_split, createChunkID | |
| # config.API_KEY = "your_api_key_here" | |
| # config.API_KEY = "your_api_key" | |
| # # Loading the file and splitting it into chunks | |
| # split_docs = load_file_and_split(r"./assets/AttentionPaper.pdf") | |
| # chunks1 = createChunkID(split_docs) | |
| # split_docs = load_file_and_split(r"./assets/Cover letter.pdf") | |
| # chunks2 = createChunkID(split_docs) | |
| # chunks = chunks1 + chunks2 | |
| # print("Number of chunks: \n", len(chunks)) | |
| # # print(chunks[0]) | |
| # # Creating the embeddings and adding them to the ChromaDB | |
| # db = ChromaDB(GooglePalmEmbeddings()) | |
| # for i in range(0, len(chunks), 100): | |
| # # print(f"Adding chunks {i} to {i+99}") | |
| # db.addEmbeddings_to_Chroma(chunks[i:i+100]) | |
| # db.observeDB() | |
| # db.deleteEmbeddings("AttentionPaper.pdf") | |
| # db.deleteEmbeddings("Cover letter.pdf") | |
| # db.observeDB() | |