Spaces:
Sleeping
Sleeping
File size: 5,623 Bytes
54eae24 |
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 |
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
from tqdm.asyncio import tqdm_asyncio
import asyncio
from tqdm.asyncio import tqdm
load_dotenv()
HF_LLM_ENDPOINT = os.environ["HF_LLM_ENDPOINT"]
HF_EMBED_ENDPOINT = os.environ["HF_EMBED_ENDPOINT"]
HF_TOKEN = os.environ["HF_TOKEN"]
if not all([HF_LLM_ENDPOINT, HF_EMBED_ENDPOINT, HF_TOKEN]):
raise ValueError("Missing required Hugging Face API environment variables.")
if not os.path.exists("./data/paul_graham_essays.txt"):
raise FileNotFoundError("The specified document file is missing.")
document_loader = TextLoader("./data/paul_graham_essays.txt")
documents = document_loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=30)
split_documents = text_splitter.split_documents(documents)
hf_embeddings = HuggingFaceEndpointEmbeddings(
model=HF_EMBED_ENDPOINT,
task="feature-extraction",
huggingfacehub_api_token=HF_TOKEN,
)
async def add_documents_async(vectorstore, documents):
await vectorstore.aadd_documents(documents)
async def process_batch(vectorstore, batch, is_first_batch, pbar):
if is_first_batch:
result = await FAISS.afrom_documents(batch, hf_embeddings)
else:
await add_documents_async(vectorstore, batch)
result = vectorstore
pbar.update(len(batch))
return result
async def main():
print("Indexing Files")
vectorstore = None
batch_size = 32
batches = [split_documents[i:i+batch_size] for i in range(0, len(split_documents), batch_size)]
async def process_all_batches():
nonlocal vectorstore
tasks = []
pbars = []
for i, batch in enumerate(batches):
pbar = tqdm(total=len(batch), desc=f"Batch {i+1}/{len(batches)}", position=i)
pbars.append(pbar)
if i == 0:
vectorstore = await process_batch(None, batch, True, pbar)
else:
tasks.append(process_batch(vectorstore, batch, False, pbar))
if tasks:
await asyncio.gather(*tasks)
for pbar in pbars:
pbar.close()
await process_all_batches()
hf_retriever = vectorstore.as_retriever()
print("\nIndexing complete. Vectorstore is ready for use.")
return hf_retriever
async def run():
retriever = await main()
return retriever
hf_retriever = asyncio.run(run())
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)
### 1. CREATE HUGGINGFACE ENDPOINT FOR 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.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.
"""
### BUILD LCEL RAG CHAIN THAT ONLY RETURNS TEXT
lcel_rag_chain = (
{"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}
| rag_prompt | hf_llm
)
cl.user_session.set("lcel_rag_chain", lcel_rag_chain)
await cl.Message(content="Welcome to the Paul Graham Essay Bot! Ask me anything about Paul Graham's essays.").send()
@cl.on_message
async def main(message: cl.Message):
"""
This function will be called every time a message is received 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="")
# Create a buffer to store the complete response
complete_response = ""
async for chunk in lcel_rag_chain.astream(
{"query": message.content},
config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
):
# Clean the chunk of any <|eot_id|> tokens
clean_chunk = chunk.replace("<|eot_id|>", "")
complete_response += chunk # Store original for final cleaning
await msg.stream_token(clean_chunk)
# If <|eot_id|> appears at the end of the complete response, remove it
if complete_response.endswith("<|eot_id|>"):
# Update the final message content without the token
await msg.update(content=complete_response.replace("<|eot_id|>", ""))
await msg.send() |