Ruff formatting

app.py

@@ -1,35 +1,48 @@
-import re
 import os
+import re
+
 import bm25s
-import spaces
 import gradio as gr
 import gradio_iframe
+import spaces
 from bm25s.hf import BM25HF
+from citations import inseq_citation, lxt_citation, mirage_citation, pecore_citation
+from examples import examples
+from lxt.functional import add2, mul2, softmax
+from lxt.models.llama import LlamaForCausalLM, attnlrp
 from rerankers import Reranker
+from style import custom_css
+from transformers import AutoTokenizer
 
-from inseq import register_step_function, load_model
+from inseq import load_model, register_step_function
 from inseq.attr import StepFunctionArgs
 from inseq.commands.attribute_context import visualize_attribute_context
+from inseq.commands.attribute_context.attribute_context import (
+    AttributeContextArgs,
+    attribute_context_with_model,
+)
 from inseq.utils.contrast_utils import _setup_contrast_args
-from lxt.models.llama import LlamaForCausalLM, attnlrp
-from transformers import AutoTokenizer
-from lxt.functional import softmax, add2, mul2
-from inseq.commands.attribute_context.attribute_context import attribute_context_with_model, AttributeContextArgs
-
-from style import custom_css
-from citations import pecore_citation, mirage_citation, inseq_citation, lxt_citation
-from examples import examples
 
+model = None
 model_id = "HuggingFaceTB/SmolLM-360M-Instruct"
-ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type='colbert')
+ranker = Reranker("answerdotai/answerai-colbert-small-v1", model_type="colbert")
 retriever = BM25HF.load_from_hub("xhluca/bm25s-nq-index", load_corpus=True, mmap=True)
-hf_model = LlamaForCausalLM.from_pretrained(model_id)
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-attnlrp.register(hf_model)
-model = load_model(hf_model, "saliency", tokenizer=tokenizer)
-# Needed since the <|im_start|> token is also the BOS
-model.bos_token = "<|endoftext|>"
-model.bos_token_id = 0
+
+# Model registry to store loaded models
+model_registry = {}
+
+
+def get_model(model_size):
+    model_id = f"HuggingFaceTB/SmolLM-{model_size}-Instruct"
+    if model_id not in model_registry:
+        hf_model = LlamaForCausalLM.from_pretrained(model_id)
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        attnlrp.register(hf_model)
+        model = load_model(hf_model, "saliency", tokenizer=tokenizer)
+        model.bos_token = "<|endoftext|>"
+        model.bos_token_id = 0
+        model_registry[model_id] = model
+    return model_registry[model_id]
 
 
 def lxt_probability_fn(args: StepFunctionArgs):
@@ -38,10 +51,11 @@ def lxt_probability_fn(args: StepFunctionArgs):
     logits = softmax(logits, dim=-1)
     return logits.gather(-1, target_ids).squeeze(-1)
 
