niting089 commited on
Commit
9b2a5f8
1 Parent(s): f62afeb

added solution_app.py file

Browse files
Files changed (2) hide show
  1. .gitignore +1 -2
  2. solution_app.py +190 -0
.gitignore CHANGED
@@ -1,2 +1 @@
1
- .env
2
- secrets.txt
 
1
+ .env
 
solution_app.py ADDED
@@ -0,0 +1,190 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import chainlit as cl
3
+ from dotenv import load_dotenv
4
+ from operator import itemgetter
5
+ from langchain_huggingface import HuggingFaceEndpoint
6
+ from langchain_community.document_loaders import TextLoader
7
+ from langchain_text_splitters import RecursiveCharacterTextSplitter
8
+ from langchain_community.vectorstores import FAISS
9
+ from langchain_huggingface import HuggingFaceEndpointEmbeddings
10
+ from langchain_core.prompts import PromptTemplate
11
+ from langchain.schema.output_parser import StrOutputParser
12
+ from langchain.schema.runnable import RunnablePassthrough
13
+ from langchain.schema.runnable.config import RunnableConfig
14
+ from tqdm.asyncio import tqdm_asyncio
15
+ import asyncio
16
+ from tqdm.asyncio import tqdm
17
+
18
+ # GLOBAL SCOPE - ENTIRE APPLICATION HAS ACCESS TO VALUES SET IN THIS SCOPE #
19
+ # ---- ENV VARIABLES ---- #
20
+ """
21
+ This function will load our environment file (.env) if it is present.
22
+
23
+ NOTE: Make sure that .env is in your .gitignore file - it is by default, but please ensure it remains there.
24
+ """
25
+ load_dotenv()
26
+
27
+ """
28
+ We will load our environment variables here.
29
+ """
30
+ HF_LLM_ENDPOINT = os.environ["https://v0jxvdenly3jgdw7.us-east-1.aws.endpoints.huggingface.cloud"]
31
+ HF_EMBED_ENDPOINT = os.environ["https://at1cd0u50368nxla.us-east-1.aws.endpoints.huggingface.cloud"]
32
+ HF_TOKEN = os.getenv("HF_TOKEN")
33
+
34
+ # ---- GLOBAL DECLARATIONS ---- #
35
+
36
+ # -- RETRIEVAL -- #
37
+ """
38
+ 1. Load Documents from Text File
39
+ 2. Split Documents into Chunks
40
+ 3. Load HuggingFace Embeddings (remember to use the URL we set above)
41
+ 4. Index Files if they do not exist, otherwise load the vectorstore
42
+ """
43
+ document_loader = TextLoader("./data/paul_graham_essays.txt")
44
+ documents = document_loader.load()
45
+
46
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=30)
47
+ split_documents = text_splitter.split_documents(documents)
48
+
49
+ hf_embeddings = HuggingFaceEndpointEmbeddings(
50
+ model=HF_EMBED_ENDPOINT,
51
+ task="feature-extraction",
52
+ huggingfacehub_api_token=HF_TOKEN,
53
+ )
54
+
55
+ async def add_documents_async(vectorstore, documents):
56
+ await vectorstore.aadd_documents(documents)
57
+
58
+ async def process_batch(vectorstore, batch, is_first_batch, pbar):
59
+ if is_first_batch:
60
+ result = await FAISS.afrom_documents(batch, hf_embeddings)
61
+ else:
62
+ await add_documents_async(vectorstore, batch)
63
+ result = vectorstore
64
+ pbar.update(len(batch))
65
+ return result
66
+
67
+ async def main():
68
+ print("Indexing Files")
69
+
70
+ vectorstore = None
71
+ batch_size = 32
72
+
73
+ batches = [split_documents[i:i+batch_size] for i in range(0, len(split_documents), batch_size)]
74
+
75
+ async def process_all_batches():
76
+ nonlocal vectorstore
77
+ tasks = []
78
+ pbars = []
79
+
80
+ for i, batch in enumerate(batches):
81
+ pbar = tqdm(total=len(batch), desc=f"Batch {i+1}/{len(batches)}", position=i)
82
+ pbars.append(pbar)
83
+
84
+ if i == 0:
85
+ vectorstore = await process_batch(None, batch, True, pbar)
86
+ else:
87
+ tasks.append(process_batch(vectorstore, batch, False, pbar))
88
+
89
+ if tasks:
90
+ await asyncio.gather(*tasks)
91
+
92
+ for pbar in pbars:
93
+ pbar.close()
94
+
95
+ await process_all_batches()
96
+
97
+ hf_retriever = vectorstore.as_retriever()
98
+ print("\nIndexing complete. Vectorstore is ready for use.")
99
+ return hf_retriever
100
+
101
+ async def run():
102
+ retriever = await main()
103
+ return retriever
104
+
105
+ hf_retriever = asyncio.run(run())
106
+
107
+ # -- AUGMENTED -- #
108
+ """
109
+ 1. Define a String Template
110
+ 2. Create a Prompt Template from the String Template
111
+ """
112
+ RAG_PROMPT_TEMPLATE = """\
113
+ <|start_header_id|>system<|end_header_id|>
114
+ 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|>
115
+
116
+ <|start_header_id|>user<|end_header_id|>
117
+ User Query:
118
+ {query}
119
+
120
+ Context:
121
+ {context}<|eot_id|>
122
+
123
+ <|start_header_id|>assistant<|end_header_id|>
124
+ """
125
+
126
+ rag_prompt = PromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
127
+
128
+ # -- GENERATION -- #
129
+ """
130
+ 1. Create a HuggingFaceEndpoint for the LLM
131
+ """
132
+ hf_llm = HuggingFaceEndpoint(
133
+ endpoint_url=HF_LLM_ENDPOINT,
134
+ max_new_tokens=512,
135
+ top_k=10,
136
+ top_p=0.95,
137
+ temperature=0.3,
138
+ repetition_penalty=1.15,
139
+ huggingfacehub_api_token=HF_TOKEN,
140
+ )
141
+
142
+ @cl.author_rename
143
+ def rename(original_author: str):
144
+ """
145
+ This function can be used to rename the 'author' of a message.
146
+
147
+ In this case, we're overriding the 'Assistant' author to be 'Paul Graham Essay Bot'.
148
+ """
149
+ rename_dict = {
150
+ "Assistant" : "Paul Graham Essay Bot"
151
+ }
152
+ return rename_dict.get(original_author, original_author)
153
+
154
+ @cl.on_chat_start
155
+ async def start_chat():
156
+ """
157
+ This function will be called at the start of every user session.
158
+
159
+ We will build our LCEL RAG chain here, and store it in the user session.
160
+
161
+ The user session is a dictionary that is unique to each user session, and is stored in the memory of the server.
162
+ """
163
+
164
+ lcel_rag_chain = (
165
+ {"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}
166
+ | rag_prompt | hf_llm
167
+ )
168
+
169
+ cl.user_session.set("lcel_rag_chain", lcel_rag_chain)
170
+
171
+ @cl.on_message
172
+ async def main(message: cl.Message):
173
+ """
174
+ This function will be called every time a message is recieved from a session.
175
+
176
+ We will use the LCEL RAG chain to generate a response to the user query.
177
+
178
+ 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.
179
+ """
180
+ lcel_rag_chain = cl.user_session.get("lcel_rag_chain")
181
+
182
+ msg = cl.Message(content="")
183
+
184
+ for chunk in await cl.make_async(lcel_rag_chain.stream)(
185
+ {"query": message.content},
186
+ config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
187
+ ):
188
+ await msg.stream_token(chunk)
189
+
190
+ await msg.send()