| """Read documents from disk and store them as embeddings in Chroma.""" | |
| import os | |
| import chromadb | |
| import config | |
| from embed import embed_many | |
| db = chromadb.Client() | |
| collection = db.create_collection("docs") | |
| def load_documents(folder=config.DOCS_DIR): | |
| """Read every .txt file in the folder into memory.""" | |
| docs = [] | |
| for name in os.listdir(folder): | |
| if name.endswith(".txt"): | |
| with open(os.path.join(folder, name)) as fh: | |
| docs.append(fh.read()) | |
| return docs | |
| def ingest(): | |
| """Load documents, embed them, and store them in the database.""" | |
| docs = load_documents() | |
| vectors = embed_many(docs) | |
| collection.add(documents=docs, embeddings=vectors, | |
| ids=[str(i) for i in range(len(docs))]) | |
| return len(docs) | |