Ganny/RAG-Optimization

1import os2import chainlit as cl3from dotenv import load_dotenv4from operator import itemgetter5from langchain_huggingface import HuggingFaceEndpoint6from langchain_community.document_loaders import TextLoader7from langchain_text_splitters import RecursiveCharacterTextSplitter8from langchain_community.vectorstores import FAISS9from langchain_huggingface import HuggingFaceEndpointEmbeddings10from langchain_core.prompts import PromptTemplate11from langchain.schema.output_parser import StrOutputParser12from langchain.schema.runnable import RunnablePassthrough13from langchain.schema.runnable.config import RunnableConfig14 15# GLOBAL SCOPE - ENTIRE APPLICATION HAS ACCESS TO VALUES SET IN THIS SCOPE #16# ---- ENV VARIABLES ---- # 17"""18This function will load our environment file (.env) if it is present.19 20NOTE: Make sure that .env is in your .gitignore file - it is by default, but please ensure it remains there.21"""22load_dotenv()23 24"""25We will load our environment variables here.26"""27HF_LLM_ENDPOINT = os.environ["HF_LLM_ENDPOINT"]28HF_EMBED_ENDPOINT = os.environ["HF_EMBED_ENDPOINT"]29HF_TOKEN = os.environ["HF_TOKEN"]30DATA_DIR = "./data"31VECTORSTORE_DIR = os.path.join(DATA_DIR, "vectorstore")32VECTORSTORE_PATH = os.path.join(VECTORSTORE_DIR, "index.faiss")33# ---- GLOBAL DECLARATIONS ---- #34 35# -- RETRIEVAL -- #36"""371. Load Documents from Text File382. Split Documents into Chunks393. Load HuggingFace Embeddings (remember to use the URL we set above)404. Index Files if they do not exist, otherwise load the vectorstore41"""42document_loader = TextLoader("./data/paul_graham_essays.txt")43documents = document_loader.load()44 45text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=30)46split_documents = text_splitter.split_documents(documents)47 48hf_embeddings = HuggingFaceEndpointEmbeddings(49    model=HF_EMBED_ENDPOINT,50    task="feature-extraction",51    huggingfacehub_api_token=HF_TOKEN,52)53 54if os.path.exists(VECTORSTORE_PATH):55    vectorstore = FAISS.load_local(56        VECTORSTORE_PATH, 57        hf_embeddings, 58        allow_dangerous_deserialization=True # this is necessary to load the vectorstore from disk as it's stored as a `.pkl` file.59    )60    hf_retriever = vectorstore.as_retriever()61    print("Loaded Vectorstore")62else:63    print("Indexing Files")64    os.makedirs(VECTORSTORE_PATH, exist_ok=True)65    for i in range(0, len(split_documents), 32):66        if i == 0:67            vectorstore = FAISS.from_documents(split_documents[i:i+32], hf_embeddings)68            continue69        vectorstore.add_documents(split_documents[i:i+32])70    vectorstore.save_local(VECTORSTORE_PATH)71 72hf_retriever = vectorstore.as_retriever()73 74# -- AUGMENTED -- #75"""761. Define a String Template772. Create a Prompt Template from the String Template78"""79RAG_PROMPT_TEMPLATE = """\80<|start_header_id|>system<|end_header_id|>81You 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|>82 83<|start_header_id|>user<|end_header_id|>84User Query:85{query}86 87Context:88{context}<|eot_id|>89 90<|start_header_id|>assistant<|end_header_id|>91"""92 93rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)94 95# -- GENERATION -- #96"""971. Create a HuggingFaceEndpoint for the LLM98"""99hf_llm = HuggingFaceEndpoint(100    endpoint_url=HF_LLM_ENDPOINT,101    max_new_tokens=512,102    top_k=10,103    top_p=0.95,104    temperature=0.3,105    repetition_penalty=1.15,106    huggingfacehub_api_token=HF_TOKEN,107)108 109@cl.author_rename110def rename(original_author: str):111    """112    This function can be used to rename the 'author' of a message. 113 114    In this case, we're overriding the 'Assistant' author to be 'Paul Graham Essay Bot'.115    """116    rename_dict = {117        "Assistant" : "Paul Graham Essay Bot"118    }119    return rename_dict.get(original_author, original_author)120 121@cl.on_chat_start122async def start_chat():123    """124    This function will be called at the start of every user session. 125 126    We will build our LCEL RAG chain here, and store it in the user session. 127 128    The user session is a dictionary that is unique to each user session, and is stored in the memory of the server.129    """130 131    lcel_rag_chain = ( {"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}132                      133                       | rag_prompt | hf_llm134                    )135 136    cl.user_session.set("lcel_rag_chain", lcel_rag_chain)137 138@cl.on_message  139async def main(message: cl.Message):140    """141    This function will be called every time a message is recieved from a session.142 143    We will use the LCEL RAG chain to generate a response to the user query.144 145    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.146    """147    lcel_rag_chain = cl.user_session.get("lcel_rag_chain")148 149    msg = cl.Message(content="")150 151    async for chunk in lcel_rag_chain.astream(152        {"query": message.content},153        config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),154    ):155        await msg.stream_token(chunk)156 157    await msg.send()158