CorrectiveRAGProject / ingestion.py
mikail-bitik's picture
embedding model was changed from Gemini to one of the Huggingface models
cefe798
Raw
History Blame Contribute Delete
1.27 kB
from dotenv import load_dotenv
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_huggingface import HuggingFaceEndpointEmbeddings
load_dotenv()
urls = [
"https://lilianweng.github.io/posts/2023-06-23-agent/",
"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
]
docs = [WebBaseLoader(urls).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter()
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
chunk_size=800, chunk_overlap=0
)
splits = text_splitter.split_documents(docs_list)
#embedding = OllamaEmbeddings(model="nomic-embed-text")
embedding = HuggingFaceEndpointEmbeddings(model="sentence-transformers/all-MiniLM-L6-v2")
vectorstore = Chroma.from_documents(
documents=splits,
collection_name="rag-chroma",
embedding=embedding,
persist_directory="./.chroma"
)
retriever = Chroma(
collection_name="rag-chroma",
persist_directory="./.chroma",
embedding_function=embedding,
).as_retriever()