update

app.py

@@ -3,6 +3,7 @@ import math
 import os
 import time
 from threading import Thread
+import uuid
 
 import gradio as gr
 import spaces
@@ -18,7 +19,6 @@ 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,
@@ -26,8 +26,6 @@ from utils import (
     update_retrieval_context,
 )
 
-
-
 # Initialize the model and tokenizer.
 api_token = os.getenv("HUGGING_FACE_HUB_TOKEN")
 model_name = "meta-llama/Llama-3.1-8B-Instruct"
@@ -37,12 +35,11 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.eval()
 model.to(device)
 embedding_model = HuggingFaceBgeEmbeddings(
-        model_name="BAAI/bge-large-en-v1.5",
-        model_kwargs={"device": str(device)},
-        encode_kwargs={"normalize_embeddings": True},
-        query_instruction=""
-    )
-
+    model_name="BAAI/bge-large-en-v1.5",
+    model_kwargs={"device": str(device)},
+    encode_kwargs={"normalize_embeddings": True},
+    query_instruction=""
+)
 
 # Create a chat template and split into prefix and suffix.
 content_system = ""
@@ -121,7 +118,6 @@ class FinchCache(DynamicCache):
         self._seen_tokens = new_length
 
     def save(self, path: str):
-        """Save the cache to disk, moving tensors to CPU."""
         try:
             os.makedirs(os.path.dirname(path), exist_ok=True)
             torch.save(
@@ -133,7 +129,6 @@ class FinchCache(DynamicCache):
 
     @classmethod
     def load(cls, path: str, device: str = "cpu") -> "FinchCache":
-        """Load the cache from disk and move tensors to the specified device."""
         data = torch.load(path, map_location=device)
         cache = cls()
         cache.key_cache = [k.to(device) for k in data["key_cache"]]
@@ -141,8 +136,6 @@ class FinchCache(DynamicCache):
         cache._seen_tokens = cache.value_cache[0].size(2) if cache.value_cache else 0
         return cache
 
-
-
 def convert_to_markdown(file_objs, url, do_ocr, do_table_structure):
     file_path = file_objs if file_objs is not None else url
     pipeline_options = PdfPipelineOptions()
@@ -154,12 +147,8 @@ def convert_to_markdown(file_objs, url, do_ocr, do_table_structure):
     )
     doc_converter = DocumentConverter(
         allowed_formats=[InputFormat.PDF],
-        format_options={
-            InputFormat.PDF: pdf_format_options
-        }
+        format_options={InputFormat.PDF: pdf_format_options}
     )
-
-    # Pass the custom converter to the DoclingLoader.
     loader = DoclingLoader(
         file_path=file_path,
         export_type=ExportType.MARKDOWN,
@@ -168,39 +157,51 @@ def convert_to_markdown(file_objs, url, do_ocr, do_table_structure):
     docs = loader.load()
     return docs[0].page_content
 
-
 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,
-                chunk_size=256,
-                chunk_overlap=0,
-                add_start_index=True,
-                strip_whitespace=True,
-                separators=["\n\n", "\n", ".", " ", ""],
-            )
-    # Concatenate pages and create Document objects.
+        tokenizer,
+        chunk_size=256,
+        chunk_overlap=0,
+        add_start_index=True,
+        strip_whitespace=True,
+        separators=["\n\n", "\n", ".", " ", ""],
+    )
     docs = [Document(page_content=x) for x in text_splitter.split_text(text_no_prefix)]
     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):
+    # When a new file/URL is loaded, disable chat (compression not done)
+    chat_status = "Document not compressed yet. Please compress the document to enable chat."
     if file_objs is None and (url is None or url.strip() == ""):
         return (
             gr.update(value=""),
             "Number of tokens before compression: ",
-            gr.update(), 
+            gr.update(),
             "Number of tokens after compression: ",
             0,
-            gr.update(interactive=False),  # Disable compress button when no input.
+            gr.update(interactive=False),
             False,
-            {}  # return an empty state dictionary
+            {},
+            chat_status
         )
-    # Convert the document to markdown.
     print("Converting to markdown")
