karthikrathod commited on
Commit
aa133c0
·
verified ·
1 Parent(s): 0f29960

Update Astronomy_BH_hybrid_RAG.py

Browse files
Files changed (1) hide show
  1. Astronomy_BH_hybrid_RAG.py +320 -320
Astronomy_BH_hybrid_RAG.py CHANGED
@@ -1,321 +1,321 @@
1
- #############################Imports##############################################
2
- ## General imports
3
- import os
4
- import sys
5
- import zipfile
6
- import logging
7
- import IPython
8
- from IPython.display import display
9
- from pyvis.network import Network
10
- import dotenv
11
- from google.cloud import storage
12
-
13
- dotenv.load_dotenv()
14
-
15
- os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "tidy-resolver-411707-0f032726c297.json"
16
-
17
-
18
- def download_blob(bucket_name, source_blob_name, destination_file_name):
19
- """Downloads a blob from the bucket."""
20
- storage_client = storage.Client()
21
- bucket = storage_client.bucket(bucket_name)
22
- blob = bucket.blob(source_blob_name)
23
- blob.download_to_filename(destination_file_name)
24
- print(f"Downloaded storage object {source_blob_name} from bucket {bucket_name} to local file {destination_file_name}.")
25
-
26
- if not (os.path.exists("storage_file/bm25") and os.path.exists("storage_file/kg")):
27
- # List of file names to download
28
- file_names = [
29
- "default__vector_store.json",
30
- "docstore.json",
31
- "graph_store.json",
32
- "image__vector_store.json",
33
- "index_store.json"
34
- ]
35
-
36
- # Bucket name
37
- bucket_name = "title_tailors_bucket"
38
-
39
- # Create the destination directory if it doesn't exist
40
- os.makedirs("storage_file/bm25", exist_ok=True)
41
-
42
- # Loop through the file names and download each one
43
- for file_name in file_names:
44
- source_blob_name = f"storage/bm25/{file_name}"
45
- destination_file_name = f"storage_file/bm25/{file_name}"
46
- download_blob(bucket_name, source_blob_name, destination_file_name)
47
-
48
- # List of file names to download
49
- file_names = [
50
- "default__vector_store.json",
51
- "docstore.json",
52
- "graph_store.json",
53
- "image__vector_store.json",
54
- "index_store.json"
55
- ]
56
-
57
- # Bucket name
58
- bucket_name = "title_tailors_bucket"
59
-
60
- # Create the destination directory if it doesn't exist
61
- os.makedirs("storage_file/kg", exist_ok=True)
62
-
63
- # Loop through the file names and download each one
64
- for file_name in file_names:
65
- source_blob_name = f"storage/kg/{file_name}"
66
- destination_file_name = f"storage_file/kg/{file_name}"
67
- download_blob(bucket_name, source_blob_name, destination_file_name)
68
- else:
69
- print("Files already exist in the storage_file directory.")
70
-
71
-
72
- HF_TOKEN = os.environ.get("HF_TOKEN", None)
73
-
74
- MISTRAL_API = os.environ.get("MISTRAL_API", None)
75
- print(f"MISTRAL_API: {MISTRAL_API}")
76
-
77
- ## logger
78
- logging.basicConfig(stream=sys.stdout, level=logging.INFO)
79
- logging.getLogger().handlers = []
80
- logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
81
-
82
- # Knowledge graph imports
83
- from llama_index.core import (
84
- SimpleDirectoryReader,
85
- StorageContext,
86
- KnowledgeGraphIndex,
87
- load_index_from_storage,
88
- Settings,
89
- )
90
- from llama_index.core.graph_stores import SimpleGraphStore
91
- from llama_index.embeddings.mistralai import MistralAIEmbedding
92
- from llama_index.llms.mistralai import MistralAI
93
-
94
- ## BM25 imports
95
- from llama_index.core import (
96
- VectorStoreIndex,
97
- QueryBundle,
98
- )
99
- from llama_index.retrievers.bm25 import BM25Retriever
100
- from llama_index.core.node_parser import SentenceSplitter
101
- from llama_index.core.retrievers import (
102
- BaseRetriever,
103
- VectorIndexRetriever,
104
- QueryFusionRetriever,
105
- )
106
- from llama_index.core.schema import NodeWithScore
107
- from llama_index.core.query_engine import RetrieverQueryEngine
108
- from llama_index.core.postprocessor import SentenceTransformerRerank
109
- from llama_index.core.response.notebook_utils import (
110
- display_response,
111
- display_source_node,
112
- )
113
-
114
- # Chat engine
115
- from llama_index.core import PromptTemplate
116
- from llama_index.core import chat_engine
117
- from llama_index.core import memory
118
-
119
- import nest_asyncio
120
- nest_asyncio.apply()
121
-
122
- ################### Loading the LLM via Mistral API
123
- llm = MistralAI(api_key=MISTRAL_API, model="open-mixtral-8x7b")
124
- ### Loading the Embedding via Mistral API
125
- embed_model = MistralAIEmbedding(api_key=MISTRAL_API, model = "mistral-embed")
126
-
127
- ################### Knowledge Graph Index#################
128
- ########### While loading from persist only ###############
129
- Settings.llm = llm
130
- Settings.embed_model = embed_model
131
- kg_storage_context = StorageContext.from_defaults(persist_dir="storage/kg")
132
- kg_index = load_index_from_storage(kg_storage_context)
133
-
134
- kg_retriever = kg_index.as_retriever(include_text=True,
135
- response_mode ="tree_summarize",
136
- embedding_mode="hybrid",
137
- similarity_top_k=10)
138
-
139
- ################### BM25 Index #################
140
- ####################### While loading Indices from persist #######################
141
- # Settings.llm = llm
142
- # Settings.embed_model = embed_model
143
- storage_context_v = StorageContext.from_defaults(persist_dir="storage/bm25")
144
- index_v = load_index_from_storage(storage_context_v)
145
-
146
- vector_retriever = index_v.as_retriever(similarity_top_k=10)
147
- bm25_retriever = BM25Retriever.from_defaults(index=index_v, similarity_top_k=10, verbose=False) ## loading after persist
148
-
149
-
150
- ###################### Reranker from HuggingFace 🤗
151
- reranker = SentenceTransformerRerank(top_n=5, model="mixedbread-ai/mxbai-rerank-base-v1", keep_retrieval_score=False)
152
-
153
- ################ Query Fusion Retriever ##################
154
- QUERY_GEN_PROMPT = (
155
- "You are a helpful assistant that generates multiple search queries on astronomy based on a "
156
- "single input query. Generate {num_queries} search queries for astronomy, one on each line, "
157
- "related to the following input query:\n"
158
- "Query: {query}\n"
159
- "Queries:\n"
160
- )
161
-
162
- #### Hybrid (Dense + BM25 + KG) Retriever ####
163
- hybrid_retriever = QueryFusionRetriever(
164
- retrievers = [vector_retriever, bm25_retriever, kg_retriever],
165
- retriever_weights = [0.25, 0.25, 0.50],
166
- similarity_top_k=5,
167
- llm=llm,
168
- num_queries=3, # set this to 1 to disable query generation
169
- mode="reciprocal_rerank",
170
- use_async=True,
171
- verbose=False,
172
- query_gen_prompt=QUERY_GEN_PROMPT,
173
- )
174
-
175
- #### KG Retriever ####
176
- kg_retriever = QueryFusionRetriever(
177
- retrievers = [kg_retriever],
178
- # retriever_weights = [0.25, 0.25, 0.50],
179
- similarity_top_k=5,
180
- llm=llm,
181
- num_queries=3, # set this to 1 to disable query generation
182
- mode="reciprocal_rerank",
183
- use_async=True,
184
- verbose=False,
185
- query_gen_prompt=QUERY_GEN_PROMPT,
186
- )
187
-
188
- #### Hybrid BM25 (Dense + BM25) Retriever ####
189
- hybrid_bm25_retriever = QueryFusionRetriever(
190
- retrievers = [vector_retriever, bm25_retriever],
191
- # retriever_weights = [0.25, 0.25, 0.50],
192
- similarity_top_k=5,
193
- llm=llm,
194
- num_queries=3, # set this to 1 to disable query generation
195
- mode="reciprocal_rerank",
196
- use_async=True,
197
- verbose=False,
198
- query_gen_prompt=QUERY_GEN_PROMPT,
199
- )
200
-
201
- ################# Hybrid Chat Engine ###################
202
- system_prompt_template = """You are a helpful AI assistant for reaearcher or enthusiast in the domain of astronomy.
203
- Please check if the following pieces of context has any mention of the keywords provided in the Question. If not then don't know the answer, just say that you don't know.Stop there. Please donot try to make up an answer."""
204
-
205
- from llama_index.core import chat_engine
206
- from llama_index.core import memory
207
-
208
- ######## Hybrid (Dense + BM25 + KG) chat engine #########
209
- hybrid_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
210
- Hybrid_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=hybrid_retriever,
211
- llm=llm,
212
- memory=hybrid_memory,
213
- # context_prompt=,
214
- # condense_prompt=,
215
- system_prompt=system_prompt_template,
216
- node_postprocessors=[reranker],
217
- verbose=False,
218
- )
219
-
220
- ######## Hybrid (Dense + BM25 + KG) chat engine #########
221
- kg_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
222
- kg_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=kg_retriever,
223
- llm=llm,
224
- memory=kg_memory,
225
- # context_prompt=,
226
- # condense_prompt=,
227
- system_prompt=system_prompt_template,
228
- node_postprocessors=[reranker],
229
- verbose=False,
230
- )
231
-
232
- ######## Hybrid (Dense + BM25 + KG) chat engine #########
233
- hybrid_bm25_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
234
- hybrid_bm25_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=hybrid_bm25_retriever,
235
- llm=llm,
236
- memory=hybrid_bm25_memory,
237
- # context_prompt=,
238
- # condense_prompt=,
239
- system_prompt=system_prompt_template,
240
- node_postprocessors=[reranker],
241
- verbose=False,
242
- )
243
-
244
- ################## Post Processors ###################
245
- def get_documents(response):
246
- """Get the reference documents for a chat engine response.
247
-
248
- Args:
249
- response (ChatResponse): The chat engine response object.
250
-
251
- Returns:
252
- str: A formatted string containing the reference document paths and page numbers.
253
- """
254
- plist = []
255
- reference_docs={}
256
- for idx, content in enumerate(response.sources[0].content.split("\n\n")):
257
- if "page_label" in content:
258
- plist.append(content)
259
- for i, entry in enumerate(plist, start=1):
260
- lines = entry.split('\n')
261
- page_label = int(lines[0].split(': ')[1])
262
- file_path = lines[1].split(': ')[1]
263
-
264
- reference_docs[f"doc{i}"] = {
265
- "document_path": file_path,
266
- "page": page_label
267
- }
268
- reference = f"""
269
- \n Reference Docs (document paths, page number):
270
- ...{reference_docs['doc1']['document_path']}, page {reference_docs['doc1']['page']}
271
- ...{reference_docs['doc2']['document_path']}, page {reference_docs['doc2']['page']}
272
- ...{reference_docs['doc3']['document_path']}, page {reference_docs['doc3']['page']}
273
- ...{reference_docs['doc4']['document_path']}, page {reference_docs['doc4']['page']}
274
- ...{reference_docs['doc5']['document_path']}, page {reference_docs['doc5']['page']}
275
- """
276
- return reference
277
-
278
- ##################### Query Function #####################
279
- def get_query(query=""):
280
- """Get a response from the Hybrid chat engine and format the result with reference documents.
281
-
282
- Args:
283
- query (str): The query to send to the chat engine.
284
-
285
- Returns:
286
- str: The chat engine response with the reference documents appended.
287
- """
288
- response_1 = Hybrid_chat_engine.chat(query)
289
- response_2 = kg_chat_engine.chat(query)
290
- response_3 = hybrid_bm25_chat_engine.chat(query)
291
- reference_1 = get_documents(response_1)
292
- reference_2 = get_documents(response_2)
293
- reference_3 = get_documents(response_3)
294
- reply_1 = response_1.response
295
- reply_2 = response_2.response
296
- reply_3 = response_3.response
297
- result_hybrid = reply_1 + reference_1
298
- result_kg = reply_2 + reference_2
299
- result_hybrid_bm25 = reply_3 + reference_3
300
- return result_hybrid, result_kg, result_hybrid_bm25
301
-
302
-
303
-
304
- # print(get_query("what is difference between class I and class II ?"))
305
- # Based on the provided documents, Class I and Class II are two distinct evolutionary stages of young stellar objects (YSOs) in the process of forming stars. The documents describe the evolutionary stages of YSOs as follows:
306
-
307
- # * Class 0: The formation of a YSO in the central region of a protostellar core with an envelope mass that is much in excess of the YSO mass.
308
- # * Class I: The collapse of the envelope onto the central object, with the transition between Class 0 and Class I being the point in time at which the envelope mass and the mass of the protostar are nearly equal.
309
- # * Class II: The emergence of a disk around the central star.
310
- # * Class III: The dissipation of the disk by various processes such as the formation of planets, photo-evaporation, and tidal stripping.
311
-
312
- # The documents also mention that an intermediate class between Class 0 and Class I has been proposed, but it is considered to be close to Class I in terms of the evolutionary status of the YSOs.
313
-
314
- # Regarding the difference between Class I and Class II, the documents do not provide a detailed explanation, but it is suggested that the main difference lies in the presence of a disk around the central star in Class II, which is not present in Class I. The emergence of a disk around the central star is a sign of a more evolved stage in the star formation process. The documents also mention that the exact duration of each evolutionary stage for YSOs is relatively uncertain and depends on the number of objects found in each class, which may be affected by misclassifications due to YSOs being seen edge-on.
315
- # Reference Docs (document paths, page number):
316
- # .../content/documents/2405.00095v1.pdf, page 8
317
- # .../content/documents/2405.00095v1.pdf, page 3
318
- # .../content/documents/2405.00095v1.pdf, page 9
319
- # .../content/documents/2405.00095v1.pdf, page 2
320
- # .../content/documents/2405.00095v1.pdf, page 2
321
 
 
1
+ #############################Imports##############################################
2
+ ## General imports
3
+ import os
4
+ import sys
5
+ import zipfile
6
+ import logging
7
+ import IPython
8
+ from IPython.display import display
9
+ from pyvis.network import Network
10
+ import dotenv
11
+ from google.cloud import storage
12
+
13
+ dotenv.load_dotenv()
14
+
15
+ os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "tidy-resolver-411707-0f032726c297.json"
16
+
17
+
18
+ def download_blob(bucket_name, source_blob_name, destination_file_name):
19
+ """Downloads a blob from the bucket."""
20
+ storage_client = storage.Client()
21
+ bucket = storage_client.bucket(bucket_name)
22
+ blob = bucket.blob(source_blob_name)
23
+ blob.download_to_filename(destination_file_name)
24
+ print(f"Downloaded storage object {source_blob_name} from bucket {bucket_name} to local file {destination_file_name}.")
25
+
26
+ if not (os.path.exists("storage/bm25") and os.path.exists("storag/kg")):
27
+ # List of file names to download
28
+ file_names = [
29
+ "default__vector_store.json",
30
+ "docstore.json",
31
+ "graph_store.json",
32
+ "image__vector_store.json",
33
+ "index_store.json"
34
+ ]
35
+
36
+ # Bucket name
37
+ bucket_name = "title_tailors_bucket"
38
+
39
+ # Create the destination directory if it doesn't exist
40
+ os.makedirs("storage/bm25", exist_ok=True)
41
+
42
+ # Loop through the file names and download each one
43
+ for file_name in file_names:
44
+ source_blob_name = f"storage/bm25/{file_name}"
45
+ destination_file_name = f"storage/bm25/{file_name}"
46
+ download_blob(bucket_name, source_blob_name, destination_file_name)
47
+
48
+ # List of file names to download
49
+ file_names = [
50
+ "default__vector_store.json",
51
+ "docstore.json",
52
+ "graph_store.json",
53
+ "image__vector_store.json",
54
+ "index_store.json"
55
+ ]
56
+
57
+ # Bucket name
58
+ bucket_name = "title_tailors_bucket"
59
+
60
+ # Create the destination directory if it doesn't exist
61
+ os.makedirs("storage/kg", exist_ok=True)
62
+
63
+ # Loop through the file names and download each one
64
+ for file_name in file_names:
65
+ source_blob_name = f"storage/kg/{file_name}"
66
+ destination_file_name = f"storage/kg/{file_name}"
67
+ download_blob(bucket_name, source_blob_name, destination_file_name)
68
+ else:
69
+ print("Files already exist in the storage_file directory.")
70
+
71
+
72
+ HF_TOKEN = os.environ.get("HF_TOKEN", None)
73
+
74
+ MISTRAL_API = os.environ.get("MISTRAL_API", None)
75
+ print(f"MISTRAL_API: {MISTRAL_API}")
76
+
77
+ ## logger
78
+ logging.basicConfig(stream=sys.stdout, level=logging.INFO)
79
+ logging.getLogger().handlers = []
80
+ logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
81
+
82
+ # Knowledge graph imports
83
+ from llama_index.core import (
84
+ SimpleDirectoryReader,
85
+ StorageContext,
86
+ KnowledgeGraphIndex,
87
+ load_index_from_storage,
88
+ Settings,
89
+ )
90
+ from llama_index.core.graph_stores import SimpleGraphStore
91
+ from llama_index.embeddings.mistralai import MistralAIEmbedding
92
+ from llama_index.llms.mistralai import MistralAI
93
+
94
+ ## BM25 imports
95
+ from llama_index.core import (
96
+ VectorStoreIndex,
97
+ QueryBundle,
98
+ )
99
+ from llama_index.retrievers.bm25 import BM25Retriever
100
+ from llama_index.core.node_parser import SentenceSplitter
101
+ from llama_index.core.retrievers import (
102
+ BaseRetriever,
103
+ VectorIndexRetriever,
104
+ QueryFusionRetriever,
105
+ )
106
+ from llama_index.core.schema import NodeWithScore
107
+ from llama_index.core.query_engine import RetrieverQueryEngine
108
+ from llama_index.core.postprocessor import SentenceTransformerRerank
109
+ from llama_index.core.response.notebook_utils import (
110
+ display_response,
111
+ display_source_node,
112
+ )
113
+
114
+ # Chat engine
115
+ from llama_index.core import PromptTemplate
116
+ from llama_index.core import chat_engine
117
+ from llama_index.core import memory
118
+
119
+ import nest_asyncio
120
+ nest_asyncio.apply()
121
+
122
+ ################### Loading the LLM via Mistral API
123
+ llm = MistralAI(api_key=MISTRAL_API, model="open-mixtral-8x7b")
124
+ ### Loading the Embedding via Mistral API
125
+ embed_model = MistralAIEmbedding(api_key=MISTRAL_API, model = "mistral-embed")
126
+
127
+ ################### Knowledge Graph Index#################
128
+ ########### While loading from persist only ###############
129
+ Settings.llm = llm
130
+ Settings.embed_model = embed_model
131
+ kg_storage_context = StorageContext.from_defaults(persist_dir="storage/kg")
132
+ kg_index = load_index_from_storage(kg_storage_context)
133
+
134
+ kg_retriever = kg_index.as_retriever(include_text=True,
135
+ response_mode ="tree_summarize",
136
+ embedding_mode="hybrid",
137
+ similarity_top_k=10)
138
+
139
+ ################### BM25 Index #################
140
+ ####################### While loading Indices from persist #######################
141
+ # Settings.llm = llm
142
+ # Settings.embed_model = embed_model
143
+ storage_context_v = StorageContext.from_defaults(persist_dir="storage/bm25")
144
+ index_v = load_index_from_storage(storage_context_v)
145
+
146
+ vector_retriever = index_v.as_retriever(similarity_top_k=10)
147
+ bm25_retriever = BM25Retriever.from_defaults(index=index_v, similarity_top_k=10, verbose=False) ## loading after persist
148
+
149
+
150
+ ###################### Reranker from HuggingFace 🤗
151
+ reranker = SentenceTransformerRerank(top_n=5, model="mixedbread-ai/mxbai-rerank-base-v1", keep_retrieval_score=False)
152
+
153
+ ################ Query Fusion Retriever ##################
154
+ QUERY_GEN_PROMPT = (
155
+ "You are a helpful assistant that generates multiple search queries on astronomy based on a "
156
+ "single input query. Generate {num_queries} search queries for astronomy, one on each line, "
157
+ "related to the following input query:\n"
158
+ "Query: {query}\n"
159
+ "Queries:\n"
160
+ )
161
+
162
+ #### Hybrid (Dense + BM25 + KG) Retriever ####
163
+ hybrid_retriever = QueryFusionRetriever(
164
+ retrievers = [vector_retriever, bm25_retriever, kg_retriever],
165
+ retriever_weights = [0.25, 0.25, 0.50],
166
+ similarity_top_k=5,
167
+ llm=llm,
168
+ num_queries=3, # set this to 1 to disable query generation
169
+ mode="reciprocal_rerank",
170
+ use_async=True,
171
+ verbose=False,
172
+ query_gen_prompt=QUERY_GEN_PROMPT,
173
+ )
174
+
175
+ #### KG Retriever ####
176
+ kg_retriever = QueryFusionRetriever(
177
+ retrievers = [kg_retriever],
178
+ # retriever_weights = [0.25, 0.25, 0.50],
179
+ similarity_top_k=5,
180
+ llm=llm,
181
+ num_queries=3, # set this to 1 to disable query generation
182
+ mode="reciprocal_rerank",
183
+ use_async=True,
184
+ verbose=False,
185
+ query_gen_prompt=QUERY_GEN_PROMPT,
186
+ )
187
+
188
+ #### Hybrid BM25 (Dense + BM25) Retriever ####
189
+ hybrid_bm25_retriever = QueryFusionRetriever(
190
+ retrievers = [vector_retriever, bm25_retriever],
191
+ # retriever_weights = [0.25, 0.25, 0.50],
192
+ similarity_top_k=5,
193
+ llm=llm,
194
+ num_queries=3, # set this to 1 to disable query generation
195
+ mode="reciprocal_rerank",
196
+ use_async=True,
197
+ verbose=False,
198
+ query_gen_prompt=QUERY_GEN_PROMPT,
199
+ )
200
+
201
+ ################# Hybrid Chat Engine ###################
202
+ system_prompt_template = """You are a helpful AI assistant for reaearcher or enthusiast in the domain of astronomy.
203
+ Please check if the following pieces of context has any mention of the keywords provided in the Question. If not then don't know the answer, just say that you don't know.Stop there. Please donot try to make up an answer."""
204
+
205
+ from llama_index.core import chat_engine
206
+ from llama_index.core import memory
207
+
208
+ ######## Hybrid (Dense + BM25 + KG) chat engine #########
209
+ hybrid_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
210
+ Hybrid_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=hybrid_retriever,
211
+ llm=llm,
212
+ memory=hybrid_memory,
213
+ # context_prompt=,
214
+ # condense_prompt=,
215
+ system_prompt=system_prompt_template,
216
+ node_postprocessors=[reranker],
217
+ verbose=False,
218
+ )
219
+
220
+ ######## Hybrid (Dense + BM25 + KG) chat engine #########
221
+ kg_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
222
+ kg_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=kg_retriever,
223
+ llm=llm,
224
+ memory=kg_memory,
225
+ # context_prompt=,
226
+ # condense_prompt=,
227
+ system_prompt=system_prompt_template,
228
+ node_postprocessors=[reranker],
229
+ verbose=False,
230
+ )
231
+
232
+ ######## Hybrid (Dense + BM25 + KG) chat engine #########
233
+ hybrid_bm25_memory = memory.ChatMemoryBuffer.from_defaults(token_limit=30000)
234
+ hybrid_bm25_chat_engine = chat_engine.CondensePlusContextChatEngine(retriever=hybrid_bm25_retriever,
235
+ llm=llm,
236
+ memory=hybrid_bm25_memory,
237
+ # context_prompt=,
238
+ # condense_prompt=,
239
+ system_prompt=system_prompt_template,
240
+ node_postprocessors=[reranker],
241
+ verbose=False,
242
+ )
243
+
244
+ ################## Post Processors ###################
245
+ def get_documents(response):
246
+ """Get the reference documents for a chat engine response.
247
+
248
+ Args:
249
+ response (ChatResponse): The chat engine response object.
250
+
251
+ Returns:
252
+ str: A formatted string containing the reference document paths and page numbers.
253
+ """
254
+ plist = []
255
+ reference_docs={}
256
+ for idx, content in enumerate(response.sources[0].content.split("\n\n")):
257
+ if "page_label" in content:
258
+ plist.append(content)
259
+ for i, entry in enumerate(plist, start=1):
260
+ lines = entry.split('\n')
261
+ page_label = int(lines[0].split(': ')[1])
262
+ file_path = lines[1].split(': ')[1]
263
+
264
+ reference_docs[f"doc{i}"] = {
265
+ "document_path": file_path,
266
+ "page": page_label
267
+ }
268
+ reference = f"""
269
+ \n Reference Docs (document paths, page number):
270
+ ...{reference_docs['doc1']['document_path']}, page {reference_docs['doc1']['page']}
271
+ ...{reference_docs['doc2']['document_path']}, page {reference_docs['doc2']['page']}
272
+ ...{reference_docs['doc3']['document_path']}, page {reference_docs['doc3']['page']}
273
+ ...{reference_docs['doc4']['document_path']}, page {reference_docs['doc4']['page']}
274
+ ...{reference_docs['doc5']['document_path']}, page {reference_docs['doc5']['page']}
275
+ """
276
+ return reference
277
+
278
+ ##################### Query Function #####################
279
+ def get_query(query=""):
280
+ """Get a response from the Hybrid chat engine and format the result with reference documents.
281
+
282
+ Args:
283
+ query (str): The query to send to the chat engine.
284
+
285
+ Returns:
286
+ str: The chat engine response with the reference documents appended.
287
+ """
288
+ response_1 = Hybrid_chat_engine.chat(query)
289
+ response_2 = kg_chat_engine.chat(query)
290
+ response_3 = hybrid_bm25_chat_engine.chat(query)
291
+ reference_1 = get_documents(response_1)
292
+ reference_2 = get_documents(response_2)
293
+ reference_3 = get_documents(response_3)
294
+ reply_1 = response_1.response
295
+ reply_2 = response_2.response
296
+ reply_3 = response_3.response
297
+ result_hybrid = reply_1 + reference_1
298
+ result_kg = reply_2 + reference_2
299
+ result_hybrid_bm25 = reply_3 + reference_3
300
+ return result_hybrid, result_kg, result_hybrid_bm25
301
+
302
+
303
+
304
+ # print(get_query("what is difference between class I and class II ?"))
305
+ # Based on the provided documents, Class I and Class II are two distinct evolutionary stages of young stellar objects (YSOs) in the process of forming stars. The documents describe the evolutionary stages of YSOs as follows:
306
+
307
+ # * Class 0: The formation of a YSO in the central region of a protostellar core with an envelope mass that is much in excess of the YSO mass.
308
+ # * Class I: The collapse of the envelope onto the central object, with the transition between Class 0 and Class I being the point in time at which the envelope mass and the mass of the protostar are nearly equal.
309
+ # * Class II: The emergence of a disk around the central star.
310
+ # * Class III: The dissipation of the disk by various processes such as the formation of planets, photo-evaporation, and tidal stripping.
311
+
312
+ # The documents also mention that an intermediate class between Class 0 and Class I has been proposed, but it is considered to be close to Class I in terms of the evolutionary status of the YSOs.
313
+
314
+ # Regarding the difference between Class I and Class II, the documents do not provide a detailed explanation, but it is suggested that the main difference lies in the presence of a disk around the central star in Class II, which is not present in Class I. The emergence of a disk around the central star is a sign of a more evolved stage in the star formation process. The documents also mention that the exact duration of each evolutionary stage for YSOs is relatively uncertain and depends on the number of objects found in each class, which may be affected by misclassifications due to YSOs being seen edge-on.
315
+ # Reference Docs (document paths, page number):
316
+ # .../content/documents/2405.00095v1.pdf, page 8
317
+ # .../content/documents/2405.00095v1.pdf, page 3
318
+ # .../content/documents/2405.00095v1.pdf, page 9
319
+ # .../content/documents/2405.00095v1.pdf, page 2
320
+ # .../content/documents/2405.00095v1.pdf, page 2
321