File size: 1,467 Bytes
9718baa | 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 | import os
from dotenv import load_dotenv
from langchain_community.document_loaders import WebBaseLoader, TextLoader
from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
load_dotenv()
TEXT_FILE_PATH = "./data/aboutPromtior.txt"
WEBSITE_URL = os.getenv("WEBSITE_URL")
VECTOR_STORE_PATH = "./vectorstore"
USER_AGENT = os.getenv("USER_AGENT")
def main():
generate_data_store()
def generate_data_store():
documents = load_documents()
chunks = split_text(documents)
save_to_chroma(chunks)
def load_documents():
doc_text = load_text_content()
doc_web = load_web_content(WEBSITE_URL)
return doc_text + doc_web
def load_text_content():
loader = TextLoader(TEXT_FILE_PATH)
documents = loader.load()
return documents
def load_web_content(url):
loader = WebBaseLoader(url)
documents = loader.load()
return documents
def split_text(documents):
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, add_start_index=True)
chunks = text_splitter.split_documents(documents)
return chunks
def save_to_chroma(chunks):
embeddings = OllamaEmbeddings(model="llama3")
Chroma.from_documents(
chunks,
embeddings,
persist_directory=VECTOR_STORE_PATH,
)
print(f"Saved {len(chunks)} chunks to {VECTOR_STORE_PATH}")
if __name__ == "__main__":
main() |