Engg-SS_ChatBOT / initialize.py
abhivsh's picture
Update initialize.py
5a0035a verified
raw
history blame
728 Bytes
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import Chroma
import os
import gen_splits
GEMINI_API_KEY = os.environ.get('GEMINI_API_KEY')
# Using Google GenAI Text Embeddings
embedding_model = GoogleGenerativeAIEmbeddings(model="models/embedding-001", task_type="retrieval_document", google_api_key=GEMINI_API_KEY)
# Create Embeddings for Searching the Splits
persist_directory = './chroma/'
def initialize():
splits = gen_splits.gen_splits()
vectordb = Chroma.from_documents(documents=splits, persist_directory=persist_directory, embedding=embedding_model)
vectordb.persist()
return vectordb
if __name__ == "__main__":
vectordb = initialize()