Spaces:
Sleeping
Sleeping
import os | |
import chainlit as cl | |
from dotenv import load_dotenv | |
from operator import itemgetter | |
from langchain_huggingface import HuggingFaceEndpoint | |
from langchain_community.document_loaders import TextLoader | |
from langchain_text_splitters import RecursiveCharacterTextSplitter | |
from langchain_community.vectorstores import FAISS | |
from langchain_huggingface import HuggingFaceEndpointEmbeddings | |
from langchain_core.prompts import PromptTemplate | |
from langchain.schema.output_parser import StrOutputParser | |
from langchain.schema.runnable import RunnablePassthrough | |
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_EMBED_ENDPOINT = os.environ["HF_EMBED_ENDPOINT"] | |
HF_TOKEN = os.environ["HF_TOKEN"] | |
# ---- GLOBAL DECLARATIONS ---- # | |
# -- RETRIEVAL -- # | |
""" | |
1. Load Documents from Text 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 | |
""" | |
### 1. CREATE TEXT LOADER AND LOAD DOCUMENTS | |
### NOTE: PAY ATTENTION TO THE PATH THEY ARE IN. | |
text_loader = | |
documents = | |
### 2. CREATE TEXT SPLITTER AND SPLIT DOCUMENTS | |
text_splitter = | |
split_documents = | |
### 3. LOAD HUGGINGFACE EMBEDDINGS | |
hf_embeddings = | |
if os.path.exists("./data/vectorstore"): | |
vectorstore = FAISS.load_local( | |
"./data/vectorstore", | |
hf_embeddings, | |
allow_dangerous_deserialization=True # this is necessary to load the vectorstore from disk as it's stored as a `.pkl` file. | |
) | |
hf_retriever = vectorstore.as_retriever() | |
print("Loaded Vectorstore") | |
else: | |
print("Indexing Files") | |
os.makedirs("./data/vectorstore", exist_ok=True) | |
### 4. INDEX FILES | |
### NOTE: REMEMBER TO BATCH THE DOCUMENTS WITH MAXIMUM BATCH SIZE = 32 | |
hf_retriever = vectorstore.as_retriever() | |
# -- AUGMENTED -- # | |
""" | |
1. Define a String Template | |
2. Create a Prompt Template from the String Template | |
""" | |
### 1. DEFINE STRING TEMPLATE | |
RAG_PROMPT_TEMPLATE = | |
### 2. CREATE PROMPT TEMPLATE | |
rag_prompt = | |
# -- GENERATION -- # | |
""" | |
1. Create a HuggingFaceEndpoint for the LLM | |
""" | |
### 1. CREATE HUGGINGFACE ENDPOINT FOR LLM | |
hf_llm = | |
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" : "Paul Graham Essay Bot" | |
} | |
return rename_dict.get(original_author, original_author) | |
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. | |
""" | |
### BUILD LCEL RAG CHAIN THAT ONLY RETURNS TEXT | |
lcel_rag_chain = | |
cl.user_session.set("lcel_rag_chain", lcel_rag_chain) | |
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="") | |
async for chunk in lcel_rag_chain.astream( | |
{"query": message.content}, | |
config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]), | |
): | |
await msg.stream_token(chunk) | |
await msg.send() |