Spaces:
Sleeping
Sleeping
File size: 3,287 Bytes
aa879af e80787b aa879af e80787b aa879af cfbc55e aa879af e80787b aa879af cfbc55e e80787b aa879af e80787b aa879af e80787b aa879af e80787b cfbc55e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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()
|