updat

app.py

@@ -68,8 +68,6 @@ 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
@@ -227,8 +225,7 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
             gr.update(interactive=False),
             False,
             {},
-            chat_status,
-            gr.update(interactive=False)  # Disable chat interface
+            chat_status
         )
     print("Converting to markdown")
     try:
@@ -243,8 +240,7 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
             gr.update(interactive=False),
             False,
             {},
-            chat_status,
-            gr.update(interactive=False)  # Disable chat interface on error
+            chat_status
         )
 
     print("Done")
@@ -254,7 +250,7 @@ def auto_convert(file_objs, url, do_ocr, do_table_structure):
     print("Done")
     min_ratio = min(suggestions)
     max_ratio = max(suggestions)
-    default_ratio = suggestions[len(suggestions) // 2]
+    default_ratio = 6
     retrieval_tokens = int(token_count / default_ratio)
     token_count_str = f"Number of tokens before compression: {token_count}"
     retrieval_str = f"Number of tokens after compression: {retrieval_tokens}"
@@ -277,8 +273,7 @@ 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,
-        gr.update(interactive=False)  # Ensure chat remains disabled until compression
+        chat_status
     )
 
 def get_compressed_kv_cache(sink_tokens, step_size, target_token_size, context_ids, context_attention_mask, question_ids, question_attention_mask):
@@ -458,8 +453,11 @@ def run_naive_rag_query(collection_name, query, rag_token_size, prefix, task, fe
     return rag_context
 
 @spaces.GPU
-def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_local_value, task_description, few_shot, state):
+def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_local_value, task_description, few_shot, state, progress=gr.Progress()):
+    progress(0, desc="Starting compression process")
+    
     percentage = int(global_local_value.replace('%', ''))
+    progress(0.1, desc="Tokenizing text and preparing task")
     question_text = task_description + "\n" + few_shot
     context_encoding = tokenizer(combined_text, return_tensors="pt").to(device)
     question_encoding = tokenizer(question_text, return_tensors="pt").to(device)
@@ -467,41 +465,44 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
     context_attention_mask = context_encoding["attention_mask"]
     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}")
+    progress(0.2, desc=f"Token breakdown computed: {kv_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)
+        progress(0.3, desc="Starting KV cache compression")
         step_size = 2
-        start_time_prefill = time.perf_counter()
         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:
+            progress(1, desc="Compression failed")
             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)
+        progress(0.6, desc="KV cache compression completed")
     else:
-        start_time_prefill = 0
         target_token_size = 0
         past_key_values = FinchCache()
         compressed_length = past_key_values.get_seq_length()
+        progress(0.3, desc="Skipping compression as percentage is 0")
+    
     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)
     past_key_values.save(save_path)
+    progress(0.8, desc="Cache saved successfully")
+    
     collection_name = state.get("rag_index", None)
     if collection_name is None:
-        print("Collection name not found creating a new one.")
+        print("Collection name not found; creating a new one.")
         if combined_text.startswith(prefix):
             rag_text = combined_text[len(prefix):]
         else:
