-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathgenerate_embeddings.py
31 lines (22 loc) · 1.13 KB
/
generate_embeddings.py
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
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
DOC_PATH = "docs"
DB_PATH = "chromadb"
embedding_model_name="nomic-embed-text"
loader = TextLoader('docs/ewu.txt', encoding = 'UTF-8')
documents = loader.load()
chunk_size=3900
chunk_overlap=500
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
docs = text_splitter.split_documents(documents)
embeddings = OllamaEmbeddings(model = embedding_model_name)
print('| ---------------------------------------------------------|')
print(f'| Chunk size:{chunk_size}\n| Chunk overlap:{chunk_overlap}')
print(f'| {len(docs)} chunks are being embedded now, processing...')
vectordb = Chroma.from_documents(documents = docs, embedding=embeddings, persist_directory=DB_PATH)
vectordb.persist()
if vectordb:
print(f"| Embeddings are generated with chunks: {len(docs)}")
print('| ---------------------------------------------------------')