| 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() |
