Spaces:
Paused
Paused
File size: 4,678 Bytes
db9f86a 8b7bc13 db9f86a 8b7bc13 db9f86a 8b7bc13 db9f86a 8b7bc13 db9f86a 8b7bc13 db9f86a |
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 |
import os
import chainlit as cl
from dotenv import load_dotenv
from operator import itemgetter
from langchain_huggingface import HuggingFaceEndpoint
from PyPDF2 import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Qdrant
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain_core.prompts import PromptTemplate
from langchain.schema.runnable.config import RunnableConfig
# GLOBAL SCOPE - ENTIRE APPLICATION HAS ACCESS TO VALUES SET IN THIS SCOPE #
# ---- ENV VARIABLES ---- #
"""
This function will load our environment file (.env) if it is present.
NOTE: Make sure that .env is in your .gitignore file - it is by default, but please ensure it remains there.
"""
load_dotenv()
"""
We will load our environment variables here.
"""
HF_LLM_ENDPOINT = os.environ["HF_LLM_ENDPOINT"]
HF_TOKEN = os.environ["HF_TOKEN"]
VECTOR_STORE_PATH = "./data/vectorstore"
# ---- GLOBAL DECLARATIONS ---- #
# -- RETRIEVAL -- #
"""
1. Load Documents from PDF File
2. Split Documents into Chunks
3. Load HuggingFace Embeddings (remember to use the URL we set above)
4. Index Files if they do not exist, otherwise load the vectorstore
"""
document_loader = PdfReader("./data/Airbnb_Q1_Filings.pdf")
raw_text = ''
for i, page in enumerate(document_loader.pages):
content = page.extract_text()
if content:
raw_text += content
text_splitter = RecursiveCharacterTextSplitter(
chunk_size = 850,
chunk_overlap = 50,
length_function = len,
)
texts = text_splitter.split_text(raw_text)
# Note: Uses OPENAI_API_KEY env variable to make api calls
hf_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
for i in range(0, len(texts), 32):
if i == 0:
vectorstore = Qdrant.from_texts(
texts[i:i+32],
hf_embeddings,
force_recreate=True,
path=VECTOR_STORE_PATH,
collection_name="Airbnb Filings")
continue
vectorstore.add_texts(texts[i:i+32])
retriever = vectorstore.as_retriever()
# -- AUGMENTED -- #
"""
1. Define a String Template
2. Create a Prompt Template from the String Template
"""
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|>
"""
rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
# -- GENERATION -- #
"""
1. Create a HuggingFaceEndpoint for the LLM
"""
hf_llm = HuggingFaceEndpoint(
endpoint_url=HF_LLM_ENDPOINT,
max_new_tokens=512,
top_k=10,
top_p=0.95,
temperature=0.3,
repetition_penalty=1.15,
huggingfacehub_api_token=HF_TOKEN,
)
@cl.author_rename
def rename(original_author: str):
"""
This function can be used to rename the 'author' of a message.
In this case, we're overriding the 'Assistant' author to be 'Paul Graham Essay Bot'.
"""
rename_dict = {
"Assistant" : "AirBnB Auditor"
}
return rename_dict.get(original_author, original_author)
@cl.on_chat_start
async def start_chat():
"""
This function will be called at the start of every user session.
We will build our LCEL RAG chain here, and store it in the user session.
The user session is a dictionary that is unique to each user session, and is stored in the memory of the server.
"""
lcel_rag_chain = (
{"context": itemgetter("query") | retriever, "query": itemgetter("query")}
| rag_prompt | hf_llm
)
cl.user_session.set("lcel_rag_chain", lcel_rag_chain)
@cl.on_message
async def main(message: cl.Message):
"""
This function will be called every time a message is recieved from a session.
We will use the LCEL RAG chain to generate a response to the user query.
The LCEL RAG chain is stored in the user session, and is unique to each user session - this is why we can access it here.
"""
lcel_rag_chain = cl.user_session.get("lcel_rag_chain")
msg = cl.Message(content="")
for chunk in await cl.make_async(lcel_rag_chain.stream)(
{"query": message.content},
config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
):
# Note: Skip printing eot_id token at the end of response. A more elegant solution would be to fix the model's generator config.
if chunk != "<|eot_id|>":
await msg.stream_token(chunk)
await msg.send()
|