app.py

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)