-    markdown = convert_to_markdown(file_objs, url, do_ocr, do_table_structure)
+    try:
+        markdown = convert_to_markdown(file_objs, url, do_ocr, do_table_structure)
+    except Exception as e:
+        print("Error converting to markdown:", e)
+        return (
+            gr.update(value="Error converting document to markdown. Please try uploading another document format."),
+            "Number of tokens before compression: ",
+            gr.update(),
+            "Number of tokens after compression: ",
+            0,
+            gr.update(interactive=False),
+            False,
+            {},
+            chat_status
+        )
     print("Done")
     combined_text = prefix + markdown
     print("Suggestioning Compression ratio")
@@ -213,8 +214,6 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
     token_count_str = f"Number of tokens before compression: {token_count}"
     retrieval_str = f"Number of tokens after compression: {retrieval_tokens}"
     slider_update = gr.update(value=default_ratio, minimum=min_ratio, maximum=max_ratio, step=1)
-    
-    # Create the RAG index immediately.
     if combined_text.startswith(prefix):
         rag_text = combined_text[len(prefix):]
     else:
@@ -223,18 +222,17 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
     rag_index = create_rag_index(collection_name, rag_text)
     state = {"rag_index": collection_name}
     print("Done")
-    
     return (
-        combined_text, 
-        token_count_str, 
-        slider_update, 
-        retrieval_str, 
-        token_count, 
-        gr.update(interactive=True),
+        combined_text,
+        token_count_str,
+        slider_update,
+        retrieval_str,
+        token_count,
+        gr.update(interactive=True),  # Enable compress button if conversion succeeds.
         False,
-        state
+        state,
+        chat_status
     )
-    
 
 def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_ids, context_attention_mask, question_ids, question_attention_mask):
     device = model.device
@@ -250,32 +248,18 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
     max_context_tokens_allowed = model.config.max_position_embeddings - question_len
     if total_len > max_context_tokens_allowed:
         num_chunks = max(step_size, math.ceil(total_len / max_context_tokens_allowed))
-
     if total_len <= sink_tokens or num_chunks == 1:
-        # If the context is too short or only one chunk is desired, use the entire context.
         context_ids_list = [context_ids]
         context_attention_mask_list = [context_attention_mask]
     else:
-        # Calculate how many tokens remain after the sink tokens.
         remainder_len = total_len - sink_tokens
-
-        # Compute the base tokens per chunk and any leftover.
         base = remainder_len // num_chunks
         leftover = remainder_len % num_chunks
-
-        # Build a list of chunk sizes.
-        # First chunk gets the sink tokens plus base tokens.
         chunk_sizes = [sink_tokens + base]
-
-        # Chunks 2 to num_chunks-1 get base tokens each.
         for _ in range(num_chunks - 2):
             chunk_sizes.append(base)
-
-        # The last chunk gets the remaining tokens (base + leftover).
         if num_chunks > 1:
             chunk_sizes.append(base + leftover)
-
-        # Now slice the context using the calculated sizes.
         context_ids_list = []
         context_attention_mask_list = []
         offset = 0
@@ -284,33 +268,23 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
             context_ids_list.append(context_ids[:, offset:end])
             context_attention_mask_list.append(context_attention_mask[:, offset:end])
             offset = end
-
-    # (Optional) Continue with the rest of your processing…
     len_rest = max(total_len - sink_tokens, 1)
     compression_factor = len_rest // target_token_size
     if compression_factor < 1:
         compression_factor = 1
-
     tokenized_doc_chunks = []
     for ids_chunk, mask_chunk in zip(context_ids_list, context_attention_mask_list):
         tokenized_doc_chunks.append({"input_ids": ids_chunk, "attention_mask": mask_chunk})
-
     print("Number of chunks: ", len(tokenized_doc_chunks))
-
     rotary_emb = model.model.rotary_emb.to(device)
     inv_freq = rotary_emb.inv_freq
     batch_size = question_ids.size(0)
     ones_mask = torch.ones(batch_size, 1, dtype=question_attention_mask.dtype, device=device)
-
     cache = FinchCache()
     past_cache_len = 0
     past_attention_mask = torch.zeros(batch_size, 0, dtype=question_attention_mask.dtype, device=device)
     num_chunks = len(tokenized_doc_chunks)
-
-    # Prepare a shared dictionary for hook outputs.
     query_context_matrices = {}
-
-    # Define a hook function that uses a per-chunk offset stored on self.
     def query_hook_fn(module, input, output):
         layer_idx = getattr(module, "layer_idx", None)
         if layer_idx is not None:
@@ -323,26 +297,17 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
                 .transpose(1, 2)
                 .contiguous()
             )
