Spaces:
Running
on
Zero
Running
on
Zero
| # Acknowledgement: This demo code is adapted from the original Hugging Face Space "ContextCite" | |
| # (https://huggingface.co/spaces/contextcite/context-cite). | |
| import os | |
| from enum import Enum | |
| from dataclasses import dataclass | |
| from typing import Dict, List, Any, Optional | |
| import gradio as gr | |
| import numpy as np | |
| import spaces | |
| import nltk | |
| import base64 | |
| import traceback | |
| from src.utils import split_into_sentences as split_into_sentences_utils | |
| # --- AttnTrace imports (from app_full.py) --- | |
| from src.models import create_model | |
| from src.attribution import AttnTraceAttribution | |
| from src.prompts import wrap_prompt | |
| from gradio_highlightedtextbox import HighlightedTextbox | |
| from examples import run_example_1, run_example_2, run_example_3, run_example_4, run_example_5, run_example_6 | |
| from functools import partial | |
| # Load original app constants | |
| APP_TITLE = '<div class="app-title"><span class="brand">AttnTrace: </span><span class="subtitle">Attention-based Context Traceback for Long-Context LLMs</span></div>' | |
| APP_DESCRIPTION = """AttnTrace traces a model's generated statements back to specific parts of the context using attention-based traceback. Try it out with Meta-Llama-3.1-8B-Instruct here! See the [[paper](https://arxiv.org/abs/2506.04202)] and [[code](https://github.com/Wang-Yanting/TracLLM-Kit)] for more! | |
| Maintained by the AttnTrace team.""" | |
| # NEW_TEXT = """Long-context large language models (LLMs), such as Gemini-2.5-Pro and Claude-Sonnet-4, are increasingly used to empower advanced AI systems, including retrieval-augmented generation (RAG) pipelines and autonomous agents. In these systems, an LLM receives an instruction along with a context—often consisting of texts retrieved from a knowledge database or memory—and generates a response that is contextually grounded by following the instruction. Recent studies have designed solutions to trace back to a subset of texts in the context that contributes most to the response generated by the LLM. These solutions have numerous real-world applications, including performing post-attack forensic analysis and improving the interpretability and trustworthiness of LLM outputs. While significant efforts have been made, state-of-the-art solutions such as TracLLM often lead to a high computation cost, e.g., it takes TracLLM hundreds of seconds to perform traceback for a single response-context pair. In this work, we propose {\name}, a new context traceback method based on the attention weights produced by an LLM for a prompt. To effectively utilize attention weights, we introduce two techniques designed to enhance the effectiveness of {\name}, and we provide theoretical insights for our design choice. %Moreover, we perform both theoretical analysis and empirical evaluation to demonstrate their effectiveness. | |
| # We also perform a systematic evaluation for {\name}. The results demonstrate that {\name} is more accurate and efficient than existing state-of-the-art context traceback methods. We also show {\name} can improve state-of-the-art methods in detecting prompt injection under long contexts through the attribution-before-detection paradigm. As a real-world application, we demonstrate that {\name} can effectively pinpoint injected instructions in a paper designed to manipulate LLM-generated reviews. | |
| # The code and data will be open-sourced. """ | |
| # EDIT_TEXT = "Feel free to edit!" | |
| GENERATE_CONTEXT_TOO_LONG_TEXT = ( | |
| '<em style="color: red;">Context is too long for the current model.</em>' | |
| ) | |
| ATTRIBUTE_CONTEXT_TOO_LONG_TEXT = '<em style="color: red;">Context is too long for the current traceback method.</em>' | |
| CONTEXT_LINES = 20 | |
| CONTEXT_MAX_LINES = 40 | |
| SELECTION_DEFAULT_TEXT = "Click on a sentence in the response to traceback!" | |
| SELECTION_DEFAULT_VALUE = [(SELECTION_DEFAULT_TEXT, None)] | |
| SOURCES_INFO = 'These are the texts that contribute most to the response.' | |
| # SOURCES_IN_CONTEXT_INFO = ( | |
| # "This shows the important sentences highlighted within their surrounding context from the text above. Colors indicate ranking: Red (1st), Orange (2nd), Golden (3rd), Yellow (4th-5th), Light (6th+)." | |
| # ) | |
| MODEL_PATHS = [ | |
| "meta-llama/Meta-Llama-3.1-8B-Instruct", | |
| ] | |
| MAX_TOKENS = { | |
| "meta-llama/Meta-Llama-3.1-8B-Instruct": 131072, | |
| } | |
| DEFAULT_MODEL_PATH = MODEL_PATHS[0] | |
| EXPLANATION_LEVELS = ["sentence", "paragraph", "text segment"] | |
| DEFAULT_EXPLANATION_LEVEL = "sentence" | |
| class WorkflowState(Enum): | |
| WAITING_TO_GENERATE = 0 | |
| WAITING_TO_SELECT = 1 | |
| READY_TO_ATTRIBUTE = 2 | |
| class State: | |
| workflow_state: WorkflowState | |
| context: str | |
| query: str | |
| response: str | |
| start_index: int | |
| end_index: int | |
| scores: np.ndarray | |
| answer: str | |
| highlighted_context: str | |
| full_response: str | |
| explained_response_part: str | |
| last_query_used: str = "" | |
| # --- Dynamic Model and Attribution Management --- | |
| current_llm = None | |
| current_attr = None | |
| current_model_path = None | |
| current_explanation_level = None | |
| current_api_key = None | |
| def initialize_model_and_attr(): | |
| """Initialize model and attribution with default configuration""" | |
| global current_llm, current_attr, current_model_path, current_explanation_level, current_api_key | |
| try: | |
| # Check if we need to reinitialize the model | |
| need_model_update = (current_llm is None or | |
| current_model_path != DEFAULT_MODEL_PATH or | |
| current_api_key != os.getenv("HF_TOKEN")) | |
| # Check if we need to update attribution | |
| need_attr_update = (current_attr is None or | |
| current_explanation_level != DEFAULT_EXPLANATION_LEVEL or | |
| need_model_update) | |
| if need_model_update: | |
| print(f"Initializing model: {DEFAULT_MODEL_PATH}") | |
| effective_api_key = os.getenv("HF_TOKEN") | |
| current_llm = create_model(model_path=DEFAULT_MODEL_PATH, api_key=effective_api_key, device="cuda") | |
| current_model_path = DEFAULT_MODEL_PATH | |
| current_api_key = effective_api_key | |
| if need_attr_update: | |
| print(f"Initializing context traceback with explanation level: {DEFAULT_EXPLANATION_LEVEL}") | |
| current_attr = AttnTraceAttribution( | |
| current_llm, | |
| explanation_level=DEFAULT_EXPLANATION_LEVEL, | |
| K=3, | |
| q=0.4, | |
| B=30 | |
| ) | |
| current_explanation_level = DEFAULT_EXPLANATION_LEVEL | |
| return current_llm, current_attr, None | |
| except Exception as e: | |
| error_msg = f"Error initializing model/traceback: {str(e)}" | |
| print(error_msg) | |
| traceback.print_exc() | |
| return None, None, error_msg | |
| # Initialize with defaults | |
| initialize_model_and_attr() | |
| # Images replaced with CSS textures and gradients - no longer needed | |
| def clear_state(): | |
| return State( | |
| workflow_state=WorkflowState.WAITING_TO_GENERATE, | |
| context="", | |
| query="", | |
| response="", | |
| start_index=0, | |
| end_index=0, | |
| scores=np.array([]), | |
| answer="", | |
| highlighted_context="", | |
| full_response="", | |
| explained_response_part="", | |
| last_query_used="" | |
| ) | |
| def load_an_example(example_loader_func, state: State): | |
| context, query = example_loader_func() | |
| # Update both UI and state | |
| state.context = context | |
| state.query = query | |
| state.workflow_state = WorkflowState.WAITING_TO_GENERATE | |
| # Clear previous results | |
| state.response = "" | |
| state.answer = "" | |
| state.full_response = "" | |
| state.explained_response_part = "" | |
| print(f"Loaded example - Context: {len(context)} chars, Query: {query[:50]}...") | |
| return ( | |
| context, # basic_context_box | |
| query, # basic_query_box | |
| state, | |
| "", # response_input_box - clear it | |
| gr.update(value=[("Click the 'Generate/Use Response' button above to see response text here for traceback analysis.", None)]), # basic_response_box - keep visible | |
| gr.update(selected=0) # basic_context_tabs - switch to first tab | |
| ) | |
| def get_max_tokens(model_path: str): | |
| return MAX_TOKENS.get(model_path, 2048) # Default fallback | |
| def get_scroll_js_code(elem_id): | |
| return f""" | |
| function scrollToElement() {{ | |
| const element = document.getElementById("{elem_id}"); | |
| element.scrollIntoView({{ behavior: "smooth", block: "nearest" }}); | |
| }} | |
| """ | |
| def basic_update(context: str, query: str, state: State): | |
| state.context = context | |
| state.query = query | |
| state.workflow_state = WorkflowState.WAITING_TO_GENERATE | |
| return ( | |
| gr.update(value=[("Click the 'Generate/Use Response' button above to see response text here for traceback analysis.", None)]), # basic_response_box - keep visible | |
| gr.update(selected=0), # basic_context_tabs - switch to first tab | |
| state, | |
| ) | |
| def generate_model_response(state: State): | |
| # Validate inputs first with debug info | |
| print(f"Validation - Context length: {len(state.context) if state.context else 0}") | |
| print(f"Validation - Query: {state.query[:50] if state.query else 'empty'}...") | |
| if not state.context or not state.context.strip(): | |
| print("❌ Validation failed: No context") | |
| return state, gr.update(value=[("❌ Please enter context before generating response! If you just changed configuration, try reloading an example.", None)], visible=True) | |
| if not state.query or not state.query.strip(): | |
| print("❌ Validation failed: No query") | |
| return state, gr.update(value=[("❌ Please enter a query before generating response! If you just changed configuration, try reloading an example.", None)], visible=True) | |
| # Initialize model and attribution with default configuration | |
| print(f"🔧 Generating response with explanation_level: {DEFAULT_EXPLANATION_LEVEL}") | |
| llm, attr, error_msg = initialize_model_and_attr() | |
| if llm is None or attr is None: | |
| error_text = error_msg if error_msg else "Model initialization failed!" | |
| return state, gr.update(value=[(f"❌ {error_text}", None)], visible=True) | |
| prompt = wrap_prompt(state.query, [state.context]) | |
| print(f"Generated prompt for {DEFAULT_MODEL_PATH}: {prompt[:200]}...") # Debug log | |
| # Check context length | |
| if len(prompt.split()) > get_max_tokens(DEFAULT_MODEL_PATH) - 512: | |
| return state, gr.update(value=[(GENERATE_CONTEXT_TOO_LONG_TEXT, None)], visible=True) | |
| answer = llm.query(prompt) | |
| print(f"Model response: {answer}") # Debug log | |
| state.response = answer | |
| state.answer = answer | |
| state.full_response = answer | |
| state.workflow_state = WorkflowState.WAITING_TO_SELECT | |
| return state, gr.update(visible=False) | |
| def split_into_sentences(text: str): | |
| lines = text.splitlines() | |
| sentences = [] | |
| for line in lines: | |
| sentences.extend(nltk.sent_tokenize(line)) | |
| separators = [] | |
| cur_start = 0 | |
| for sentence in sentences: | |
| cur_end = text.find(sentence, cur_start) | |
| separators.append(text[cur_start:cur_end]) | |
| cur_start = cur_end + len(sentence) | |
| return sentences, separators | |
| def basic_highlight_response( | |
| response: str, selected_index: int, num_sources: int = -1 | |
| ): | |
| sentences, separators = split_into_sentences(response) | |
| ht = [] | |
| if num_sources == -1: | |
| citations_text = "Traceback!" | |
| elif num_sources == 0: | |
| citations_text = "No important text!" | |
| else: | |
| citations_text = f"[{','.join(str(i) for i in range(1, num_sources + 1))}]" | |
| for i, (sentence, separator) in enumerate(zip(sentences, separators)): | |
| label = citations_text if i == selected_index else "Traceback" | |
| # Hack to ignore punctuation | |
| if len(sentence) >= 4: | |
| ht.append((separator + sentence, label)) | |
| else: | |
| ht.append((separator + sentence, None)) | |
| color_map = {"Click to cite!": "blue", citations_text: "yellow"} | |
| return gr.HighlightedText(value=ht, color_map=color_map) | |
| def basic_highlight_response_with_visibility( | |
| response: str, selected_index: int, num_sources: int = -1, visible: bool = True | |
| ): | |
| """Version of basic_highlight_response that also sets visibility""" | |
| sentences, separators = split_into_sentences(response) | |
| ht = [] | |
| if num_sources == -1: | |
| citations_text = "Traceback!" | |
| elif num_sources == 0: | |
| citations_text = "No important text!" | |
| else: | |
| citations_text = f"[{','.join(str(i) for i in range(1, num_sources + 1))}]" | |
| for i, (sentence, separator) in enumerate(zip(sentences, separators)): | |
| label = citations_text if i == selected_index else "Traceback" | |
| # Hack to ignore punctuation | |
| if len(sentence) >= 4: | |
| ht.append((separator + sentence, label)) | |
| else: | |
| ht.append((separator + sentence, None)) | |
| color_map = {"Click to cite!": "blue", citations_text: "yellow"} | |
| return gr.update(value=ht, color_map=color_map, visible=visible) | |
| def basic_update_highlighted_response(evt: gr.SelectData, state: State): | |
| response_update = basic_highlight_response(state.response, evt.index) | |
| return response_update, state | |
| def unified_response_handler(response_text: str, state: State): | |
| """Handle both LLM generation and manual input based on whether text is provided""" | |
| # Check if instruction has changed from what was used to generate current response | |
| instruction_changed = hasattr(state, 'last_query_used') and state.last_query_used != state.query | |
| # If response_text is empty, whitespace, or instruction changed, generate from LLM | |
| if not response_text or not response_text.strip() or instruction_changed: | |
| if instruction_changed: | |
| print("📝 Instruction changed, generating new response from LLM...") | |
| else: | |
| print("🤖 Generating response from LLM...") | |
| # Validate inputs first | |
| if not state.context or not state.context.strip(): | |
| return ( | |
| state, | |
| response_text, # Keep current text box content | |
| gr.update(visible=False), # Keep response box hidden | |
| gr.update(value=[("❌ Please enter context before generating response!", None)], visible=True) | |
| ) | |
| if not state.query or not state.query.strip(): | |
| return ( | |
| state, | |
| response_text, # Keep current text box content | |
| gr.update(visible=False), # Keep response box hidden | |
| gr.update(value=[("❌ Please enter a query before generating response!", None)], visible=True) | |
| ) | |
| # Initialize model and generate response | |
| llm, attr, error_msg = initialize_model_and_attr() | |
| if llm is None: | |
| error_text = error_msg if error_msg else "Model initialization failed!" | |
| return ( | |
| state, | |
| response_text, # Keep current text box content | |
| gr.update(visible=False), # Keep response box hidden | |
| gr.update(value=[(f"❌ {error_text}", None)], visible=True) | |
| ) | |
| prompt = wrap_prompt(state.query, [state.context]) | |
| # Check context length | |
| if len(prompt.split()) > get_max_tokens(DEFAULT_MODEL_PATH) - 512: | |
| return ( | |
| state, | |
| response_text, # Keep current text box content | |
| gr.update(visible=False), # Keep response box hidden | |
| gr.update(value=[(GENERATE_CONTEXT_TOO_LONG_TEXT, None)], visible=True) | |
| ) | |
| # Generate response | |
| answer = llm.query(prompt) | |
| print(f"Generated response: {answer[:100]}...") | |
| # Update state and UI | |
| state.response = answer | |
| state.answer = answer | |
| state.full_response = answer | |
| state.last_query_used = state.query # Track which query was used for this response | |
| state.workflow_state = WorkflowState.WAITING_TO_SELECT | |
| # Create highlighted response and show it | |
| response_update = basic_highlight_response_with_visibility(state.response, -1, visible=True) | |
| return ( | |
| state, | |
| answer, # Put generated response in text box | |
| response_update, # Update clickable response content | |
| gr.update(visible=False) # Hide error box | |
| ) | |
| else: | |
| # Use provided text as manual response | |
| print("✏️ Using manual response...") | |
| manual_text = response_text.strip() | |
| # Update state with manual response | |
| state.response = manual_text | |
| state.answer = manual_text | |
| state.full_response = manual_text | |
| state.last_query_used = state.query # Track current query for this response | |
| state.workflow_state = WorkflowState.WAITING_TO_SELECT | |
| # Create highlighted response for selection | |
| response_update = basic_highlight_response_with_visibility(state.response, -1, visible=True) | |
| return ( | |
| state, | |
| manual_text, # Keep text in text box | |
| response_update, # Update clickable response content | |
| gr.update(visible=False) # Hide error box | |
| ) | |
| def get_color_by_rank(rank, total_items): | |
| """Get color based purely on rank position for better visual distinction""" | |
| if total_items == 0: | |
| return "#F0F0F0", "rgba(240, 240, 240, 0.8)" | |
| # Pure ranking-based color assignment for clear visual hierarchy | |
| if rank == 1: # Highest importance - Strong Red | |
| bg_color = "#FF4444" # Bright red | |
| rgba_color = "rgba(255, 68, 68, 0.9)" | |
| elif rank == 2: # Second highest - Orange | |
| bg_color = "#FF8C42" # Bright orange | |
| rgba_color = "rgba(255, 140, 66, 0.8)" | |
| elif rank == 3: # Third highest - Golden Yellow | |
| bg_color = "#FFD93D" # Golden yellow | |
| rgba_color = "rgba(255, 217, 61, 0.8)" | |
| elif rank <= 5: # 4th-5th - Light Yellow | |
| bg_color = "#FFF280" # Standard yellow | |
| rgba_color = "rgba(255, 242, 128, 0.7)" | |
| else: # Lower importance - Very Light Yellow | |
| bg_color = "#FFF9C4" # Very light yellow | |
| rgba_color = "rgba(255, 249, 196, 0.6)" | |
| return bg_color, rgba_color | |
| def basic_get_scores_and_sources_full_response(state: State): | |
| """Traceback the entire response instead of a selected segment""" | |
| # Use the entire response as the explained part | |
| state.explained_response_part = state.full_response | |
| # Attribution using default configuration | |
| _, attr, error_msg = initialize_model_and_attr() | |
| if attr is None: | |
| error_text = error_msg if error_msg else "Traceback initialization failed!" | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[(f"❌ {error_text}", None)], visible=True), | |
| state, | |
| ) | |
| try: | |
| # Validate attribution inputs | |
| if not state.context or not state.context.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No context available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| if not state.query or not state.query.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No query available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| if not state.full_response or not state.full_response.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No response available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| print(f"start full response traceback with explanation_level: {DEFAULT_EXPLANATION_LEVEL}") | |
| print(f"context length: {len(state.context)}, query: {state.query[:100]}...") | |
| print(f"full response: {state.full_response[:100]}...") | |
| print(f"tracing entire response (length: {len(state.full_response)} chars)") | |
| texts, important_ids, importance_scores, _, _ = attr.attribute( | |
| state.query, [state.context], state.full_response, state.full_response | |
| ) | |
| print("end full response traceback") | |
| print(f"explanation_level: {DEFAULT_EXPLANATION_LEVEL}") | |
| print(f"texts count: {len(texts)} (how context was segmented)") | |
| if len(texts) > 0: | |
| print(f"sample text segments: {[text[:50] + '...' if len(text) > 50 else text for text in texts[:3]]}") | |
| print(f"important_ids: {important_ids}") | |
| print("importance_scores: ", importance_scores) | |
| if not importance_scores: | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No traceback scores generated for full response!", None)], visible=True), | |
| state, | |
| ) | |
| state.scores = np.array(importance_scores) | |
| # Highlighted sources with ranking-based colors | |
| highlighted_text = [] | |
| sorted_indices = np.argsort(state.scores)[::-1] | |
| total_sources = len(important_ids) | |
| for rank, i in enumerate(sorted_indices): | |
| source_text = texts[important_ids[i]] | |
| _ = get_color_by_rank(rank + 1, total_sources) | |
| highlighted_text.append( | |
| ( | |
| source_text, | |
| f"rank_{rank+1}", | |
| ) | |
| ) | |
| # In-context highlights with ranking-based colors - show ALL text | |
| in_context_highlighted_text = [] | |
| ranks = {important_ids[i]: rank for rank, i in enumerate(sorted_indices)} | |
| for i in range(len(texts)): | |
| source_text = texts[i] | |
| # Skip or don't highlight segments that are only newlines or whitespace | |
| if source_text.strip() == "": | |
| # For whitespace-only segments, add them without highlighting | |
| in_context_highlighted_text.append((source_text, None)) | |
| elif i in important_ids: | |
| # Only highlight if the segment has actual content (not just newlines) | |
| if source_text.strip(): # Has non-whitespace content | |
| rank = ranks[i] + 1 | |
| # Split the segment to separate leading/trailing newlines from content | |
| # This prevents newlines from being highlighted | |
| leading_whitespace = "" | |
| trailing_whitespace = "" | |
| content = source_text | |
| # Extract leading newlines/whitespace | |
| while content and content[0] in ['\n', '\r', '\t', ' ']: | |
| leading_whitespace += content[0] | |
| content = content[1:] | |
| # Extract trailing newlines/whitespace | |
| while content and content[-1] in ['\n', '\r', '\t', ' ']: | |
| trailing_whitespace = content[-1] + trailing_whitespace | |
| content = content[:-1] | |
| # Add the parts separately: whitespace unhighlighted, content highlighted | |
| if leading_whitespace: | |
| in_context_highlighted_text.append((leading_whitespace, None)) | |
| if content: | |
| in_context_highlighted_text.append((content, f"rank_{rank}")) | |
| if trailing_whitespace: | |
| in_context_highlighted_text.append((trailing_whitespace, None)) | |
| else: | |
| # Even if marked as important, don't highlight whitespace-only segments | |
| in_context_highlighted_text.append((source_text, None)) | |
| else: | |
| # Add unhighlighted text for non-important segments | |
| in_context_highlighted_text.append((source_text, None)) | |
| # Enhanced color map with ranking-based colors | |
| color_map = {} | |
| for rank in range(len(important_ids)): | |
| _, rgba_color = get_color_by_rank(rank + 1, total_sources) | |
| color_map[f"rank_{rank+1}"] = rgba_color | |
| dummy_update = gr.update( | |
| value=f"AttnTrace_{state.response}_{state.start_index}_{state.end_index}" | |
| ) | |
| attribute_error_update = gr.update(visible=False) | |
| # Combine sources and highlighted context into a single display | |
| # Sources at the top | |
| combined_display = [] | |
| # Add sources header (no highlighting for UI elements) | |
| combined_display.append(("═══ FULL RESPONSE TRACEBACK RESULTS ═══\n", None)) | |
| combined_display.append(("These are the text segments that contribute most to the entire response:\n\n", None)) | |
| # Add sources using available data | |
| for rank, i in enumerate(sorted_indices): | |
| if i < len(important_ids): | |
| source_text = texts[important_ids[i]] | |
| # Strip leading/trailing whitespace from source text to avoid highlighting newlines | |
| clean_source_text = source_text.strip() | |
| if clean_source_text: # Only add if there's actual content | |
| # Add the source text with highlighting, then add spacing without highlighting | |
| combined_display.append((clean_source_text, f"rank_{rank+1}")) | |
| combined_display.append(("\n\n", None)) | |
| # Add separator (no highlighting for UI elements) | |
| combined_display.append(("\n" + "═"*50 + "\n", None)) | |
| combined_display.append(("FULL CONTEXT WITH HIGHLIGHTS\n", None)) | |
| combined_display.append(("Scroll down to see the complete context with important segments highlighted:\n\n", None)) | |
| # Add highlighted context using in_context_highlighted_text | |
| combined_display.extend(in_context_highlighted_text) | |
| # Use only the ranking colors (no highlighting for UI elements) | |
| enhanced_color_map = color_map.copy() | |
| combined_sources_update = HighlightedTextbox( | |
| value=combined_display, color_map=enhanced_color_map, visible=True | |
| ) | |
| # Switch to the highlighted context tab and show results | |
| basic_context_tabs_update = gr.update(selected=1) | |
| basic_sources_in_context_tab_update = gr.update(visible=True) | |
| return ( | |
| combined_sources_update, | |
| basic_context_tabs_update, | |
| basic_sources_in_context_tab_update, | |
| dummy_update, | |
| attribute_error_update, | |
| state, | |
| ) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[(f"❌ Error: {str(e)}", None)], visible=True), | |
| state, | |
| ) | |
| def basic_get_scores_and_sources( | |
| evt: gr.SelectData, | |
| highlighted_response: List[Dict[str, str]], | |
| state: State, | |
| ): | |
| # Get the selected sentence | |
| print("highlighted_response: ", highlighted_response[evt.index]) | |
| selected_text = highlighted_response[evt.index]['token'] | |
| state.explained_response_part = selected_text | |
| # Attribution using default configuration | |
| _, attr, error_msg = initialize_model_and_attr() | |
| if attr is None: | |
| error_text = error_msg if error_msg else "Traceback initialization failed!" | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[(f"❌ {error_text}", None)], visible=True), | |
| state, | |
| ) | |
| try: | |
| # Validate attribution inputs | |
| if not state.context or not state.context.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No context available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| if not state.query or not state.query.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No query available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| if not state.full_response or not state.full_response.strip(): | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No response available for traceback!", None)], visible=True), | |
| state, | |
| ) | |
| print(f"start traceback with explanation_level: {DEFAULT_EXPLANATION_LEVEL}") | |
| print(f"context length: {len(state.context)}, query: {state.query[:100]}...") | |
| print(f"response: {state.full_response[:100]}...") | |
| print(f"selected part: {state.explained_response_part[:100]}...") | |
| texts, important_ids, importance_scores, _, _ = attr.attribute( | |
| state.query, [state.context], state.full_response, state.explained_response_part | |
| ) | |
| print("end traceback") | |
| print(f"explanation_level: {DEFAULT_EXPLANATION_LEVEL}") | |
| print(f"texts count: {len(texts)} (how context was segmented)") | |
| if len(texts) > 0: | |
| print(f"sample text segments: {[text[:50] + '...' if len(text) > 50 else text for text in texts[:3]]}") | |
| print(f"important_ids: {important_ids}") | |
| print("importance_scores: ", importance_scores) | |
| if not importance_scores: | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[("❌ No traceback scores generated! Try a different text segment.", None)], visible=True), | |
| state, | |
| ) | |
| state.scores = np.array(importance_scores) | |
| # Highlighted sources with ranking-based colors | |
| highlighted_text = [] | |
| sorted_indices = np.argsort(state.scores)[::-1] | |
| total_sources = len(important_ids) | |
| for rank, i in enumerate(sorted_indices): | |
| source_text = texts[important_ids[i]] | |
| _ = get_color_by_rank(rank + 1, total_sources) | |
| highlighted_text.append( | |
| ( | |
| source_text, | |
| f"rank_{rank+1}", | |
| ) | |
| ) | |
| # In-context highlights with ranking-based colors - show ALL text | |
| in_context_highlighted_text = [] | |
| ranks = {important_ids[i]: rank for rank, i in enumerate(sorted_indices)} | |
| for i in range(len(texts)): | |
| source_text = texts[i] | |
| # Skip or don't highlight segments that are only newlines or whitespace | |
| if source_text.strip() == "": | |
| # For whitespace-only segments, add them without highlighting | |
| in_context_highlighted_text.append((source_text, None)) | |
| elif i in important_ids: | |
| # Only highlight if the segment has actual content (not just newlines) | |
| if source_text.strip(): # Has non-whitespace content | |
| rank = ranks[i] + 1 | |
| # Split the segment to separate leading/trailing newlines from content | |
| # This prevents newlines from being highlighted | |
| leading_whitespace = "" | |
| trailing_whitespace = "" | |
| content = source_text | |
| # Extract leading newlines/whitespace | |
| while content and content[0] in ['\n', '\r', '\t', ' ']: | |
| leading_whitespace += content[0] | |
| content = content[1:] | |
| # Extract trailing newlines/whitespace | |
| while content and content[-1] in ['\n', '\r', '\t', ' ']: | |
| trailing_whitespace = content[-1] + trailing_whitespace | |
| content = content[:-1] | |
| # Add the parts separately: whitespace unhighlighted, content highlighted | |
| if leading_whitespace: | |
| in_context_highlighted_text.append((leading_whitespace, None)) | |
| if content: | |
| in_context_highlighted_text.append((content, f"rank_{rank}")) | |
| if trailing_whitespace: | |
| in_context_highlighted_text.append((trailing_whitespace, None)) | |
| else: | |
| # Even if marked as important, don't highlight whitespace-only segments | |
| in_context_highlighted_text.append((source_text, None)) | |
| else: | |
| # Add unhighlighted text for non-important segments | |
| in_context_highlighted_text.append((source_text, None)) | |
| # Enhanced color map with ranking-based colors | |
| color_map = {} | |
| for rank in range(len(important_ids)): | |
| _, rgba_color = get_color_by_rank(rank + 1, total_sources) | |
| color_map[f"rank_{rank+1}"] = rgba_color | |
| dummy_update = gr.update( | |
| value=f"AttnTrace_{state.response}_{state.start_index}_{state.end_index}" | |
| ) | |
| attribute_error_update = gr.update(visible=False) | |
| # Combine sources and highlighted context into a single display | |
| # Sources at the top | |
| combined_display = [] | |
| # Add sources header (no highlighting for UI elements) | |
| combined_display.append(("═══ TRACEBACK RESULTS ═══\n", None)) | |
| combined_display.append(("These are the text segments that contribute most to the response:\n\n", None)) | |
| # Add sources using available data | |
| for rank, i in enumerate(sorted_indices): | |
| if i < len(important_ids): | |
| source_text = texts[important_ids[i]] | |
| # Strip leading/trailing whitespace from source text to avoid highlighting newlines | |
| clean_source_text = source_text.strip() | |
| if clean_source_text: # Only add if there's actual content | |
| # Add the source text with highlighting, then add spacing without highlighting | |
| combined_display.append((clean_source_text, f"rank_{rank+1}")) | |
| combined_display.append(("\n\n", None)) | |
| # Add separator (no highlighting for UI elements) | |
| combined_display.append(("\n" + "═"*50 + "\n", None)) | |
| combined_display.append(("FULL CONTEXT WITH HIGHLIGHTS\n", None)) | |
| combined_display.append(("Scroll down to see the complete context with important segments highlighted:\n\n", None)) | |
| # Add highlighted context using in_context_highlighted_text | |
| combined_display.extend(in_context_highlighted_text) | |
| # Use only the ranking colors (no highlighting for UI elements) | |
| enhanced_color_map = color_map.copy() | |
| combined_sources_update = HighlightedTextbox( | |
| value=combined_display, color_map=enhanced_color_map, visible=True | |
| ) | |
| # Switch to the highlighted context tab and show results | |
| basic_context_tabs_update = gr.update(selected=1) | |
| basic_sources_in_context_tab_update = gr.update(visible=True) | |
| return ( | |
| combined_sources_update, | |
| basic_context_tabs_update, | |
| basic_sources_in_context_tab_update, | |
| dummy_update, | |
| attribute_error_update, | |
| state, | |
| ) | |
| except Exception as e: | |
| traceback.print_exc() | |
| return ( | |
| gr.update(value=[("", None)], visible=False), | |
| gr.update(selected=0), | |
| gr.update(visible=False), | |
| gr.update(value=""), | |
| gr.update(value=[(f"❌ Error: {str(e)}", None)], visible=True), | |
| state, | |
| ) | |
| def load_custom_css(): | |
| """Load CSS from external file""" | |
| try: | |
| with open("assets/app_styles.css", "r") as f: | |
| css_content = f.read() | |
| return css_content | |
| except FileNotFoundError: | |
| print("Warning: CSS file not found, using minimal CSS") | |
| return "" | |
| except Exception as e: | |
| print(f"Error loading CSS: {e}") | |
| return "" | |
| # Load CSS from external file | |
| custom_css = load_custom_css() | |
| theme = gr.themes.Citrus( | |
| text_size="lg", | |
| spacing_size="md", | |
| ) | |
| with gr.Blocks(theme=theme, css=custom_css) as demo: | |
| gr.Markdown(f"# {APP_TITLE}") | |
| gr.Markdown(APP_DESCRIPTION, elem_classes="app-description") | |
| # gr.Markdown(NEW_TEXT, elem_classes="app-description-2") | |
| gr.Markdown(""" | |
| <div style="font-size: 18px;"> | |
| AttnTrace is an efficient context traceback method for long contexts (e.g., full papers). It is over 15× faster than the state-of-the-art context traceback method TracLLM. Compared to previous attention-based approaches, AttnTrace is more accurate, reliable, and memory-efficient. | |
| """, elem_classes="feature-highlights") | |
| # Feature highlights | |
| gr.Markdown(""" | |
| <div style="font-size: 18px;"> | |
| AttnTrace can be used in many real-world applications, such as tracing back to: | |
| - 📄 prompt injection instructions that manipulate LLM-generated paper reviews. | |
| - 💻 malicious comment & code hiding in the codebase that misleads the AI coding assistant. | |
| - 🤖 malicious instructions that mislead the action of the LLM agent. | |
| - 🖋 source texts in the context from an AI summary. | |
| - 🔍 evidence that supports the LLM-generated answer for a question. | |
| - ❌ misinformation (corrupted knowledge) that manipulates LLM output for a question. | |
| - And a lot more... | |
| </div> | |
| """, elem_classes="feature-highlights") | |
| # Example buttons with topic-relevant images - moved here for better positioning | |
| gr.Markdown("### 🚀 Try These Examples!", elem_classes="example-title") | |
| with gr.Row(elem_classes=["example-button-container"]): | |
| with gr.Column(scale=1): | |
| example_1_btn = gr.Button( | |
| "📄 Prompt Injection Attacks in AI Paper Review", | |
| elem_classes=["example-button", "example-paper"], | |
| elem_id="example_1_button", | |
| scale=None, | |
| size="sm" | |
| ) | |
| with gr.Column(scale=1): | |
| example_2_btn = gr.Button( | |
| "💻 Malicious Comments & Code in Codebase", | |
| elem_classes=["example-button", "example-movie"], | |
| elem_id="example_2_button" | |
| ) | |
| with gr.Column(scale=1): | |
| example_3_btn = gr.Button( | |
| "🤖 Malicious Instructions Misleading the LLM Agent", | |
| elem_classes=["example-button", "example-code"], | |
| elem_id="example_3_button" | |
| ) | |
| with gr.Row(elem_classes=["example-button-container"]): | |
| with gr.Column(scale=1): | |
| example_4_btn = gr.Button( | |
| "🖋 Source Texts for an AI Summary", | |
| elem_classes=["example-button", "example-paper-alt"], | |
| elem_id="example_4_button" | |
| ) | |
| with gr.Column(scale=1): | |
| example_5_btn = gr.Button( | |
| "🔍 Evidence that Support Question Answering", | |
| elem_classes=["example-button", "example-movie-alt"], | |
| elem_id="example_5_button" | |
| ) | |
| with gr.Column(scale=1): | |
| example_6_btn = gr.Button( | |
| "❌ Misinformation (Corrupted Knowledge) in Question Answering", | |
| elem_classes=["example-button", "example-code-alt"], | |
| elem_id="example_6_button" | |
| ) | |
| state = gr.State( | |
| value=clear_state() | |
| ) | |
| basic_tab = gr.Tab("Demo") | |
| with basic_tab: | |
| # gr.Markdown("## Demo") | |
| gr.Markdown( | |
| "Enter your context and instruction below to try out AttnTrace! You can also click on the example buttons above to load pre-configured examples." | |
| ) | |
| gr.Markdown( | |
| '**Color Legend for Context Traceback (by ranking):** <span style="background-color: #FF4444; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Red</span> = 1st (most important) | <span style="background-color: #FF8C42; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Orange</span> = 2nd | <span style="background-color: #FFD93D; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Golden</span> = 3rd | <span style="background-color: #FFF280; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Yellow</span> = 4th-5th | <span style="background-color: #FFF9C4; color: black; padding: 2px 6px; border-radius: 4px; font-weight: 600;">Light</span> = 6th+' | |
| ) | |
| # Top section: Wide Context box with tabs | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| with gr.Tabs() as basic_context_tabs: | |
| with gr.TabItem("Context", id=0): | |
| basic_context_box = gr.Textbox( | |
| placeholder="Enter context...", | |
| show_label=False, | |
| value="", | |
| lines=6, | |
| max_lines=6, | |
| elem_id="basic_context_box", | |
| autoscroll=False, | |
| ) | |
| with gr.TabItem("Context with highlighted traceback results", id=1, visible=True) as basic_sources_in_context_tab: | |
| basic_sources_in_context_box = HighlightedTextbox( | |
| value=[("Click on a sentence in the response below to see highlighted traceback results here.", None)], | |
| show_legend_label=False, | |
| show_label=False, | |
| show_legend=False, | |
| interactive=False, | |
| elem_id="basic_sources_in_context_box", | |
| ) | |
| # Error messages | |
| basic_generate_error_box = HighlightedTextbox( | |
| show_legend_label=False, | |
| show_label=False, | |
| show_legend=False, | |
| visible=False, | |
| interactive=False, | |
| container=False, | |
| ) | |
| # Bottom section: Left (instruction + button + response), Right (response selection) | |
| with gr.Row(equal_height=True): | |
| # Left: Instruction + Button + Response | |
| with gr.Column(scale=1): | |
| basic_query_box = gr.Textbox( | |
| label="Instruction", | |
| placeholder="Enter an instruction...", | |
| value="", | |
| lines=3, | |
| max_lines=3, | |
| ) | |
| unified_response_button = gr.Button( | |
| "Generate/Use Response", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| response_input_box = gr.Textbox( | |
| label="Response (Editable)", | |
| placeholder="Response will appear here after generation, or type your own response for traceback...", | |
| lines=8, | |
| max_lines=8, | |
| info="Leave empty and click button to generate from LLM, or type your own response to use for traceback" | |
| ) | |
| # Right: Response for attribution selection | |
| with gr.Column(scale=1): | |
| basic_response_box = gr.HighlightedText( | |
| label="Click to select text for traceback!", | |
| value=[("Click the 'Generate/Use Response' button on the left to see response text here for traceback analysis.", None)], | |
| interactive=False, | |
| combine_adjacent=False, | |
| show_label=True, | |
| show_legend=False, | |
| elem_id="basic_response_box", | |
| visible=True, | |
| ) | |
| # Button for full response traceback | |
| full_response_traceback_button = gr.Button( | |
| "🔍 Traceback Entire Response", | |
| variant="secondary", | |
| size="sm" | |
| ) | |
| # Hidden error box and dummy elements | |
| basic_attribute_error_box = HighlightedTextbox( | |
| show_legend_label=False, | |
| show_label=False, | |
| show_legend=False, | |
| visible=False, | |
| interactive=False, | |
| container=False, | |
| ) | |
| dummy_basic_sources_box = gr.Textbox( | |
| visible=False, interactive=False, container=False | |
| ) | |
| # Only a single (AttnTrace) method and model in this simplified version | |
| def basic_clear_state(): | |
| state = clear_state() | |
| return ( | |
| "", # basic_context_box | |
| "", # basic_query_box | |
| "", # response_input_box | |
| gr.update(value=[("Click the 'Generate/Use Response' button above to see response text here for traceback analysis.", None)]), # basic_response_box - keep visible | |
| gr.update(selected=0), # basic_context_tabs - switch to first tab | |
| state, | |
| ) | |
| # Defining behavior of various interactions for the basic tab | |
| basic_tab.select( | |
| fn=basic_clear_state, | |
| inputs=[], | |
| outputs=[ | |
| basic_context_box, | |
| basic_query_box, | |
| response_input_box, | |
| basic_response_box, | |
| basic_context_tabs, | |
| state, | |
| ], | |
| ) | |
| for component in [basic_context_box, basic_query_box]: | |
| component.change( | |
| basic_update, | |
| [basic_context_box, basic_query_box, state], | |
| [ | |
| basic_response_box, | |
| basic_context_tabs, | |
| state, | |
| ], | |
| ) | |
| # Example button event handlers - now update both UI and state | |
| outputs_for_examples = [ | |
| basic_context_box, | |
| basic_query_box, | |
| state, | |
| response_input_box, | |
| basic_response_box, | |
| basic_context_tabs, | |
| ] | |
| example_1_btn.click( | |
| fn=partial(load_an_example, run_example_1), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| example_2_btn.click( | |
| fn=partial(load_an_example, run_example_2), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| example_3_btn.click( | |
| fn=partial(load_an_example, run_example_3), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| example_4_btn.click( | |
| fn=partial(load_an_example, run_example_4), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| example_5_btn.click( | |
| fn=partial(load_an_example, run_example_5), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| example_6_btn.click( | |
| fn=partial(load_an_example, run_example_6), | |
| inputs=[state], | |
| outputs=outputs_for_examples | |
| ) | |
| unified_response_button.click( | |
| fn=lambda: None, | |
| inputs=[], | |
| outputs=[], | |
| js=get_scroll_js_code("basic_response_box"), | |
| ) | |
| basic_response_box.change( | |
| fn=lambda: None, | |
| inputs=[], | |
| outputs=[], | |
| js=get_scroll_js_code("basic_sources_in_context_box"), | |
| ) | |
| # Add immediate tab switch on response selection | |
| def immediate_tab_switch(): | |
| return ( | |
| gr.update(value=[("🔄 Processing traceback... Please wait...", None)]), # Show progress message | |
| gr.update(selected=1), # Switch to annotation tab immediately | |
| ) | |
| basic_response_box.select( | |
| fn=immediate_tab_switch, | |
| inputs=[], | |
| outputs=[basic_sources_in_context_box, basic_context_tabs], | |
| queue=False, # Execute immediately without queue | |
| ) | |
| basic_response_box.select( | |
| fn=basic_get_scores_and_sources, | |
| inputs=[basic_response_box, state], | |
| outputs=[ | |
| basic_sources_in_context_box, | |
| basic_context_tabs, | |
| basic_sources_in_context_tab, | |
| dummy_basic_sources_box, | |
| basic_attribute_error_box, | |
| state, | |
| ], | |
| show_progress="full", | |
| ) | |
| basic_response_box.select( | |
| fn=basic_update_highlighted_response, | |
| inputs=[state], | |
| outputs=[basic_response_box, state], | |
| ) | |
| # Full response traceback button | |
| full_response_traceback_button.click( | |
| fn=immediate_tab_switch, | |
| inputs=[], | |
| outputs=[basic_sources_in_context_box, basic_context_tabs], | |
| queue=False, # Execute immediately without queue | |
| ) | |
| full_response_traceback_button.click( | |
| fn=basic_get_scores_and_sources_full_response, | |
| inputs=[state], | |
| outputs=[ | |
| basic_sources_in_context_box, | |
| basic_context_tabs, | |
| basic_sources_in_context_tab, | |
| dummy_basic_sources_box, | |
| basic_attribute_error_box, | |
| state, | |
| ], | |
| show_progress="full", | |
| ) | |
| dummy_basic_sources_box.change( | |
| fn=lambda: None, | |
| inputs=[], | |
| outputs=[], | |
| js=get_scroll_js_code("basic_sources_in_context_box"), | |
| ) | |
| # Unified response handler | |
| unified_response_button.click( | |
| fn=unified_response_handler, | |
| inputs=[response_input_box, state], | |
| outputs=[state, response_input_box, basic_response_box, basic_generate_error_box] | |
| ) | |
| # gr.Markdown( | |
| # "Please do not interact with elements while generation/attribution is in progress. This may cause errors. You can refresh the page if you run into issues because of this." | |
| # ) | |
| demo.launch(show_api=False, share=True) | |