Initial commit

.gitignore

@@ -0,0 +1,4 @@
+*.pyc
+*.html
+*.json
+.DS_Store

README.md

@@ -1,13 +1,21 @@
 ---
-title: Gemma Lxt Mirage
-emoji: 📚
-colorFrom: red
-colorTo: red
+title: MIRAGE
+emoji: 🌴
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 4.41.0
+sdk_version: 4.42.0
 app_file: app.py
-pinned: false
+pinned: true
 license: apache-2.0
+short_description: Model Internals to generate RAG citations
+tags:
+- answer-attribution
+- interpretability
+- context-usage
+- language-modeling
+- arxiv:2406.13663
+- arxiv:2402.05602
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Demo for the paper [Model Internals-based Answer Attribution for Trustworthy Retrieval-Augmented Generation](https://arxiv.org/abs/2406.13663) using the AttnLRP method from [AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers](https://arxiv.org/abs/2402.05602).

app.py

@@ -0,0 +1,268 @@
+import re
+import os
+import bm25s
+import spaces
+import gradio as gr
+import gradio_iframe
+from bm25s.hf import BM25HF
+from rerankers import Reranker
+
+from inseq import register_step_function, load_model
+from inseq.attr import StepFunctionArgs
+from inseq.commands.attribute_context import visualize_attribute_context
+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_id = "HuggingFaceTB/SmolLM-360M-Instruct"
+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
+
+
+def lxt_probability_fn(args: StepFunctionArgs):
+    logits = args.attribution_model.output2logits(args.forward_output)
+    target_ids = args.target_ids.reshape(logits.shape[0], 1).to(logits.device)
+    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_targets_alignments: list[list[tuple[int, int]]] | None = None,
+    contrast_force_inputs: bool = False,
+    skip_special_tokens: bool = False,
+):
+    c_args = _setup_contrast_args(
+        args,
+        contrast_sources=contrast_sources,
+        contrast_targets=contrast_targets,
+        contrast_targets_alignments=contrast_targets_alignments,
+        contrast_force_inputs=contrast_force_inputs,
+        skip_special_tokens=skip_special_tokens,
+    )
+    return lxt_probability_fn(c_args)
+
+def lxt_contrast_prob_diff_fn(
+    args: StepFunctionArgs,
+    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,
+):
+    model_probs = lxt_probability_fn(args)
+    contrast_probs = lxt_contrast_prob_fn(
+        args=args,
+        contrast_sources=contrast_sources,
+        contrast_targets=contrast_targets,
+        contrast_targets_alignments=contrast_targets_alignments,
+        contrast_force_inputs=contrast_force_inputs,
+        skip_special_tokens=skip_special_tokens,
+    ).to(model_probs.device)
+    return add2(model_probs, mul2(contrast_probs, -1))
+
+
+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)
+    elif rag_setting == "Use Custom Context":
+        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()):
+    global model, model_id
+    if rag_setting == "Use Custom Context":
+        docs = custom_context.split("\n\n")
+        progress(0.1, desc="Using custom context...")
+    else:
+        if not query:
+            raise gr.Error("Please enter a query.")
+        progress(0, desc="Retrieving with BM25...")
+        q = bm25s.tokenize(query)
+        results = retriever.retrieve(q, k=retrieve_k)
+        if rag_setting == "Rerank with ColBERT":
+            progress(0.1, desc="Reranking with ColBERT...")
+            docs = [x["text"] for x in results.documents[0]]
+            out = ranker.rank(query=query, docs=docs)
+            docs = [out.results[i].document.text for i in range(top_k)]
+        else:
+            docs = [results.documents[0][i]["text"] for i in range(top_k)]
+        docs = [re.sub(r"\[\d+\]", "", doc) for doc in docs]
+    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)
+    progress(0.3, desc="Attributing with LXT...")
+    lm_rag_prompting_example = AttributeContextArgs(
+        model_name_or_path=model_id,
+        input_context_text="\n\n".join(docs),
+        input_current_text=query,
+        output_template="{current}",
+        attributed_fn="lxt_contrast_prob_diff",
+        input_template="<|im_start|>user\n### Context\n{context}\n\n### Query\n{current}<|im_end|>\n<|im_start|>assistant\n",
+        contextless_input_current_text="<|im_start|>user\n### Query\n{current}<|im_end|>\n<|im_start|>assistant\n",
+        attribution_method="saliency",
+        show_viz=False,
+        show_intermediate_outputs=False,
+        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},
+        attribution_aggregators=["sum"],
+        rescale_attributions=True,
+        save_path=os.path.join(os.path.dirname(__file__), "outputs/output.json"),
+        viz_path=os.path.join(os.path.dirname(__file__), "outputs/output.html"),
+    )
+    out = attribute_context_with_model(lm_rag_prompting_example, model)
+    html = visualize_attribute_context(out, show_viz=False, return_html=True)
+    return [
+        gradio_iframe.iFrame(html, height=500, visible=True),
+        gr.DownloadButton(
+            label="📂 Download output",
+            value=os.path.join(os.path.dirname(__file__), "outputs/output.json"),
+            visible=True,
+        ),
+        gr.DownloadButton(
+            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)
+
+
+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>')
+    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>"
+        "📋 <i>Retrieval is performed on <a href='https://huggingface.co/datasets/google-research-datasets/natural_questions' target='_blank'>Natural Questions</a> using <a href='https://github.com/xhluca/bm25s' target='_blank'>BM25S</a>, with optional reranking via <a href='https://huggingface.co/answerdotai/answerai-colbert-small-v1' target='_blank'>ColBERT</a>."
+        " <a href='https://huggingface.co/blog/smollm' target='_blank'>SmolLM</a> models are used for generation, while <a href='https://github.com/inseq-team/inseq' target='_blank'>Inseq</a> and <a href='https://github.com/rachtibat/LRP-eXplains-Transformers' target='_blank'>LXT</a> are used for attribution.</i><br>"
+        "➡️ <i>For more details, see also our <a href='https://huggingface.co/spaces/gsarti/pecore' target='_blank'>PECoRe Demo</a>",
+    )
+    with gr.Row():
+        with gr.Column():
+            query = gr.Textbox(
+                placeholder="Insert a query for the language model...",
+                label="Model query",
+                interactive=True,
+                lines=2,
+            )
+            attribute_input_examples = gr.Examples(
+                examples,
+                inputs=[query],
+                examples_per_page=2,
+            )
+            with gr.Accordion("⚙️ Parameters", open=False):
+                with gr.Row():
+                    model_size = gr.Radio(
+                        ["135M", "360M", "1.7B"],
+                        value="360M",
+                        label="Model size",
+                        interactive=True
+                    )
+                with gr.Row():
+                    rag_setting = gr.Radio(
+                        ["Retrieve with BM25", "Rerank with ColBERT", "Use Custom Context"],
+                        value="Rerank with ColBERT",
+                        label="Mode",
+                        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)
+                custom_context = gr.Textbox(
+                    placeholder="Context will be retrieved automatically. Change mode to 'Use Custom Context' to specify your own.",
+                    label="Custom context",
+                    interactive=False,
+                    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)
+            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("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)")
+            btn = gr.Button("Submit", variant="primary")
+        with gr.Column():
+            attribute_context_out = gradio_iframe.iFrame(height=400, visible=True)
+            with gr.Row(equal_height=True):
+                download_output_file_button = gr.DownloadButton(
+                    "📂 Download output",
+                    visible=False,
+                )
+                download_output_html_button = gr.DownloadButton(
+                    "🔍 Download HTML",
+                    visible=False,
+                    value=os.path.join(
+                        os.path.dirname(__file__), "outputs/output.html"
+                    ),
+                )
+    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>""")
+        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>""")
+                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>""")
+                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>""")
+
+    rag_setting.change(
+        fn=set_interactive_settings,
+        inputs=[rag_setting, retrieve_k, top_k, custom_context],
+        outputs=[retrieve_k, top_k, custom_context],
+    )
+
+    btn.click(
+        fn=generate,
+        inputs=[
+            query,
+            max_new_tokens,
+            top_p,
+            temperature,
+            retrieve_k,
+            top_k,
+            rag_setting,
+            custom_context,
+            model_size,
+        ],
+        outputs=[
+            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)

citations.py

@@ -0,0 +1,53 @@
+pecore_citation = """@inproceedings{sarti-etal-2023-quantifying,
+    title = "Quantifying the Plausibility of Context Reliance in Neural Machine Translation",
+    author = "Sarti, Gabriele and
+        Chrupa{\l}a, Grzegorz and
+        Nissim, Malvina and
+        Bisazza, Arianna",
+    booktitle = "The Twelfth International Conference on Learning Representations (ICLR 2024)",
+    month = may,
+    year = "2024",
+    address = "Vienna, Austria",
+    publisher = "OpenReview",
+    url = "https://openreview.net/forum?id=XTHfNGI3zT"
+}"""
+
+inseq_citation = """@inproceedings{sarti-etal-2023-inseq,
+    title = "Inseq: An Interpretability Toolkit for Sequence Generation Models",
+    author = "Sarti, Gabriele  and
+        Feldhus, Nils  and
+        Sickert, Ludwig  and
+        van der Wal, Oskar and
+        Nissim, Malvina and
+        Bisazza, Arianna",
+    booktitle = "Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 3: System Demonstrations)",
+    month = jul,
+    year = "2023",
+    address = "Toronto, Canada",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2023.acl-demo.40",
+    pages = "421--435",
+}"""
+
+mirage_citation = """@article{qi-sarti-etal-2024-mirage,
+    title = "Model Internals-based Answer Attribution for Trustworthy Retrieval-Augmented Generation",
+    author = "Qi, Jirui and Sarti, Gabriele and Fern{\'a}ndez, Raquel and Bisazza, Arianna",
+    journal = "ArXiv",
+    month = jun,
+    year = "2024",
+    volume = {abs/2406.13663},
+    url = {https://arxiv.org/abs/2406.13663},
+}"""
+
+lxt_citation = """@inproceedings{achtibat-etal-2024-attnlrp,
+  title = {{A}ttn{LRP}: Attention-Aware Layer-Wise Relevance Propagation for Transformers},
+  author = {Achtibat, Reduan and Hatefi, Sayed Mohammad Vakilzadeh and Dreyer, Maximilian and Jain, Aakriti and Wiegand, Thomas and Lapuschkin, Sebastian and Samek, Wojciech},
+  booktitle = {Proceedings of the 41st International Conference on Machine Learning},
+  pages = {135--168},
+  year = {2024},
+  editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix},
+  volume = {235},
+  series = {Proceedings of Machine Learning Research},
+  month = {21--27 Jul},
+  publisher = {PMLR}
+}"""

examples.py

@@ -0,0 +1,4 @@
+examples = [
+    "Who was the greek gooddess of spring growth?",
+    "When was the heir to the throne of the United Kingdom born?",
+]

requirements.txt

@@ -0,0 +1,5 @@
+spaces
+git+https://github.com/inseq-team/inseq.git@main
+bm25s
+rerankers[transformers]
+git+https://github.com/rachtibat/LRP-eXplains-Transformers

style.py

@@ -0,0 +1,37 @@
+custom_css = """
+h1 > img {
+    text-align: center;
+    display: block;
+    margin-bottom: 0;
+    font-size: 1.7em;
+}
+
+iframe {
+    overflow: scroll;
+    border: 2px solid green;
+}
+
+.summary-label {
+    display: inline;
+}
+.prose a:visited {
+  color: var(--link-text-color);
+}
+.footer-container {
+    align-items: center;
+}
+.footer-custom-block {
+    display: flex;
+    justify-content: center;
+    align-items: center;
+}
+.footer-custom-block b {
+    margin-right: 10px;
+}
+.footer-custom-block img {
+    margin-right: 15px;
+}
+ol {
+    padding-left: 30px;
+}
+"""