-            # Use self._current_chunk_offset to select only the new tokens.
             query_context_matrices[layer_idx] = query_states[:, :, _current_chunk_offset:, :].clone()
-
-    # Pre-register hooks for all layers only once.
     hooks = []
     for i, layer in enumerate(model.model.layers):
-        layer.self_attn.q_proj.layer_idx = i  # For tracking.
+        layer.self_attn.q_proj.layer_idx = i
         layer.self_attn.q_proj.num_query_heads = layer.self_attn.config.num_attention_heads
         hook = layer.self_attn.q_proj.register_forward_hook(query_hook_fn)
         hooks.append(hook)
-
-    # Process each document chunk sequentially.
     for j, tokenized_doc_chunk in enumerate(tokenized_doc_chunks):
         current_seq_length = tokenized_doc_chunk["input_ids"].size(1)
-        # Save the offset in an attribute the hook can access.
         _current_chunk_offset = current_seq_length
-        # Clear the dictionary from any previous chunk.
         query_context_matrices.clear()
-
-        # These chunks are already on the device.
         chunk_input_ids = tokenized_doc_chunk["input_ids"].contiguous()
         chunk_attention_mask = tokenized_doc_chunk["attention_mask"].contiguous()
         segment_attention_mask = torch.cat(
@@ -350,7 +315,6 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
         ).contiguous()
         current_input_ids = torch.cat([chunk_input_ids, question_ids], dim=-1).contiguous()
         current_attention_mask = torch.cat([segment_attention_mask, question_attention_mask], dim=-1).contiguous()
-
         past_seen_tokens = cache.get_seq_length() if cache is not None else 0
         cache_position = torch.arange(
             past_seen_tokens + chunk_input_ids.shape[1],
@@ -366,7 +330,6 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
             cache_position=cache_position,
             batch_size=current_input_ids.size(0),
         ).contiguous()
-
         with torch.no_grad():
             outputs = model.model(
                 input_ids=current_input_ids,
@@ -374,9 +337,7 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
                 past_key_values=cache,
             )
             cache = outputs.past_key_values
-
         len_question = question_ids.size(1)
-        # Now, for each transformer layer, update the cache using the query/key attention.
         for layer_idx in range(len(model.model.layers)):
             key_matrix = cache.key_cache[layer_idx]
             query_matrix = query_context_matrices[layer_idx]
@@ -392,13 +353,11 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
             query_matrix = (query_matrix * cos) + (rotate_half(query_matrix) * sin)
             num_repeats = model.config.num_attention_heads // model.config.num_key_value_heads
             key_matrix = repeat_kv(key_matrix, num_repeats)
-
             scaling = math.sqrt(model.config.head_dim)
             attention_matrix = torch.matmul(query_matrix, key_matrix.transpose(2, 3)) / scaling
             causal_mask_sliced = causal_mask[:, :, :, : key_matrix.shape[-2]]
             attention_matrix = attention_matrix + causal_mask_sliced
             attention_matrix = torch.nn.functional.softmax(attention_matrix, dim=-1, dtype=torch.float32).to(query_matrix.dtype)
-            # Normalization
             tol = 1e-8
             binary_mask = (torch.abs(causal_mask_sliced.to(torch.float32)) < tol).to(torch.float32)
             non_zero_counts = binary_mask.sum(dim=3, keepdim=True)