@@ -509,27 +510,23 @@ def prepare_compression_and_rag(combined_text, retrieval_slider_value, global_lo
         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,
-        "compressed_length": compressed_length,
         "rag_index": collection_name,
-        "target_token_size": target_token_size,
         "global_local": percentage,
-        "combined_text": combined_text,
         "task_description": task_description,
         "few_shot": few_shot,
         "retrieval_slider": retrieval_context_length,
-        "prefill_time": time.perf_counter() - start_time_prefill,
-        "compression_done": True,
-        "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
+    progress(1, desc="Compression complete")
+    return state, "Document compressed successfully. You can now chat.", True
+
 
 @spaces.GPU
-def chat_response_stream(message: str, history: list, state: dict):
+def chat_response_stream(message: str, history: list, state: dict, compression_done: bool):
     # Check if the document is compressed before allowing chat
-    if not state.get("compression_done", False) or "compressed_cache" not in state:
+    if not compression_done or "compressed_cache" not in state:
         yield "Document not compressed yet. Please compress the document first to enable chat."
         return
     user_message = message
@@ -589,7 +586,7 @@ def update_token_breakdown(token_count, retrieval_slider, global_local_value):
     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."
+    return f"Token Breakdown: {kv_tokens} tokens (KV compression), {rag_tokens} tokens (RAG retrieval)"
 
 ##########################################################################
 # Gradio Interface
@@ -651,33 +648,22 @@ body {
     display: flex;
     align-items: center;
 }
-@media (min-width: 768px) {
-    .main-container {
-        display: flex;
-        justify-content: space-between;
-        gap: 20px;
-    }
-    .upload-section {
-        flex: 3;
-    }
-    .chatbot-container {
-        flex: 1;
-        margin-top: 0;
-    }
-}
+
 """
+def reset_chat_state():
+    return gr.update(value="Document not compressed yet. Please compress the document to enable chat."), False
 
 with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
     gr.HTML("<h1><center>Beyond RAG with LLama 3.1-8B-Instruct Model</center></h1>")
     gr.HTML("<center><p>Compress your document and chat with it.</p></center>")
     
+    # Define chat_status_text as a Textbox with a set elem_id for custom styling.
+    chat_status_text = gr.Textbox(value="Document not compressed yet. Please compress the document to enable chat.", interactive=False, show_label=False, render=False, lines=5)
+    
     hidden_token_count = gr.State(value=0)
     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")
@@ -696,37 +682,90 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
             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: ")
             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%")
+            global_local_slider = gr.Radio(label="Hybrid Retrieval (0 is all RAG, 100 is all global)", 
+                                           choices=["0%", "25%", "50%", "75%", "100%"], value="100%")
             compress_button = gr.Button("Compress Document", interactive=False, elem_classes="upload-button")
-            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: Run auto_convert then chain reset_chat_state.
             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, gr.State().update(interactive=False)]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, 
+                         hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
+            ).then(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
             )
+            
+            # URL input: Run auto_convert then chain reset_chat_state.
             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, gr.State().update(interactive=False)]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, 
+                         hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
+            ).then(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
             )
+            
+            # OCR checkbox: Run auto_convert then chain reset_chat_state.
             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, gr.State().update(interactive=False)]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, 
+                         hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
+            ).then(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
             )
+            
+            # Table structure checkbox: Run auto_convert then chain reset_chat_state.
             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, gr.State().update(interactive=False)]
+                outputs=[markdown_output, token_count_text, retrieval_slider, retrieval_info_text, 
+                         hidden_token_count, compress_button, compression_done, compressed_doc_state, chat_status_text]
+            ).then(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
+            )
+            
+            # Reset chat state when prompt designer fields change.
+            task_description_input.change(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
+            )
+            few_shot_input.change(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
+            )
+            
+            # Reset chat state when the Markdown output changes.
+            markdown_output.change(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
             )
+            
+            # When sliders change, reset chat state.
             retrieval_slider.change(
-                fn=update_retrieval_context,
-                inputs=[hidden_token_count, retrieval_slider],
-                outputs=retrieval_info_text
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
+            )
+            global_local_slider.change(
+                fn=reset_chat_state,
+                inputs=None,
+                outputs=[chat_status_text, compression_done]
             )
+            
+            # Update token breakdown when sliders change.
             retrieval_slider.change(
                 fn=update_token_breakdown,
                 inputs=[hidden_token_count, retrieval_slider, global_local_slider],
@@ -737,25 +776,22 @@ with gr.Blocks(css=CSS, theme=gr.themes.Soft()) as demo:
                 inputs=[hidden_token_count, retrieval_slider, global_local_slider],
                 outputs=tokens_breakdown_text
             )
+            
+            # Compress button: Prepare compression and then update chat status.
             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
-            ).then(
-                fn=lambda: gr.update(interactive=True),
-                outputs=lambda: chat_interface  # Re-enable chat interface after successful compression.
+                outputs=[compressed_doc_state, chat_status_text, compression_done]
             )
-        
         with gr.Column(elem_classes="chatbot-container"):
+            chat_status_text.render()
             gr.Markdown("## Chat")
             chat_interface = gr.ChatInterface(
                 fn=chat_response_stream,
-                additional_inputs=[compressed_doc_state],
+                additional_inputs=[compressed_doc_state, compression_done],
                 type="messages",
-                interactive=False
+                fill_height=True 
             )
 
-demo.queue(max_size=16).launch()
+    demo.queue(max_size=16).launch()
+

cache.py

@@ -1,75 +0,0 @@
-from transformers import DynamicCache
-import torch
-import os
-
-class FinchCache(DynamicCache):
-    def __init__(self) -> None:
-        super().__init__()
-        self.key_cache = []
-        self.value_cache = []
-
-    @staticmethod
-    def _rotate_half(x):
-        x1 = x[..., : x.shape[-1] // 2]
-        x2 = x[..., x.shape[-1] // 2 :]
-        return torch.cat((-x2, x1), dim=-1)
-
-    def _apply_key_rotary_pos_emb(self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-        return (key_states * cos) + (self._rotate_half(key_states) * sin)
-
-    @staticmethod
-    def _rerotate_cos_sin(x, inv_freq, important_pos_batch):
-        B, H, L = important_pos_batch.shape
-        device = important_pos_batch.device
-        device_type = x.device.type
-        dtype = x.dtype
-        idx = torch.arange(0, L, device=device)
-        idx = idx.unsqueeze(0)
-        inv_freq = inv_freq[None, None, :, None].float().expand(B, H, -1, 1) # (B, H, M, 1)
-        idx = idx[:, None, :].float().expand(B, H, L) # (B, H, L)
-        delta_pos =  idx - important_pos_batch
-        delta_pos = delta_pos.unsqueeze(2) # (B, H, 1, L)
-
-        device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
-
-        with torch.autocast(device_type=device_type, enabled=False):
-            freqs = delta_pos.float() * inv_freq.float()
-            freqs = freqs.transpose(2, 3)
-            emb = torch.cat((freqs, freqs), dim=-1)
-            cos = emb.cos().contiguous()
-            sin = emb.sin().contiguous()
-        return cos.to(dtype=dtype), sin.to(dtype=dtype)
-
-    @staticmethod
-    def gather_important_tokens(states, indices):
-        return torch.gather(states, 2, indices.unsqueeze(-1).expand(-1, -1, -1, states.size(3))).contiguous()
-
-    def compress_cache(self, layer_index, important_pos, inv_freq):
-        new_length = important_pos.size(2)
-        new_cos, new_sin = self._rerotate_cos_sin(self.key_cache[layer_index], inv_freq, important_pos)
-        gathered_keys = self.gather_important_tokens(self.key_cache[layer_index], important_pos).clone()
-        self.key_cache[layer_index] = self._apply_key_rotary_pos_emb(gathered_keys, new_cos, new_sin)
-        gathered_values = self.gather_important_tokens(self.value_cache[layer_index], important_pos).clone()
-        self.value_cache[layer_index] = gathered_values
-        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(
-                {"key_cache": [k.cpu() for k in self.key_cache], "value_cache": [v.cpu() for v in self.value_cache]},
-                path,
-            )
-        except Exception as e:
-            print(f"Error occurred while saving: {e}")
-
-    @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"]]
-        cache.value_cache = [v.to(device) for v in data["value_cache"]]
-        cache._seen_tokens = cache.value_cache[0].size(2) if cache.value_cache else 0
-        return cache

global_compression.py

@@ -1,211 +0,0 @@
-import math
-import torch
-from cache import FinchCache
-from utils import repeat_kv
-from transformers.models.llama.modeling_llama import rotate_half
-import spaces
-
-@spaces.GPU
-def get_compressed_kv_cache(model, 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:
-        # 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
-        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
-
-    # (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:
-            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()
-            )
-            # 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.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(
-            [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)
-        # 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]
-            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)
-            # 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)
-            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# TODO remember to include the question tokens
-                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
-

preprocess_document.py

@@ -1,34 +0,0 @@
-from langchain_docling import DoclingLoader
-from langchain_docling.loader import ExportType
-
-# Import required classes for building a custom converter
-from docling.document_converter import DocumentConverter, PdfFormatOption, InputFormat
-from docling.datamodel.pipeline_options import PdfPipelineOptions
-from docling.backend.pypdfium2_backend import PyPdfiumDocumentBackend
-import spaces
-
-@spaces.GPU
-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()
-    pipeline_options.do_ocr = do_ocr
-    pipeline_options.do_table_structure = do_table_structure
-    pdf_format_options = PdfFormatOption(
-        pipeline_options=pipeline_options,
-        backend=PyPdfiumDocumentBackend,
-    )
-    doc_converter = DocumentConverter(
-        allowed_formats=[InputFormat.PDF],
-        format_options={
-            InputFormat.PDF: pdf_format_options
-        }
-    )
-
-    # Pass the custom converter to the DoclingLoader.
-    loader = DoclingLoader(
-        file_path=file_path,
-        export_type=ExportType.MARKDOWN,
-        converter=doc_converter
-    )
-    docs = loader.load()
-    return docs[0].page_content

rag.py

@@ -1,53 +0,0 @@
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.schema.document import Document
-from langchain_community.embeddings import HuggingFaceBgeEmbeddings
-from langchain_chroma import Chroma
-import spaces
-from langchain_text_splitters import MarkdownHeaderTextSplitter
-import os
-from transformers import AutoTokenizer
-api_token = os.getenv("HF_TOKEN")
-model_name = "meta-llama/Llama-3.1-8B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name, token=api_token)
-    
-embedding_model = HuggingFaceBgeEmbeddings(
-        model_name="BAAI/bge-large-en-v1.5",
-        model_kwargs={"device": "cuda"},
-        encode_kwargs={"normalize_embeddings": True},
-        query_instruction=""
-    )
-
-
-def create_rag_index(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.
-    docs = [Document(page_content=x) for x in text_splitter.split_text(text_no_prefix)]
-    
-    vectorstore = Chroma.from_documents(documents=docs, embedding=embedding_model)
-    return vectorstore
-
-def run_naive_rag_query(vectorstore, 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})
-    retrieved_docs = retriever.invoke(query)
-    for doc in retrieved_docs:
-        print("=================")
-        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