saving collection

app.py

@@ -18,6 +18,7 @@ from langchain_docling.loader import ExportType
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache, TextIteratorStreamer
 from transformers.models.llama.modeling_llama import rotate_half
+import uuid
 
 from utils import (
     calculate_tokens_suggest_compression_ratio,
@@ -28,7 +29,7 @@ from utils import (
 
 
 # Initialize the model and tokenizer.
-api_token = os.getenv("HF_TOKEN")
+api_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 model_name = "meta-llama/Llama-3.1-8B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_token)
 model = AutoModelForCausalLM.from_pretrained(model_name, token=api_token, torch_dtype=torch.float16)
@@ -68,8 +69,6 @@ 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__()
@@ -169,7 +168,8 @@ def convert_to_markdown(file_objs, url, do_ocr, do_table_structure):
     docs = loader.load()
     return docs[0].page_content
 
-def create_rag_index(text_no_prefix):
+
+def create_rag_index(collection_name, text_no_prefix):
     """Loads the PDF, splits its text, and builds a vectorstore for naive RAG."""
     text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
                 tokenizer,
@@ -181,13 +181,12 @@ def create_rag_index(text_no_prefix):
             )
     # Concatenate pages and create Document objects.
     docs = [Document(page_content=x) for x in text_splitter.split_text(text_no_prefix)]
-    vectorstore = Chroma.from_documents(documents=docs, embedding=embedding_model)
+    vectorstore = Chroma.from_documents(collection_name=collection_name, persist_directory="./chroma_db", documents=docs, embedding=embedding_model)
     return vectorstore
 
 
 @spaces.GPU
 def auto_convert(file_objs, url, do_ocr, do_table_structure):
-    global global_rag_index
     if file_objs is None and (url is None or url.strip() == ""):
         return (
             gr.update(value=""),
@@ -204,7 +203,7 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
     markdown = convert_to_markdown(file_objs, url, do_ocr, do_table_structure)
     print("Done")
     combined_text = prefix + markdown
-    print("Calculating tokens")
+    print("Suggestioning Compression ratio")
     token_count, suggestions, _ = calculate_tokens_suggest_compression_ratio(combined_text, tokenizer, model)
     print("Done")
     min_ratio = min(suggestions)
@@ -220,10 +219,10 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
         rag_text = combined_text[len(prefix):]
     else:
         rag_text = combined_text
-    print("Creating RAG index")
-    global_rag_index = create_rag_index(rag_text)
+    collection_name = "default_collection_" + uuid.uuid4().hex[:6]
+    rag_index = create_rag_index(collection_name, rag_text)
+    state = {"rag_index": collection_name}
     print("Done")
-    state = {}
     
     return (
         combined_text, 
@@ -441,14 +440,14 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
     return cache
 
 
-def run_naive_rag_query(query, rag_token_size, prefix, task, few_shot_examples):
+def run_naive_rag_query(collection_name, 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.
     """
-    global global_rag_index
     k = max(1, rag_token_size // 256)
-    retriever = global_rag_index.as_retriever(search_type="similarity", search_kwargs={"k": k})
+    vectorstore = Chroma(persist_directory="./chroma_db", embedding=embedding_model, collection_name=collection_name)
+    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": k})
     retrieved_docs = retriever.invoke(query)
     for doc in retrieved_docs:
         print("=================")
@@ -466,7 +465,6 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
     """
     Prepares the compressed KV cache. Uses the precomputed rag_index from state.
     """
-    global global_rag_index
     percentage = int(global_local_value.replace('%', ''))
     question_text = task_description + "\n" + few_shot
     context_encoding = tokenizer(combined_text, return_tensors="pt").to(device)
@@ -482,11 +480,9 @@ 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)
@@ -497,17 +493,21 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         compressed_length = past_key_values.get_seq_length()
         
         
-
-    if global_rag_index is None:
+    # Use the precomputed rag_index from state.
+    collection_name = state.get("rag_index", None)
+    if collection_name is None:
+        print("Collection name not found creating a new one.")
         if combined_text.startswith(prefix):
             rag_text = combined_text[len(prefix):]
         else:
             rag_text = combined_text
-        global_rag_index = create_rag_index(rag_text, device)
-    
+        collection_name = "default_collection_" + uuid.uuid4().hex[:6]
+        rag_index = create_rag_index(collection_name, rag_text)
+
     state.update({
         "compressed_cache": past_key_values,
         "compressed_length": compressed_length,
+        "rag_index": collection_name,
         "target_token_size": target_token_size,
         "global_local": percentage,
         "combined_text": combined_text,
@@ -528,6 +528,7 @@ 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()
+    collection_name = state["rag_index"]
     retrieval_slider_value = state["retrieval_slider"]
     percentage = state["global_local"]
     
@@ -544,7 +545,8 @@ 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(user_message, rag_retrieval_size, rag_prefix, rag_task, rag_few_shot)
+        print("Running RAG query")
+        rag_context = run_naive_rag_query(collection_name, 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:"
@@ -724,4 +726,4 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
                 type="messages"
             )
 
-demo.queue().launch()
+demo.queue(max_size=16).launch()