update

app.py

@@ -19,7 +19,9 @@ 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 threading
+import shutil
+import time
 from utils import (
     calculate_tokens_suggest_compression_ratio,
     repeat_kv,
@@ -66,6 +68,44 @@ question: Prior to playing for Michigan State, Keith Nichol played football for
 answer: Norman
 """
 
+
+
+CHROMA_DB_DIR = "./chroma_db"
+CACHE_DIR = "./cache_dir"
+EXPIRATION_SECONDS = 3600
+
+def background_cleanup():
+    while True:
+        current_time = int(time.time())
+
+        # Clean Chroma collections
+        if os.path.exists(CHROMA_DB_DIR):
+            for dirname in os.listdir(CHROMA_DB_DIR):
+                parts = dirname.split("_")
+                if len(parts) >= 3 and parts[1].isdigit():
+                    timestamp = int(parts[1])
+                    if current_time - timestamp > EXPIRATION_SECONDS:
+                        path = os.path.join(CHROMA_DB_DIR, dirname)
+                        shutil.rmtree(path, ignore_errors=True)
+                        print(f"[Cleanup] Deleted Chroma collection: {path}")
+
+        # Clean cache files
+        if os.path.exists(CACHE_DIR):
+            for filename in os.listdir(CACHE_DIR):
+                parts = filename.split("_")
+                if len(parts) >= 3 and parts[1].isdigit():
+                    timestamp = int(parts[1])
+                    if current_time - timestamp > EXPIRATION_SECONDS:
+                        path = os.path.join(CACHE_DIR, filename)
+                        os.remove(path)
+                        print(f"[Cleanup] Deleted cache file: {path}")
+
+        time.sleep(600)
+
+cleanup_thread = threading.Thread(target=background_cleanup, daemon=True)
+cleanup_thread.start()
+
+
 class FinchCache(DynamicCache):
     def __init__(self) -> None:
         super().__init__()
@@ -154,8 +194,11 @@ def convert_to_markdown(file_objs, url, do_ocr, do_table_structure):
         export_type=ExportType.MARKDOWN,
         converter=doc_converter
     )
-    docs = loader.load()
-    return docs[0].page_content
+    try:
+        docs = loader.load()
+        return docs[0].page_content
+    except Exception as e:
+        raise RuntimeError(f"Failed to convert document to markdown: {e}")
 
 def create_rag_index(collection_name, text_no_prefix):
     text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
@@ -184,15 +227,15 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
             gr.update(interactive=False),
             False,
             {},
-            chat_status
+            chat_status,
+            gr.update(interactive=False)  # Disable chat interface
         )
     print("Converting to markdown")
     try:
         markdown = convert_to_markdown(file_objs, url, do_ocr, do_table_structure)
-    except Exception as e:
-        print("Error converting to markdown:", e)
+    except RuntimeError as e:
         return (
-            gr.update(value="Error converting document to markdown. Please try uploading another document format."),
+            gr.update(value=f"{str(e)} Please try uploading another document format."),
             "Number of tokens before compression: ",
             gr.update(),
             "Number of tokens after compression: ",
@@ -200,8 +243,10 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
             gr.update(interactive=False),
             False,
             {},
-            chat_status
+            chat_status,
+            gr.update(interactive=False)  # Disable chat interface on error
         )
+
     print("Done")
     combined_text = prefix + markdown
     print("Suggestioning Compression ratio")
@@ -218,7 +263,8 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
         rag_text = combined_text[len(prefix):]
     else:
         rag_text = combined_text
-    collection_name = "default_collection_" + uuid.uuid4().hex[:6]
+    current_timestamp = int(time.time())
+    collection_name = f"default_{current_timestamp}_{uuid.uuid4().hex[:6]}"
     rag_index = create_rag_index(collection_name, rag_text)
     state = {"rag_index": collection_name}
     print("Done")
@@ -231,168 +277,172 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
         gr.update(interactive=True),  # Enable compress button if conversion succeeds.
         False,
         state,
-        chat_status
+        chat_status,
+        gr.update(interactive=False)  # Ensure chat remains disabled until compression
     )
 
 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
-    dtype = model.dtype
-    sink_tokens = sink_tokens
-    num_chunks = step_size
-    context_ids = context_ids.to(device)
-    context_attention_mask = context_attention_mask.to(device)
-    question_ids = question_ids.to(device)
-    question_attention_mask = question_attention_mask.to(device)
-    question_len = question_ids.size(1)
-    total_len = context_ids.size(1)
-    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:
-        context_ids_list = [context_ids]
-        context_attention_mask_list = [context_attention_mask]
-    else:
-        remainder_len = total_len - sink_tokens
-        base = remainder_len // num_chunks
-        leftover = remainder_len % num_chunks
-        chunk_sizes = [sink_tokens + base]
-        for _ in range(num_chunks - 2):
-            chunk_sizes.append(base)
-        if num_chunks > 1:
-            chunk_sizes.append(base + leftover)
-        context_ids_list = []
-        context_attention_mask_list = []
-        offset = 0
-        for size in chunk_sizes:
-            end = offset + size
-            context_ids_list.append(context_ids[:, offset:end])
-            context_attention_mask_list.append(context_attention_mask[:, offset:end])
-            offset = end
-    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)
-    query_context_matrices = {}
-    def query_hook_fn(module, input, output):
-        layer_idx = getattr(module, "layer_idx", None)
-        if layer_idx is not None:
-            query_states = output.detach()
-            bsz, seq_len, hidden_dim = query_states.size()
-            num_query_heads = module.num_query_heads
-            head_dim = hidden_dim // num_query_heads
-            query_states = (
-                query_states.view(bsz, seq_len, num_query_heads, head_dim)
-                .transpose(1, 2)
-                .contiguous()
-            )
-            query_context_matrices[layer_idx] = query_states[:, :, _current_chunk_offset:, :].clone()
-    hooks = []
-    for i, layer in enumerate(model.model.layers):
-        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)
-    for j, tokenized_doc_chunk in enumerate(tokenized_doc_chunks):
-        current_seq_length = tokenized_doc_chunk["input_ids"].size(1)
-        _current_chunk_offset = current_seq_length
-        query_context_matrices.clear()
-        chunk_input_ids = tokenized_doc_chunk["input_ids"].contiguous()
-        chunk_attention_mask = tokenized_doc_chunk["attention_mask"].contiguous()
-        segment_attention_mask = torch.cat(
-            [past_attention_mask, chunk_attention_mask, ones_mask], dim=-1
-        ).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],
-            past_seen_tokens + current_input_ids.shape[1],
-            device=device
-        )
-        causal_mask = model.model._prepare_4d_causal_attention_mask_with_cache_position(
-            current_attention_mask,
-            sequence_length=question_ids.size(1),
-            target_length=current_attention_mask.size(-1),
-            dtype=dtype,
-            device=device,
-            cache_position=cache_position,
-            batch_size=current_input_ids.size(0),
-        ).contiguous()
-        with torch.no_grad():
-            outputs = model.model(
-                input_ids=current_input_ids,
-                use_cache=True,
-                past_key_values=cache,
-            )
-            cache = outputs.past_key_values
-        len_question = question_ids.size(1)
-        for layer_idx in range(len(model.model.layers)):
-            key_matrix = cache.key_cache[layer_idx]
-            query_matrix = query_context_matrices[layer_idx]
-            layer_cache_pos = torch.arange(
-                past_cache_len + current_seq_length,
-                past_cache_len + current_seq_length + len_question,
+    try:
+        device = model.device
+        dtype = model.dtype
+        sink_tokens = sink_tokens
+        num_chunks = step_size
+        context_ids = context_ids.to(device)
+        context_attention_mask = context_attention_mask.to(device)
+        question_ids = question_ids.to(device)
+        question_attention_mask = question_attention_mask.to(device)
+        question_len = question_ids.size(1)
+        total_len = context_ids.size(1)
+        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:
+            context_ids_list = [context_ids]
+            context_attention_mask_list = [context_attention_mask]
+        else:
+            remainder_len = total_len - sink_tokens
+            base = remainder_len // num_chunks
+            leftover = remainder_len % num_chunks
+            chunk_sizes = [sink_tokens + base]
+            for _ in range(num_chunks - 2):
+                chunk_sizes.append(base)
+            if num_chunks > 1:
+                chunk_sizes.append(base + leftover)
+            context_ids_list = []
+            context_attention_mask_list = []
+            offset = 0
+            for size in chunk_sizes:
+                end = offset + size
+                context_ids_list.append(context_ids[:, offset:end])
+                context_attention_mask_list.append(context_attention_mask[:, offset:end])
+                offset = end
+        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)
+        query_context_matrices = {}
+        def query_hook_fn(module, input, output):
+            layer_idx = getattr(module, "layer_idx", None)
+            if layer_idx is not None:
+                query_states = output.detach()
+                bsz, seq_len, hidden_dim = query_states.size()
+                num_query_heads = module.num_query_heads
+                head_dim = hidden_dim // num_query_heads
+                query_states = (
+                    query_states.view(bsz, seq_len, num_query_heads, head_dim)
+                    .transpose(1, 2)
+                    .contiguous()
+                )
+                query_context_matrices[layer_idx] = query_states[:, :, _current_chunk_offset:, :].clone()
+        hooks = []
+        for i, layer in enumerate(model.model.layers):
+            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)
+        for j, tokenized_doc_chunk in enumerate(tokenized_doc_chunks):
+            current_seq_length = tokenized_doc_chunk["input_ids"].size(1)
+            _current_chunk_offset = current_seq_length
+            query_context_matrices.clear()
+            chunk_input_ids = tokenized_doc_chunk["input_ids"].contiguous()
+            chunk_attention_mask = tokenized_doc_chunk["attention_mask"].contiguous()
+            segment_attention_mask = torch.cat(
+                [past_attention_mask, chunk_attention_mask, ones_mask], dim=-1
+            ).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],
+                past_seen_tokens + current_input_ids.shape[1],
                 device=device
             )
-            position_ids = layer_cache_pos.unsqueeze(0)
-            cos, sin = rotary_emb(query_matrix, position_ids)
-            cos = cos.unsqueeze(1)
-            sin = sin.unsqueeze(1)
-            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)
-            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)
-            non_zero_counts = torch.clamp_min(non_zero_counts, 1.0).to(attention_matrix.dtype)
-            attention_matrix = attention_matrix / non_zero_counts
-            if j != num_chunks - 1: 
-                attention_matrix = attention_matrix[:, :, :, : past_cache_len + current_seq_length].clone().contiguous()
-            else:
-                attention_matrix = attention_matrix[:, :, :, : past_cache_len + current_seq_length + len_question].clone().contiguous()
-            attention_matrix = torch.sum(attention_matrix, dim=-2)
-            attention_matrix = attention_matrix.view(
-                attention_matrix.size(0), model.config.num_key_value_heads, num_repeats, -1
-            ).sum(dim=2)
-            full_context_size = attention_matrix.size(-1)
-            attention_matrix[..., :sink_tokens] = float("inf")
-            if j == num_chunks - 1:
-                attention_matrix[..., -len_question:] = float("inf")
-            if j == 0:
-                k = int(sink_tokens + (max(0, current_seq_length - sink_tokens) // compression_factor))
-                k = min(k + past_cache_len, full_context_size)
-            elif j < num_chunks - 1:
-                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
-                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)
-    for hook in hooks:
-        hook.remove()
-    return cache
+            causal_mask = model.model._prepare_4d_causal_attention_mask_with_cache_position(
+                current_attention_mask,
+                sequence_length=question_ids.size(1),
+                target_length=current_attention_mask.size(-1),
+                dtype=dtype,
+                device=device,
+                cache_position=cache_position,
+                batch_size=current_input_ids.size(0),
+            ).contiguous()
+            with torch.no_grad():
+                outputs = model.model(
+                    input_ids=current_input_ids,
+                    use_cache=True,
+                    past_key_values=cache,
+                )
+                cache = outputs.past_key_values
+            len_question = question_ids.size(1)
+            for layer_idx in range(len(model.model.layers)):
+                key_matrix = cache.key_cache[layer_idx]
+                query_matrix = query_context_matrices[layer_idx]
+                layer_cache_pos = torch.arange(
+                    past_cache_len + current_seq_length,
+                    past_cache_len + current_seq_length + len_question,
+                    device=device
+                )
+                position_ids = layer_cache_pos.unsqueeze(0)
+                cos, sin = rotary_emb(query_matrix, position_ids)
+                cos = cos.unsqueeze(1)
+                sin = sin.unsqueeze(1)
+                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)
+                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)
+                non_zero_counts = torch.clamp_min(non_zero_counts, 1.0).to(attention_matrix.dtype)
+                attention_matrix = attention_matrix / non_zero_counts
+                if j != num_chunks - 1: 
+                    attention_matrix = attention_matrix[:, :, :, : past_cache_len + current_seq_length].clone().contiguous()
+                else:
+                    attention_matrix = attention_matrix[:, :, :, : past_cache_len + current_seq_length + len_question].clone().contiguous()
+                attention_matrix = torch.sum(attention_matrix, dim=-2)
+                attention_matrix = attention_matrix.view(
+                    attention_matrix.size(0), model.config.num_key_value_heads, num_repeats, -1
+                ).sum(dim=2)
+                full_context_size = attention_matrix.size(-1)
+                attention_matrix[..., :sink_tokens] = float("inf")
+                if j == num_chunks - 1:
+                    attention_matrix[..., -len_question:] = float("inf")
+                if j == 0:
+                    k = int(sink_tokens + (max(0, current_seq_length - sink_tokens) // compression_factor))
+                    k = min(k + past_cache_len, full_context_size)
+                elif j < num_chunks - 1:
+                    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
+                    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)
+        for hook in hooks:
+            hook.remove()
+        return cache
+    except Exception as e:
+        raise RuntimeError(f"Failed to compress KV cache: {e}")
 
 def run_naive_rag_query(collection_name, query, rag_token_size, prefix, task, few_shot_examples):
     k = max(1, rag_token_size // 256)
@@ -443,7 +493,8 @@ 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] + ".pt"
+    current_timestamp = int(time.time())
+    cache_name = f"cache_{current_timestamp}_{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)
@@ -455,7 +506,8 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
             rag_text = combined_text[len(prefix):]
         else:
             rag_text = combined_text
-        collection_name = "default_collection_" + uuid.uuid4().hex[:6]
+        current_timestamp = int(time.time())
+        collection_name = f"default_{current_timestamp}_{uuid.uuid4().hex[:6]}"
         rag_index = create_rag_index(collection_name, rag_text)
     state.update({
         "compressed_cache": save_path,
@@ -469,7 +521,7 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         "retrieval_slider": retrieval_context_length,
         "prefill_time": time.perf_counter() - start_time_prefill,
         "compression_done": True,
-        "tokens_breakdown": f"KV Compress Tokens: {kv_tokens}, RAG Tokens: {rag_tokens}",
+        "tokens_breakdown": f"RAG tokens: {rag_tokens} (for retrieval), {kv_tokens} tokens (for KV compression)",
         "chat_feedback": "Document compressed successfully. You can now chat."
     })
     return state, True
@@ -530,20 +582,14 @@ def chat_response_stream(message: str, history: list, state: dict):
         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."
+def update_token_breakdown(token_count, retrieval_slider, global_local_value):
+    retrieval_context_length = int(token_count / retrieval_slider)
+    percentage = int(global_local_value.replace('%', ''))
+    rag_tokens = int(retrieval_context_length * (1.0 - (percentage / 100)))
+    kv_tokens = retrieval_context_length - rag_tokens
+    return f"Token Breakdown: {rag_tokens} tokens will be used for RAG retrieval, and {kv_tokens} tokens for KV compression."
 
 ##########################################################################
 # Gradio Interface
@@ -629,6 +675,9 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
     compression_done = gr.State(value=False)
     compressed_doc_state = gr.State(value={})
     
+    def toggle_chat_interactivity(compression_done):
+        return gr.update(interactive=compression_done)
+    
     with gr.Row(elem_classes="main-container"):
         with gr.Column(elem_classes="upload-section"):
             gr.Markdown("## Document Preprocessing")
@@ -646,40 +695,38 @@ 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.")
             
+            # When document parameters change, disable the chat interface.
             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, chat_status_text]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text, gr.State().update(interactive=False)]
             )
             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, chat_status_text]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text, gr.State().update(interactive=False)]
             )
             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, chat_status_text]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text, gr.State().update(interactive=False)]
             )
             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, chat_status_text]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text, gr.State().update(interactive=False)]
             )
             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],
@@ -697,6 +744,9 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
             ).then(
                 fn=lambda state: gr.update(value="Document compressed successfully. You can now chat."),
                 outputs=chat_status_text
+            ).then(
+                fn=lambda: gr.update(interactive=True),
+                outputs=lambda: chat_interface  # Re-enable chat interface after successful compression.
             )
         
         with gr.Column(elem_classes="chatbot-container"):
@@ -704,7 +754,8 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
             chat_interface = gr.ChatInterface(
                 fn=chat_response_stream,
                 additional_inputs=[compressed_doc_state],
-                type="messages"
+                type="messages",
+                interactive=False
             )
 
 demo.queue(max_size=16).launch()