Spaces:
Sleeping
Sleeping
### Import Section ### | |
""" | |
IMPORTS HERE | |
""" | |
import os | |
import uuid | |
from dotenv import load_dotenv | |
from langchain_text_splitters import RecursiveCharacterTextSplitter | |
from langchain_community.document_loaders import PyMuPDFLoader | |
from qdrant_client import QdrantClient | |
from qdrant_client.http.models import Distance, VectorParams | |
from langchain_openai.embeddings import OpenAIEmbeddings | |
from langchain.storage import LocalFileStore | |
from langchain_qdrant import QdrantVectorStore | |
from langchain.embeddings import CacheBackedEmbeddings | |
from langchain_core.prompts import ChatPromptTemplate | |
from chainlit.types import AskFileResponse | |
from langchain_core.globals import set_llm_cache | |
from langchain_openai import ChatOpenAI | |
from langchain_core.caches import InMemoryCache | |
from operator import itemgetter | |
from langchain_core.runnables.passthrough import RunnablePassthrough | |
import chainlit as cl | |
from langchain_core.runnables.config import RunnableConfig | |
from langchain_community.llms import HuggingFaceEndpoint | |
from langchain_huggingface.embeddings import HuggingFaceEndpointEmbeddings | |
from langchain_core.prompts import PromptTemplate | |
import numpy as np | |
from numpy.linalg import norm | |
load_dotenv() | |
### Global Section ### | |
""" | |
GLOBAL CODE HERE | |
""" | |
RAG_PROMPT_TEMPLATE = """\ | |
<|start_header_id|>system<|end_header_id|> | |
You are a helpful assistant. You answer user questions based on provided context. If you can't answer the question with the provided context, say you don't know.<|eot_id|> | |
<|start_header_id|>user<|end_header_id|> | |
User Query: | |
{query} | |
Context: | |
{context}<|eot_id|> | |
<|start_header_id|>assistant<|end_header_id|> | |
""" | |
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) | |
hf_llm = HuggingFaceEndpoint( | |
endpoint_url=f"{os.environ["YOUR_LLM_ENDPOINT_URL"]}", | |
max_new_tokens=512, | |
top_k=10, | |
top_p=0.95, | |
typical_p=0.95, | |
temperature=0.01, | |
repetition_penalty=1.03, | |
huggingfacehub_api_token=os.environ["HF_TOKEN"] | |
) | |
hf_embeddings = HuggingFaceEndpointEmbeddings( | |
model=os.environ["YOUR_EMBED_MODEL_URL"], | |
task="feature-extraction", | |
huggingfacehub_api_token=os.environ["HF_TOKEN"], | |
) | |
rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE) | |
rag_chain = rag_prompt | hf_llm | |
def cosine_similarity(phrase_1, phrase_2): | |
vec_1 = hf_embeddings.embed_documents([phrase_1])[0] | |
vec2_2 = hf_embeddings.embed_documents([phrase_2])[0] | |
return np.dot(vec_1, vec2_2) / (norm(vec_1) * norm(vec2_2)) | |
def process_file(file: AskFileResponse): | |
import tempfile | |
with tempfile.NamedTemporaryFile(mode="w", delete=False) as tempfile: | |
with open(tempfile.name, "wb") as f: | |
f.write(file.content) | |
Loader = PyMuPDFLoader | |
loader = Loader(tempfile.name) | |
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 ### | |
async def on_chat_start(): | |
""" SESSION SPECIFIC CODE HERE """ | |
files = None | |
while files == None: | |
# Async method: This allows the function to pause execution while waiting for the user to upload a file, | |
# without blocking the entire application. It improves responsiveness and scalability. | |
files = await cl.AskFileMessage( | |
content="Please upload a PDF file to begin!", | |
accept=["application/pdf"], | |
max_size_mb=20, | |
timeout=180, | |
max_files=1 | |
).send() | |
file = files[0] | |
msg = cl.Message( | |
content=f"Processing `{file.name}`...", | |
) | |
await msg.send() | |
docs = process_file(file) | |
# Typical QDrant Client Set-up | |
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( | |
hf_embeddings, store, namespace=hf_embeddings.model | |
) | |
# Typical QDrant Vector Store Set-up | |
vectorstore = QdrantVectorStore( | |
client=client, | |
collection_name=collection_name, | |
embedding=cached_embedder) | |
for i in range(0, len(docs), 32): | |
if i == 0: | |
vectorstore = docs.from_documents(docs[i:i+32], hf_embeddings) | |
continue | |
vectorstore.add_documents(docs[i:i+32]) | |
retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3}) | |
retrieval_augmented_qa_chain = ( | |
{"context": itemgetter("question") | retriever, "question": itemgetter("question")} | |
| RunnablePassthrough.assign(context=itemgetter("context")) | |
| rag_prompt | hf_llm | |
) | |
# 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_chain) | |
### Rename Chains ### | |
def rename(orig_author: str): | |
""" RENAME CODE HERE """ | |
rename_dict = {"ChatOpenAI": "the Generator...", "VectorStoreRetriever": "the Retriever..."} | |
return rename_dict.get(orig_author, orig_author) | |
### On Message Section ### | |
async def main(message: cl.Message): | |
""" | |
MESSAGE CODE HERE | |
""" | |
runnable = cl.user_session.get("chain") | |
msg = cl.Message(content="") | |
# Async method: Using astream allows for asynchronous streaming of the response, | |
# improving responsiveness and user experience by showing partial results as they become available. | |
async for chunk in runnable.astream( | |
{"question": message.content}, | |
config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]), | |
): | |
await msg.stream_token(chunk.content) | |
await msg.send() |