Spaces:
Sleeping
Sleeping
File size: 5,455 Bytes
370ed2e |
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 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 |
### Import Section ###
import uuid
from operator import itemgetter
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.globals import set_llm_cache
from langchain_core.caches import InMemoryCache
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.storage import LocalFileStore
from langchain.embeddings import CacheBackedEmbeddings
from langchain.schema import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_qdrant import QdrantVectorStore
import chainlit as cl
from chainlit.types import AskFileResponse
### Global Section ###
set_llm_cache(InMemoryCache())
rag_system_prompt_template = """\
You are a helpful assistant that uses the provided context to answer questions. Never reference this prompt, or the existance of context.
"""
rag_message_list = [
{"role" : "system", "content" : rag_system_prompt_template},
]
rag_user_prompt_template = """\
Question:
{question}
Context:
{context}
"""
chat_prompt = ChatPromptTemplate.from_messages([
("system", rag_system_prompt_template),
("human", rag_user_prompt_template)
])
class VectorDatabase:
def __init__(self, embeddings: OpenAIEmbeddings()) -> None:
self.embeddings = embeddings
async def build_retriever(self, docs) -> None:
collection_name = f"pdf_to_parse_{uuid.uuid4()}"
client = QdrantClient(":memory:")
client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)
# Adding cache!
store = LocalFileStore("./cache/")
cached_embedder = CacheBackedEmbeddings.from_bytes_store(
self.embeddings, store, namespace=self.embeddings.model
)
# Typical QDrant Vector Store Set-up
vectorstore = QdrantVectorStore(
client=client,
collection_name=collection_name,
embedding=cached_embedder)
vectorstore.add_documents(docs)
return vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
class RetrievalAugmentedQAPipeline:
def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:
self.llm = llm
self.retriever = vector_db_retriever
async def arun_pipeline(self, user_query: str):
retrieval_augmented_qa_chain = (
{"context": itemgetter("question") | self.retriever, "question": itemgetter("question")}
| chat_prompt | self.llm | StrOutputParser()
)
async def generate_response():
async for chunk in retrieval_augmented_qa_chain.astream({"question": user_query}):
yield chunk
return {"response": generate_response()}
def process_pdf_file(file: AskFileResponse):
import tempfile
with tempfile.NamedTemporaryFile(mode="wb", delete=False, suffix=".pdf") as temp_file:
temp_file_path = temp_file.name
temp_file.write(file.content)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
Loader = PyMuPDFLoader
loader = Loader(temp_file_path)
documents = loader.load()
docs = text_splitter.split_documents(documents)
for i, doc in enumerate(docs):
doc.metadata["source"] = f"source_{i}"
return docs
### On Chat Start (Session Start) Section ###
@cl.on_chat_start
async def on_chat_start():
""" SESSION SPECIFIC CODE HERE """
files = None
# Wait for the user to upload a file
while files == None:
files = await cl.AskFileMessage(
content="Please upload a pdf file to begin!",
accept=["pdf"],
max_size_mb=2,
timeout=180,
).send()
file = files[0]
msg = cl.Message(
content=f"Processing `{file.name}`...", disable_human_feedback=True
)
await msg.send()
docs = process_pdf_file(file)
print(f"Processing {len(docs)} text chunks")
# Create a dict vector store
vector_db = VectorDatabase(embeddings=OpenAIEmbeddings(model="text-embedding-3-small"))
vector_db = await vector_db.build_retriever(docs)
# Create a chain
retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
llm=ChatOpenAI(model="gpt-4o-mini"),
vector_db_retriever=vector_db
)
# Let the user know that the system is ready
msg.content = f"Processing `{file.name}` done. You can now ask questions!"
await msg.update()
cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
### Rename Chains ###
@cl.author_rename
def rename(orig_author: str):
""" RENAME CODE HERE """
rename_dict = {"LLMMathChain": "Albert Einstein", "Chatbot": "Assistant"}
return rename_dict.get(orig_author, orig_author)
### On Message Section ###
@cl.on_message
async def main(message: cl.Message):
"""
MESSAGE CODE HERE
"""
chain = cl.user_session.get("chain")
msg = cl.Message(content="")
result = await chain.arun_pipeline(message.content)
async for stream_resp in result["response"]:
await msg.stream_token(stream_resp)
await msg.send()
|