basic_rag_dense_retrival.py

import os


os.environ["COHERE_API_KEY"] = "w7f2pzS4Ku5JSkVidvDeyDtZJWLlxZHwI5YmGZZP"


from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_cohere import CohereEmbeddings
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA

from langchain_cohere import ChatCohere

llm = ChatCohere(model="command-r", temperature=0)


embeddings = CohereEmbeddings(model="embed-english-v3.0")

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/",
    # bez tego llm nie radzi sobie z odpowiedzenim na ptyanie o autorow
    "https://www.sciencedirect.com/science/article/pii/S0378437119303528",
    # "https://milvus.io/ai-quick-reference/what-is-the-difference-between-sparse-and-dense-retrieval?utm_source=chatgpt.com"
]

# tutaj uzywamy dokumentow zaladowanych z neta dlatego uzywamy WebBaseLoader
docs = [WebBaseLoader(url).load() for url in urls]

# Tutaj tworzy sie jedna plaska lista, potencjalne miejsce gdzie mozna by strukturyzowac dane do grafu
docs_list = [item for sublist in docs for item in sublist]


text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=512,
    chunk_overlap=0,
    # chunk_size=50, chunk_overlap=0
)

chunks = text_splitter.split_documents(docs_list)

# here we add to vector_store
chroma_vectorstore = Chroma.from_documents(documents=chunks, embedding=embeddings)

retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 4})

# query = "What are authors of: 'A review of cellular automata models for crowd evacuation' paper?"
# query = "What are authors of A review of cellular automata evacuation paper?"
# query = "How do simulation methods replicate panic-driven group behavior in confined areas?"
# query = "Which computational framework divides the environment into discrete cells to simulate evacuation patterns?"
# query = "What kind of spatial logic is used to simulate individual flow in emergency exits?"
query = "Who wrote the evacuation review?"

retrieved = retriever.get_relevant_documents(query)

qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, chain_type="stuff")

response = qa_chain.invoke(query)


print("Answer: ", response)