Spaces:
Sleeping
Sleeping
File size: 2,879 Bytes
dc89a4d 80e747e dc89a4d 80e747e dc89a4d 80e747e dc89a4d |
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 93 94 |
import chainlit as cl
from operator import itemgetter
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.vectorstores import Qdrant
# Constants (you can adjust these as per your environment)
# DATA LOADER
DATA_LINK1 = "https://www.whitehouse.gov/wp-content/uploads/2022/10/Blueprint-for-an-AI-Bill-of-Rights.pdf"
DATA_LINK2 = "https://nvlpubs.nist.gov/nistpubs/ai/NIST.AI.600-1.pdf"
# CHUNKING CONFIGS
CHUNK_SIZE = 500
CHUNK_OVERLAP = 50
# RETRIEVER CONFIGS
COLLECTION_NAME = "AI Bill of Rights"
EMBEDDING_MODEL = "text-embedding-3-small"
# FINAL RAG CONFIGS
QA_MODEL = "gpt-4o"
RAG_PROMPT = """\
Given a provided context and question, you must answer the question based only on context.
If you cannot answer the question based on the context - you must say "I don't know".
Context: {context}
Question: {question}
"""
# Function to chunk documents
def chunk_documents(unchunked_documents, chunk_size, chunk_overlap):
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
)
return text_splitter.split_documents(unchunked_documents)
# Function to build retriever
def build_retriever(chunked_documents, embeddings, collection_name):
vectorstore = Qdrant.from_documents(
documents=chunked_documents,
embedding=embeddings,
location=":memory:", # Storing in-memory for demonstration
collection_name=collection_name,
)
retriever = vectorstore.as_retriever()
return retriever
# Load documents and prepare retriever
rag_documents_1 = PyMuPDFLoader(file_path=DATA_LINK1).load()
rag_documents_2 = PyMuPDFLoader(file_path=DATA_LINK2).load()
chunked_rag_documents = chunk_documents(rag_documents_1, CHUNK_SIZE, CHUNK_OVERLAP) + \
chunk_documents(rag_documents_2, CHUNK_SIZE, CHUNK_OVERLAP)
@cl.on_chat_start
async def on_chat_start():
embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)
retriever = build_retriever(chunked_rag_documents, embeddings, COLLECTION_NAME)
rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
qa_llm = ChatOpenAI(model=QA_MODEL)
rag_chain = (
{"context": itemgetter("question") | retriever, "question": itemgetter("question")}
| rag_prompt | qa_llm | StrOutputParser()
)
cl.user_session.set("chain", rag_chain)
# Chainlit app
@cl.on_message
async def main(message):
chain = cl.user_session.get("chain")
result = chain.invoke({"question" : message.content})
await cl.Message(
content=result, # Extract the response from the chain
author="AI"
).send()
|