make global rag index

app.py

@@ -68,6 +68,8 @@ question: Prior to playing for Michigan State, Keith Nichol played football for
 answer: Norman
 """
 
+global_rag_index = None
+
 class FinchCache(DynamicCache):
     def __init__(self) -> None:
         super().__init__()
@@ -218,9 +220,9 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
     else:
         rag_text = combined_text
     print("Creating RAG index")
-    rag_index = create_rag_index(rag_text)
+    global_rag_index = create_rag_index(rag_text)
     print("Done")
-    state = {"rag_index": rag_index}
+    state = {}
     
     return (
         combined_text, 
@@ -438,13 +440,13 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
     return cache
 
 
-def run_naive_rag_query(vectorstore, query, rag_token_size, prefix, task, few_shot_examples):
+def run_naive_rag_query(query, rag_token_size, prefix, task, few_shot_examples):
     """
     For naive RAG, retrieves top-k chunks (k based on target token size)
     and generates an answer using those chunks.
     """
     k = max(1, rag_token_size // 256)
-    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": k})
+    retriever = global_rag_index.as_retriever(search_type="similarity", search_kwargs={"k": k})
     retrieved_docs = retriever.invoke(query)
     for doc in retrieved_docs:
         print("=================")
@@ -477,9 +479,11 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         print("Target token size for compression: ", target_token_size)
         step_size = 2
         start_time_prefill = time.perf_counter()
+        print("Compressing KV Cache")
         past_key_values = copy.deepcopy(get_compressed_kv_cache(sink_tokens, step_size, target_token_size,
                                                                 context_ids, context_attention_mask,
                                                                 question_ids, question_attention_mask))
+        print("Done")
         compressed_length = past_key_values.get_seq_length()
         print("Context size after compression: ", compressed_length)
         print("Compression rate: ", context_ids.size(1) / compressed_length)
@@ -490,19 +494,17 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         compressed_length = past_key_values.get_seq_length()
         
         
-    # Use the precomputed rag_index from state.
-    rag_index = state.get("rag_index", None)
-    if rag_index is None:
+
+    if global_rag_index is None:
         if combined_text.startswith(prefix):
             rag_text = combined_text[len(prefix):]
         else:
             rag_text = combined_text
-        rag_index = create_rag_index(rag_text, device)
+        global_rag_index = create_rag_index(rag_text, device)
     
     state.update({
         "compressed_cache": past_key_values,
         "compressed_length": compressed_length,
-        "rag_index": rag_index,
         "target_token_size": target_token_size,
         "global_local": percentage,
         "combined_text": combined_text,
@@ -523,7 +525,6 @@ def chat_response_stream(message: str, history: list, state: dict):
     user_message = message
     past_key_values = state["compressed_cache"]
     compressed_length = past_key_values.get_seq_length()
-    rag_index = state["rag_index"]
     retrieval_slider_value = state["retrieval_slider"]
     percentage = state["global_local"]
     
@@ -540,7 +541,7 @@ def chat_response_stream(message: str, history: list, state: dict):
         rag_few_shot = ""
     print("user message: ", user_message)
     if rag_retrieval_size != 0:
-        rag_context = run_naive_rag_query(rag_index, user_message, rag_retrieval_size, rag_prefix, rag_task, rag_few_shot)
+        rag_context = run_naive_rag_query(user_message, rag_retrieval_size, rag_prefix, rag_task, rag_few_shot)
         new_input = rag_context + "\nquestion: " + user_message + suffix + "answer:"
     else:
         new_input = "\nquestion: " + user_message + suffix + "answer:"