generated from phoughton/python_package_template
-
Notifications
You must be signed in to change notification settings - Fork 1
/
db_pop.py
24 lines (19 loc) · 758 Bytes
/
db_pop.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
import chromadb
from textract import extract_from_txt, get_file_names, Document
import config
books = get_file_names(config.DATA_IN_LOCATION, ".txt")
chroma_client = chromadb.PersistentClient(path=config.DB_LOCATION)
collection = chroma_client.get_or_create_collection(
name=config.DB_NAME,
metadata={"hnsw:search_ef": 20, "hnsw:construction_ef": 200})
for book_name in books:
print(f"Parsing: {book_name}")
book = Document(config.DATA_IN_LOCATION, book_name,
extractor=extract_from_txt)
print(f"Adding: {book_name}")
print(f"Book chunks: {len(book.get_book_chunks())}")
collection.add(
documents=book.get_book_chunks(),
metadatas=book.get_metadata(),
ids=book.get_ids()
)