+
 def lxt_contrast_prob_fn(
     args: StepFunctionArgs,
-    contrast_sources = None,
-    contrast_targets = None,
+    contrast_sources=None,
+    contrast_targets=None,
     contrast_targets_alignments: list[list[tuple[int, int]]] | None = None,
     contrast_force_inputs: bool = False,
     skip_special_tokens: bool = False,
@@ -56,10 +70,11 @@ def lxt_contrast_prob_fn(
     )
     return lxt_probability_fn(c_args)
 
+
 def lxt_contrast_prob_diff_fn(
     args: StepFunctionArgs,
-    contrast_sources = None,
-    contrast_targets = None,
+    contrast_sources=None,
+    contrast_targets=None,
     contrast_targets_alignments: list[list[tuple[int, int]]] | None = None,
     contrast_force_inputs: bool = False,
     skip_special_tokens: bool = False,
@@ -78,12 +93,35 @@ def lxt_contrast_prob_diff_fn(
 
 def set_interactive_settings(rag_setting, retrieve_k, top_k, custom_context):
     if rag_setting in ("Retrieve with BM25", "Rerank with ColBERT"):
-        return gr.Slider(interactive=True), gr.Slider(interactive=True), gr.Textbox(placeholder="Context will be retrieved automatically. Change mode to 'Use Custom Context' to specify your own.", interactive=False)
+        return (
+            gr.Slider(interactive=True),
+            gr.Slider(interactive=True),
+            gr.Textbox(
+                placeholder="Context will be retrieved automatically. Change mode to 'Use Custom Context' to specify your own.",
+                interactive=False,
+            ),
+        )
     elif rag_setting == "Use Custom Context":
-        return gr.Slider(interactive=False), gr.Slider(interactive=False), gr.Textbox(placeholder="Insert a custom context...", interactive=True)
+        return (
+            gr.Slider(interactive=False),
+            gr.Slider(interactive=False),
+            gr.Textbox(placeholder="Insert a custom context...", interactive=True),
+        )
+
 
 @spaces.GPU()
-def generate(query, max_new_tokens, top_p, temperature, retrieve_k, top_k, rag_setting, custom_context, model_size, progress=gr.Progress()):
+def generate(
+    query,
+    max_new_tokens,
+    top_p,
+    temperature,
+    retrieve_k,
+    top_k,
+    rag_setting,
+    custom_context,
+    model_size,
+    progress=gr.Progress(),
+):
     global model, model_id
     if rag_setting == "Use Custom Context":
         docs = custom_context.split("\n\n")
@@ -105,11 +143,7 @@ def generate(query, max_new_tokens, top_p, temperature, retrieve_k, top_k, rag_s
     curr_model_id = f"HuggingFaceTB/SmolLM-{model_size}-Instruct"
     if model is None or model.model_name != curr_model_id:
         progress(0.2, desc="Loading model...")
-        model_id = curr_model_id
-        hf_model = LlamaForCausalLM.from_pretrained(model_id)
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
-        attnlrp.register(hf_model)
-        model = load_model(hf_model, "saliency", tokenizer=tokenizer)
+        model = get_model(model_size)
     progress(0.3, desc="Attributing with LXT...")
     lm_rag_prompting_example = AttributeContextArgs(
         model_name_or_path=model_id,
@@ -125,7 +159,11 @@ def generate(query, max_new_tokens, top_p, temperature, retrieve_k, top_k, rag_s
         context_sensitivity_std_threshold=1,
         decoder_input_output_separator=" ",
         special_tokens_to_keep=["<|im_start|>", "<|endoftext|>"],
-        generation_kwargs={"max_new_tokens": max_new_tokens, "top_p": top_p, "temperature": temperature},
+        generation_kwargs={
+            "max_new_tokens": max_new_tokens,
+            "top_p": top_p,
+            "temperature": temperature,
+        },
         attribution_aggregators=["sum"],
         rescale_attributions=True,
         save_path=os.path.join(os.path.dirname(__file__), "outputs/output.json"),
@@ -144,17 +182,21 @@ def generate(query, max_new_tokens, top_p, temperature, retrieve_k, top_k, rag_s
             label="🔍 Download HTML",
             value=os.path.join(os.path.dirname(__file__), "outputs/output.html"),
             visible=True,
-        )
+        ),
     ]
 
 
-register_step_function(lxt_contrast_prob_diff_fn, "lxt_contrast_prob_diff", overwrite=True)
+register_step_function(
+    lxt_contrast_prob_diff_fn, "lxt_contrast_prob_diff", overwrite=True
+)
 
 
 with gr.Blocks(css=custom_css) as demo:
     with gr.Row():
         with gr.Column(min_width=500):
-            gr.HTML(f'<h1><img src="file/img/mirage_logo_white_contour.png" width=300px /></h1>')
+            gr.HTML(
+                '<h1><img src="file/img/mirage_logo_white_contour.png" width=300px /></h1>'
+            )
     text = gr.Markdown(
         "This demo showcases an end-to-end usage of model internals for RAG answer attribution with the <a href='https://openreview.net/forum?id=XTHfNGI3zT' target='_blank'>PECoRe</a> framework, as described in our <a href='https://arxiv.org/abs/2406.13663' target='_blank'>MIRAGE</a> paper.<br>"
         "Insert a query to retrieve relevant contexts, generate an answer and attribute its context-sensitive components. An interactive <a href='https://github.com/google-deepmind/treescope' target='_blank'>Treescope</a> visualization will appear in the green square.<br>"
@@ -182,18 +224,36 @@ with gr.Blocks(css=custom_css) as demo:
                         ["135M", "360M", "1.7B"],
                         value="360M",
                         label="Model size",
-                        interactive=True
+                        interactive=True,
                     )
                 with gr.Row():
                     rag_setting = gr.Radio(
-                        ["Retrieve with BM25", "Rerank with ColBERT", "Use Custom Context"],
+                        [
+                            "Retrieve with BM25",
+                            "Rerank with ColBERT",
+                            "Use Custom Context",
+                        ],
                         value="Rerank with ColBERT",
                         label="Mode",
-                        interactive=True
+                        interactive=True,
                     )
                 with gr.Row():
-                    retrieve_k = gr.Slider(1, 500, value=100, step=1, label="# Docs to Retrieve", interactive=True)
-                    top_k = gr.Slider(1, 10, value=3, step=1, label="# Docs in Context", interactive=True)
+                    retrieve_k = gr.Slider(
+                        1,
+                        500,
+                        value=100,
+                        step=1,
+                        label="# Docs to Retrieve",
+                        interactive=True,
+                    )
+                    top_k = gr.Slider(
+                        1,
+                        10,
+                        value=3,
+                        step=1,
+                        label="# Docs in Context",
+                        interactive=True,
+                    )
                 custom_context = gr.Textbox(
                     placeholder="Context will be retrieved automatically. Change mode to 'Use Custom Context' to specify your own.",
                     label="Custom context",
@@ -201,18 +261,51 @@ with gr.Blocks(css=custom_css) as demo:
                     lines=4,
                 )
                 with gr.Row():
-                    max_new_tokens = gr.Slider(0, 500, value=50, step=5.0, label="Max new tokens", interactive=True)
-                    top_p = gr.Slider(0, 1, value=1, step=0.01, label="Top P", interactive=True)
-                    temperature = gr.Slider(0, 1, value=0, step=0.01, label="Temperature", interactive=True)
+                    max_new_tokens = gr.Slider(
+                        0,
+                        500,
+                        value=50,
+                        step=5.0,
+                        label="Max new tokens",
+                        interactive=True,
+                    )
+                    top_p = gr.Slider(
+                        0, 1, value=1, step=0.01, label="Top P", interactive=True
+                    )
+                    temperature = gr.Slider(
+                        0, 1, value=0, step=0.01, label="Temperature", interactive=True
+                    )
             with gr.Accordion("📝 Citation", open=False):
-                gr.Markdown("Using PECoRe for model internals-based RAG answer attribution is discussed in:")
-                gr.Code(mirage_citation, interactive=False, label="MIRAGE (Qi, Sarti et al., 2024)")
+                gr.Markdown(
+                    "Using PECoRe for model internals-based RAG answer attribution is discussed in:"
+                )
+                gr.Code(
+                    mirage_citation,
+                    interactive=False,
+                    label="MIRAGE (Qi, Sarti et al., 2024)",
+                )
                 gr.Markdown("To refer to the original PECoRe paper, cite:")
-                gr.Code(pecore_citation, interactive=False, label="PECoRe (Sarti et al., 2024)")
-                gr.Markdown("The Inseq implementation used in this work (<a href=\"https://inseq.org/en/latest/main_classes/cli.html#attribute-context\"><code>inseq attribute-context</code></a>, including this demo) can be cited with:")
-                gr.Code(inseq_citation, interactive=False, label="Inseq (Sarti et al., 2023)")
-                gr.Markdown("The AttnLRP attribution method used in this demo via the LXT library can be cited with:")
-                gr.Code(lxt_citation, interactive=False, label="AttnLRP (Achtibat et al., 2024)")
+                gr.Code(
+                    pecore_citation,
+                    interactive=False,
+                    label="PECoRe (Sarti et al., 2024)",
+                )
+                gr.Markdown(
+                    'The Inseq implementation used in this work (<a href="https://inseq.org/en/latest/main_classes/cli.html#attribute-context"><code>inseq attribute-context</code></a>, including this demo) can be cited with:'
+                )
+                gr.Code(
+                    inseq_citation,
+                    interactive=False,
+                    label="Inseq (Sarti et al., 2023)",
+                )
+                gr.Markdown(
+                    "The AttnLRP attribution method used in this demo via the LXT library can be cited with:"
+                )
+                gr.Code(
+                    lxt_citation,
+                    interactive=False,
+                    label="AttnLRP (Achtibat et al., 2024)",
+                )
         with gr.Column():
             attribute_context_out = gradio_iframe.iFrame(height=400, visible=True)
             with gr.Row(equal_height=True):
@@ -229,15 +322,23 @@ with gr.Blocks(css=custom_css) as demo:
                 )
     with gr.Row(elem_classes="footer-container"):
         with gr.Column():
-            gr.Markdown("""<div class="footer-custom-block"><b>Powered by</b> <a href='https://github.com/inseq-team/inseq' target='_blank'><img src="file/img/inseq_logo_white_contour.png" width=150px /></a> <a href='https://github.com/rachtibat/LRP-eXplains-Transformers' target='_blank'><img src="file/img/lxt_logo.png" width=150px /></a></div>""")
+            gr.Markdown(
+                """<div class="footer-custom-block"><b>Powered by</b> <a href='https://github.com/inseq-team/inseq' target='_blank'><img src="file/img/inseq_logo_white_contour.png" width=150px /></a> <a href='https://github.com/rachtibat/LRP-eXplains-Transformers' target='_blank'><img src="file/img/lxt_logo.png" width=150px /></a></div>"""
+            )
         with gr.Column():
             with gr.Row(elem_classes="footer-custom-block"):
                 with gr.Column(scale=0.30, min_width=150):
-                    gr.Markdown("""<b>Built by <a href="https://gsarti.com" target="_blank">Gabriele Sarti</a><br> with the support of</b>""")
+                    gr.Markdown(
+                        """<b>Built by <a href="https://gsarti.com" target="_blank">Gabriele Sarti</a><br> with the support of</b>"""
+                    )
                 with gr.Column(scale=0.30, min_width=120):
-                    gr.Markdown("""<a href='https://www.rug.nl/research/clcg/research/cl/' target='_blank'><img src="file/img/rug_logo_white_contour.png" width=170px /></a>""")
+                    gr.Markdown(
+                        """<a href='https://www.rug.nl/research/clcg/research/cl/' target='_blank'><img src="file/img/rug_logo_white_contour.png" width=170px /></a>"""
+                    )
                 with gr.Column(scale=0.30, min_width=120):
-                    gr.Markdown("""<a href='https://projects.illc.uva.nl/indeep/' target='_blank'><img src="file/img/indeep_logo_white_contour.png" width=100px /></a>""")
+                    gr.Markdown(
+                        """<a href='https://projects.illc.uva.nl/indeep/' target='_blank'><img src="file/img/indeep_logo_white_contour.png" width=100px /></a>"""
+                    )
 
     rag_setting.change(
         fn=set_interactive_settings,
@@ -262,7 +363,9 @@ with gr.Blocks(css=custom_css) as demo:
             attribute_context_out,
             download_output_file_button,
             download_output_html_button,
-        ]
+        ],
     )
 
-demo.queue(api_open=False, max_size=20).launch(allowed_paths=["img/", "outputs/"], show_api=False)
+demo.queue(api_open=False, max_size=20).launch(
+    allowed_paths=["img/", "outputs/"], show_api=False
+)