SecureLLMSys commited on
Commit
9006816
·
1 Parent(s): 7279d70
Files changed (1) hide show
  1. app.py +3 -149
app.py CHANGED
@@ -19,7 +19,7 @@ from examples import run_example_1, run_example_2, run_example_3, run_example_4,
19
  from functools import partial
20
 
21
  # Load original app constants
22
- APP_TITLE = '<div class="app-title"><span class="brand">AttnTrace</span><span class="subtitle">Attention-based Context Traceback for Long-Context LLMs</span></div>'
23
  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!
24
  Maintained by the AttnTrace team."""
25
  # 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.
@@ -839,153 +839,7 @@ def load_custom_css():
839
  return css_content
840
  except FileNotFoundError:
841
  print("Warning: CSS file not found, using minimal CSS")
842
- return """
843
- /* Add global page margins */
844
- .gradio-container {
845
- padding-left: 12rem !important;
846
- padding-right: 12rem !important;
847
- }
848
-
849
- /* App title styling */
850
- .app-title {
851
- text-align: center !important;
852
- margin: 2rem 0 !important;
853
- }
854
-
855
- .app-title .brand {
856
- color: #333333 !important;
857
- font-weight: 700 !important;
858
- font-size: 3rem !important;
859
- margin-right: 12px !important;
860
- }
861
-
862
- .app-title .subtitle {
863
- color: #666666 !important;
864
- font-weight: 400 !important;
865
- font-size: 1.6rem !important;
866
- display: block !important;
867
- margin-top: 12px !important;
868
- }
869
-
870
- /* App description styling */
871
- .app-description p {
872
- font-size: 1.25rem !important;
873
- color: #555555 !important;
874
- line-height: 1.6 !important;
875
- }
876
-
877
- /* Feature highlights */
878
- .feature-highlights {
879
- font-size: 1.1rem !important;
880
- color: #444444 !important;
881
- line-height: 1.5 !important;
882
- }
883
-
884
- /* Example title */
885
- .example-title {
886
- text-align: center !important;
887
- margin: 2rem 0 1rem 0 !important;
888
- font-size: 1.5rem !important;
889
- font-weight: 600 !important;
890
- color: #333333 !important;
891
- }
892
-
893
- /* Example button container */
894
- .example-button-container {
895
- display: flex !important;
896
- justify-content: center !important;
897
- align-items: center !important;
898
- gap: 1rem !important;
899
- margin: 1rem 0 !important;
900
- flex-wrap: wrap !important;
901
- }
902
-
903
- /* Example buttons */
904
- .example-button button {
905
- background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
906
- color: white !important;
907
- border: none !important;
908
- border-radius: 10px !important;
909
- padding: 12px 20px !important;
910
- font-size: 0.9rem !important;
911
- font-weight: 600 !important;
912
- cursor: pointer !important;
913
- transition: all 0.3s ease !important;
914
- box-shadow: 0 4px 15px rgba(0,0,0,0.1) !important;
915
- min-width: 200px !important;
916
- text-align: center !important;
917
- }
918
-
919
- .example-button button:hover {
920
- transform: translateY(-2px) !important;
921
- box-shadow: 0 6px 20px rgba(0,0,0,0.15) !important;
922
- }
923
-
924
- /* Color legend classes */
925
- .color-red {
926
- background-color: #FF4444 !important;
927
- color: black !important;
928
- padding: 2px 6px !important;
929
- border-radius: 4px !important;
930
- font-weight: 600 !important;
931
- }
932
-
933
- .color-orange {
934
- background-color: #FF8C42 !important;
935
- color: black !important;
936
- padding: 2px 6px !important;
937
- border-radius: 4px !important;
938
- font-weight: 600 !important;
939
- }
940
-
941
- .color-golden {
942
- background-color: #FFD93D !important;
943
- color: black !important;
944
- padding: 2px 6px !important;
945
- border-radius: 4px !important;
946
- font-weight: 600 !important;
947
- }
948
-
949
- .color-yellow {
950
- background-color: #FFF280 !important;
951
- color: black !important;
952
- padding: 2px 6px !important;
953
- border-radius: 4px !important;
954
- font-weight: 600 !important;
955
- }
956
-
957
- .color-light {
958
- background-color: #FFF9C4 !important;
959
- color: black !important;
960
- padding: 2px 6px !important;
961
- border-radius: 4px !important;
962
- font-weight: 600 !important;
963
- }
964
-
965
- /* Responsive design */
966
- @media (max-width: 768px) {
967
- .gradio-container {
968
- padding-left: 1rem !important;
969
- padding-right: 1rem !important;
970
- }
971
-
972
- .app-title .brand {
973
- font-size: 2rem !important;
974
- }
975
-
976
- .app-title .subtitle {
977
- font-size: 1.2rem !important;
978
- }
979
-
980
- .example-button-container {
981
- flex-direction: column !important;
982
- }
983
-
984
- .example-button button {
985
- min-width: 100% !important;
986
- }
987
- }
988
- """
989
  except Exception as e:
990
  print(f"Error loading CSS: {e}")
991
  return ""
@@ -1087,7 +941,7 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
1087
  )
1088
 
1089
  gr.Markdown(
1090
- '**Color Legend for Context Traceback (by ranking):** <span class="color-red">Red</span> = 1st (most important) | <span class="color-orange">Orange</span> = 2nd | <span class="color-golden">Golden</span> = 3rd | <span class="color-yellow">Yellow</span> = 4th-5th | <span class="color-light">Light</span> = 6th+'
1091
  )
1092
 
1093
 
 
19
  from functools import partial
20
 
21
  # Load original app constants
22
+ APP_TITLE = '<div class="app-title"><span class="brand">AttnTrace: </span><span class="subtitle">Attention-based Context Traceback for Long-Context LLMs</span></div>'
23
  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!
24
  Maintained by the AttnTrace team."""
25
  # 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.
 
839
  return css_content
840
  except FileNotFoundError:
841
  print("Warning: CSS file not found, using minimal CSS")
842
+ return ""
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
843
  except Exception as e:
844
  print(f"Error loading CSS: {e}")
845
  return ""
 
941
  )
942
 
943
  gr.Markdown(
944
+ '**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+'
945
  )
946
 
947