@@ -423,30 +382,21 @@ def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_i
                 to_keep_new = int(current_seq_length // compression_factor)
                 k = min(past_cache_len + to_keep_new, full_context_size)
             else:
-                desired_final = sink_tokens + target_token_size + len_question# TODO remember to include the question tokens
+                desired_final = sink_tokens + target_token_size + len_question
                 k = desired_final if full_context_size >= desired_final else full_context_size
             k = max(k, sink_tokens)
             selected_indices = torch.topk(attention_matrix, k, dim=-1).indices
             selected_indices, _ = torch.sort(selected_indices, dim=-1)
             cache.compress_cache(layer_idx, selected_indices, inv_freq)
-
         past_cache_len = cache._seen_tokens
         past_attention_mask = torch.ones(1, past_cache_len, device=device)
-
-    # Remove the hooks once after all chunks are processed.
     for hook in hooks:
         hook.remove()
-
     return cache
 
-
 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.
-    """
     k = max(1, rag_token_size // 256)
-    vectorstore = Chroma(persist_directory="./chroma_db", embedding=embedding_model, collection_name=collection_name)
+    vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=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:
@@ -454,17 +404,11 @@ def run_naive_rag_query(collection_name, query, rag_token_size, prefix, task, fe
         print(doc.page_content)
         print("=================")
     formatted_context = "\n\n".join([doc.page_content for doc in retrieved_docs])
-    
     rag_context = prefix + "Retrieved context: \n" + formatted_context + task + few_shot_examples
-    
     return rag_context
 
-
 @spaces.GPU
 def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_local_value, task_description, few_shot, state):
-    """
-    Prepares the compressed KV cache. Uses the precomputed rag_index from state.
-    """
     percentage = int(global_local_value.replace('%', ''))
     question_text = task_description + "\n" + few_shot
     context_encoding = tokenizer(combined_text, return_tensors="pt").to(device)
@@ -474,15 +418,23 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
     question_ids = question_encoding["input_ids"]
     question_attention_mask = question_encoding["attention_mask"]
     retrieval_context_length = int(context_ids.size(1) / retrieval_slider_value)
-    
+    # Compute token breakdown for display (KV compress vs RAG tokens)
+    rag_tokens = int(retrieval_context_length * (1.0 - (percentage / 100)))
+    kv_tokens = retrieval_context_length - rag_tokens
+    print(f"KV Compress Tokens: {kv_tokens}, RAG Tokens: {rag_tokens}")
     if percentage > 0:
         target_token_size = int(retrieval_context_length * (percentage / 100))
         print("Target token size for compression: ", target_token_size)
         step_size = 2
         start_time_prefill = time.perf_counter()
-        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))
+        try:
+            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))
+        except Exception as e:
+            print("Error during KV cache compression:", e)
+            state["error"] = "Error during KV cache compression. Please try lowering the compression ratio and try again."
+            return state, False
         compressed_length = past_key_values.get_seq_length()
         print("Context size after compression: ", compressed_length)
         print("Compression rate: ", context_ids.size(1) / compressed_length)
@@ -491,15 +443,11 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         target_token_size = 0
         past_key_values = FinchCache()
         compressed_length = past_key_values.get_seq_length()
-
-    cache_name = "default_cache_" + uuid.uuid4().hex[:6]
     cache_name = "default_cache_" + uuid.uuid4().hex[:6] + ".pt"
     save_dir = "./cache_dir"
     os.makedirs(save_dir, exist_ok=True)
     save_path = os.path.join(save_dir, cache_name)
     past_key_values.save(save_path)
-        
-    # 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.")
@@ -509,7 +457,6 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
             rag_text = combined_text
         collection_name = "default_collection_" + uuid.uuid4().hex[:6]
         rag_index = create_rag_index(collection_name, rag_text)
-
     state.update({
         "compressed_cache": save_path,
         "compressed_length": compressed_length,
@@ -520,32 +467,28 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         "task_description": task_description,
         "few_shot": few_shot,
         "retrieval_slider": retrieval_context_length,
-        "prefill_time": time.perf_counter() - start_time_prefill
+        "prefill_time": time.perf_counter() - start_time_prefill,
+        "compression_done": True,
+        "tokens_breakdown": f"KV Compress Tokens: {kv_tokens}, RAG Tokens: {rag_tokens}",
+        "chat_feedback": "Document compressed successfully. You can now chat."
     })
     return state, True
 
-
 @spaces.GPU
 def chat_response_stream(message: str, history: list, state: dict):
-    """
-    Generates a chat response with streaming output.
-    Returns a simple string (not a list of message dicts) for ChatInterface.
-    """
+    # Check if the document is compressed before allowing chat
+    if not state.get("compression_done", False) or "compressed_cache" not in state:
+        yield "Document not compressed yet. Please compress the document first to enable chat."
+        return
     user_message = message
     save_path = state["compressed_cache"]
     past_key_values = FinchCache.load(save_path, device=model.device)
-    try:
-        os.remove(save_path)
-    except Exception as e:
-        print(f"Error removing cache file: {e}")
     compressed_length = past_key_values.get_seq_length()
     collection_name = state["rag_index"]
     retrieval_slider_value = state["retrieval_slider"]
     percentage = state["global_local"]
-    
     rag_retrieval_size = int(retrieval_slider_value * (1.0 - (percentage / 100)))
     print("RAG retrieval size: ", rag_retrieval_size)
-    
     if percentage == 0:
         rag_prefix = prefix
         rag_task = state["task_description"]
@@ -565,7 +508,6 @@ def chat_response_stream(message: str, history: list, state: dict):
     eos_block = torch.full((1, compressed_length), tokenizer.eos_token_id, device=device, dtype=torch.long)
     new_input_ids = torch.cat([eos_block, tokenized_new_input["input_ids"]], dim=-1)
     new_attention_mask = torch.cat([torch.ones((1, compressed_length), device=device), tokenized_new_input["attention_mask"]], dim=-1)
-    
     print("New input is: ", new_input)
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
     generate_kwargs = dict(
@@ -583,18 +525,28 @@ def chat_response_stream(message: str, history: list, state: dict):
     )
     t = Thread(target=model.generate, kwargs=generate_kwargs)
     t.start()
-    
     full_output = ""
     for text in streamer:
         full_output += text
         time.sleep(0.05)
         yield full_output
-        
     state["compressed_cache"] = past_key_values
     return full_output
 
+def update_token_breakdown(token_count, retrieval_slider_value, global_local_value):
+    try:
+        token_count = int(token_count)
+        slider_val = float(retrieval_slider_value)
+        percentage = int(global_local_value.replace('%', ''))
+        retrieval_context_length = int(token_count / slider_val)
+        rag_tokens = int(retrieval_context_length * (1 - (percentage / 100)))
+        kv_tokens = retrieval_context_length - rag_tokens
+        return f"KV Compress Tokens: {kv_tokens}, RAG Tokens: {rag_tokens}"
+    except Exception as e:
+        return "Token breakdown unavailable."
+
 ##########################################################################
-# Gradio Interface: note that we now use ChatInterface instead of a Chatbot.
+# Gradio Interface
 ##########################################################################
 CSS = """
 body {
@@ -694,39 +646,57 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
             token_count_text = gr.Markdown("Number of tokens before compression: ")
             retrieval_slider = gr.Slider(label="Select Compression Rate", minimum=1, maximum=32, step=1, value=2)
             retrieval_info_text = gr.Markdown("Number of tokens after compression: ")
+            # New widget for token breakdown (KV vs RAG)
+            tokens_breakdown_text = gr.Markdown("Token breakdown will appear here.")
             global_local_slider = gr.Radio(label="Global vs Local (0 is all RAG, 100 is all global)", 
                                            choices=["0%", "25%", "50%", "75%", "100%"], value="75%")
             compress_button = gr.Button("Compress Document", interactive=False, elem_classes="upload-button")
+            # New widget for chat status feedback
+            chat_status_text = gr.Markdown("Document not compressed yet. Please compress the document to enable chat.")
             
             file_input.change(
                 fn=auto_convert,
                 inputs=[file_input, url_input, do_ocr, do_table],
-                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
             )
             url_input.change(
                 fn=auto_convert,
                 inputs=[file_input, url_input, do_ocr, do_table],
-                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
             )
             do_ocr.change(
                 fn=auto_convert,
                 inputs=[file_input, url_input, do_ocr, do_table],
-                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
             )
             do_table.change(
                 fn=auto_convert,
                 inputs=[file_input, url_input, do_ocr, do_table],
-                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
             )
             retrieval_slider.change(
                 fn=update_retrieval_context,
                 inputs=[hidden_token_count, retrieval_slider],
                 outputs=retrieval_info_text
             )
+            # Update token breakdown when slider or global/local changes
+            retrieval_slider.change(
+                fn=update_token_breakdown,
+                inputs=[hidden_token_count, retrieval_slider, global_local_slider],
+                outputs=tokens_breakdown_text
+            )
+            global_local_slider.change(
+                fn=update_token_breakdown,
+                inputs=[hidden_token_count, retrieval_slider, global_local_slider],
+                outputs=tokens_breakdown_text
+            )
             compress_button.click(
                 fn=prepare_compression_and_rag,
                 inputs=[markdown_output, retrieval_slider, global_local_slider, task_description_input, few_shot_input, compressed_doc_state],
                 outputs=[compressed_doc_state, compression_done]
+            ).then(
+                fn=lambda state: gr.update(value="Document compressed successfully. You can now chat."),
+                outputs=chat_status_text
             )
         
         with gr.Column(elem_classes="chatbot-container"):