Spaces:
Running on Zero
Running on Zero
File size: 2,692 Bytes
9273228 | 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 89 90 91 92 | # Import Libraries
import os
import glob
import torch
from pathlib import Path
from dotenv import load_dotenv
from huggingface_hub import login
from langchain_chroma import Chroma
from langchain_community.document_loaders import (
DirectoryLoader,
TextLoader
)
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
# Create HF_TOKEN and EMBEDDING_MODELS
load_dotenv(override=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
HF_TOKEN = os.getenv("HF_TOKEN")
EMBEDDING_MODELS = HuggingFaceEmbeddings(
model_name=os.getenv("EMBEDDING_MODELS"),
model_kwargs={"device": device},
encode_kwargs={"normalize_embeddings": True}
)
login(token=HF_TOKEN,
add_to_git_credential=True)
# Create path for DB_NAME and KNOWLEDGE_BASE
DB_NAME = str(Path(__file__).parent.parent/"vector_db")
KNOWLEDGE_BASE = str(Path(__file__).parent.parent/"python_doc_md")
def fetch_documents(): # Create fetch_documents function
documents = []
folders = glob.glob(str(Path(KNOWLEDGE_BASE)/"*"))
for folder in folders:
doc_type = os.path.basename(folder)
loader = DirectoryLoader(
path=folder,
glob="**/*.md",
loader_cls=TextLoader,
loader_kwargs={'encoding': 'utf-8'}
)
folder_docs = loader.load()
for doc in folder_docs:
doc.metadata["doc_type"] = doc_type
doc.page_content = f"{doc_type}\n\n{doc.page_content}"
documents.append(doc)
return documents
def create_chunks(documents): # Create create_chunks function
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=700,
chunk_overlap=100,
)
chunks = text_splitter.split_documents(documents=documents)
return chunks
def create_embedding(chunks):
if os.path.exists(DB_NAME):
Chroma(persist_directory=DB_NAME,
embedding_function=EMBEDDING_MODELS).delete_collection()
vector_store = Chroma.from_documents(
documents=chunks,
embedding=EMBEDDING_MODELS,
persist_directory=DB_NAME
)
collection = vector_store._collection
count = collection.count()
sample_embedding = collection.get(
limit=1,
include=["embeddings"]
)["embeddings"][0]
dimension = len(sample_embedding)
print(
f"There are: {count:,} vector with {dimension:,} dimension in vector store")
return vector_store
# Run Ingestion
if __name__ == "__main__":
documents = fetch_documents()
chunks = create_chunks(documents=documents)
create_embedding(chunks=chunks)
print(f"Ingestion Complete")
|