Finished code generator

app.py

@@ -27,11 +27,11 @@ from presets import (
     set_zephyr_preset,
     set_gemma_preset,
     set_mistral_instruct_preset,
+    update_code_snippets_fn,
 )
 from style import custom_css
 from utils import get_formatted_attribute_context_results
 
-from inseq import list_feature_attribution_methods, list_step_functions
 from inseq.commands.attribute_context.attribute_context import (
     AttributeContextArgs,
     attribute_context_with_model,
@@ -291,21 +291,37 @@ with gr.Blocks(css=custom_css) as demo:
                 value="kl_divergence",
                 label="Context sensitivity metric",
                 info="Metric to use to measure context sensitivity of generated tokens.",
-                choices=list_step_functions(),
+                choices=[
+                    "probability",
+                    "logit",
+                    "kl_divergence",
+                    "contrast_logits_diff",
+                    "contrast_prob_diff",
+                    "pcxmi"
+                ],
                 interactive=True,
             )
             attribution_method = gr.Dropdown(
                 value="saliency",
                 label="Attribution method",
                 info="Attribution method identifier to identify relevant context tokens.",
-                choices=list_feature_attribution_methods(),
+                choices=[
+                    "saliency",
+                    "input_x_gradient",
+                    "value_zeroing",
+                ],
                 interactive=True,
             )
             attributed_fn = gr.Dropdown(
                 value="contrast_prob_diff",
                 label="Attributed function",
                 info="Function of model logits to use as target for the attribution method.",
-                choices=list_step_functions(),
+                choices=[
+                    "probability",
+                    "logit",
+                    "contrast_logits_diff",
+                    "contrast_prob_diff",
+                ],
                 interactive=True,
             )
         gr.Markdown("#### Results Selection Parameters")
@@ -330,7 +346,7 @@ with gr.Blocks(css=custom_css) as demo:
                 maximum=10,
             )
             attribution_std_threshold = gr.Number(
-                value=1.0,
+                value=2.0,
                 label="Attribution threshold",
                 info="Select N to keep attributed tokens with scores above N * std. 0 = above mean.",
                 precision=1,
@@ -461,8 +477,23 @@ with gr.Blocks(css=custom_css) as demo:
                 gr.Markdown(cci_explanation)
     with gr.Tab("🔧 Usage Guide"):
         gr.Markdown(how_to_use)
-        gr.HTML('<img src="file/img/pecore_ui_output_example.png" width=100% />')
         gr.Markdown(example_explanation)
+        update_code_snippets = gr.Button("Update code snippets", variant="primary")
+        with gr.Row(equal_height=True):
+            python_code_snippet = gr.Code(
+                value="""Generate Python code snippet by pressing the button.""",
+                language="python",
+                label="Python",
+                interactive=False,
+                show_label=True,
+            )
+            shell_code_snippet = gr.Code(
+                value="""Generate Shell code snippet by pressing the button.""",
+                language="shell",
+                label="Shell",
+                interactive=False,
+                show_label=True,
+            )
     with gr.Tab("📚 Citing PECoRe"):
         gr.Markdown(citation)
     with gr.Row(elem_classes="footer-container"):
@@ -479,36 +510,38 @@ with gr.Blocks(css=custom_css) as demo:
         tokenizer_kwargs,
     ]
 
+    pecore_args = [
+        input_current_text,
+        input_context_text,
+        output_current_text,
+        output_context_text,
+        model_name_or_path,
+        attribution_method,
+        attributed_fn,
+        context_sensitivity_metric,
+        context_sensitivity_std_threshold,
+        context_sensitivity_topk,
+        attribution_std_threshold,
+        attribution_topk,
+        input_template,
+        output_template,
+        contextless_input_template,
+        contextless_output_template,
+        special_tokens_to_keep,
+        decoder_input_output_separator,
+        model_kwargs,
+        tokenizer_kwargs,
+        generation_kwargs,
+        attribution_kwargs,
+    ]
+
     attribute_input_button.click(
         lambda *args: [gr.DownloadButton(visible=False), gr.DownloadButton(visible=False)],
         inputs=[],
         outputs=[download_output_file_button, download_output_html_button],
     ).then(
         pecore,
-        inputs=[
-            input_current_text,
-            input_context_text,
-            output_current_text,
-            output_context_text,
-            model_name_or_path,
-            attribution_method,
-            attributed_fn,
-            context_sensitivity_metric,
-            context_sensitivity_std_threshold,
-            context_sensitivity_topk,
-            attribution_std_threshold,
-            attribution_topk,
-            input_template,
-            output_template,
-            contextless_input_template,
-            contextless_output_template,
-            special_tokens_to_keep,
-            decoder_input_output_separator,
-            model_kwargs,
-            tokenizer_kwargs,
-            generation_kwargs,
-            attribution_kwargs,
-        ],
+        inputs=pecore_args,
         outputs=[
             pecore_output_highlights,
             download_output_file_button,
@@ -617,4 +650,11 @@ with gr.Blocks(css=custom_css) as demo:
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 
+    update_code_snippets.click(
+        update_code_snippets_fn,
+        inputs=pecore_args,
+        outputs=[python_code_snippet, shell_code_snippet],
+    
+    )
+
 demo.launch(allowed_paths=["outputs/", "img/"])

contents.py

@@ -33,18 +33,31 @@ cci_explanation = """
 """
 
 how_to_use = """
-<h3>How to use this demo</h3>
+<h2>How to use this demo</h3>
 
 <p>This demo provides a convenient UI for the Inseq implementation of PECoRe (the <a href="https://inseq.org/en/latest/main_classes/cli.html#attribute-context"><code>inseq attribute-context</code></a> CLI command).</p>
 <p>In the demo tab, fill in the input and context fields with the text you want to analyze, and click the <code>Run PECoRe</code> button to produce an output where the tokens selected by PECoRe in the model generation and context are highlighted. For more details on the parameters and their meaning, check the <code>Parameters</code> tab.</p>
 
-<h3>Interpreting PECoRe results</h3>
+<h2>Interpreting PECoRe results</h3>
 """
 
 example_explanation = """
-<p>The example shows the output of the <a href='https://huggingface.co/gsarti/cora_mgen' target='_blank'>CORA Multilingual QA</a> model used as default in the interface, using default settings.</p>
-<p>
-
+<p>Consider the following example, showing inputs and outputs of the <a href='https://huggingface.co/gsarti/cora_mgen' target='_blank'>CORA Multilingual QA</a> model provided as default in the interface, using default settings.</p>
+<img src="file/img/pecore_ui_output_example.png" width=100% />
+<p>The PECoRe CTI step identified two context-sensitive tokens in the generation (<code>287</code> and <code>,</code>), while the CCI step associated each of those with the most influential tokens in the context. It can be observed that in both cases similar tokens from the passage stating the number of inhabitants are identified as salient (<code>235</code> and <code>,</code> for the generated <code>287</code>, while <code>had</code> is also found salient for the generated <code>,</code>).</p>
+<h2>Usage tips</h3>
+<ol>
+    <li>The <code>📂 Download output</code> button allows you to download the full JSON output produced by the Inseq CLI. It includes, among other things, the full set of CTI and CCI scores produced by PECoRe, tokenized versions of the input context and generated output and the full arguments used for the CLI call.</li>
+    <li>The <code>🔍 Download HTML</code> button allows you to download an HTML view of the output similar to the one visualized in the demo.
+    <li>By default, all generated tokens <b>above the mean CTI score</b> for the generated text are highlighted as context-sensitive. This might be reasonable for short answers, but the threshold can be raised by increasing the <code>Context sensitivity threshold</code> parameter to ensure only very sensitive tokens are picked up in longer replies.</li>
+    <li>Relatedly, all context tokens receiving <b>CCI scores >2 standard deviations</b> above the context mean are highlighted as influential. This might be reasonable for contexts with at least 50-100 tokens, but the threshold can be lowered by decreasing the <code>Attribution threshold</code> parameter to be more lenient in the selection for shorter contexts.</li>
+    <li>When using a model, make sure that the <b>contextual and contextless templates are set to match the expected format</b>. You can use presets to auto-fill these for the provided models.</li>
+    <li>If you are using an encoder-decoder expecting an output context (e.g. the multilingual MT preset), the <b>output context should be provided manually</b> before running PECoRe in the <code>Generation context</code> parameter. This is a requirement for the demo because the splitting between output context and current cannot be reliably performed in an automatic way. However, the <code>inseq attribute-context</code> CLI command actually support various strategies, including prompting users for a split and/or trying an automatic source-target alignment. </li>
+</ol>
+<h2>Using PECoRe from Python with Inseq</h3>
+<p>This demo is useful for testing out various models and methods for PECoRe attribution, but the <a href="https://inseq.org/en/latest/main_classes/cli.html#attribute-context"><code>inseq attribute-context</code></a> CLI command is the way to go if you want to run experiments on several examples, or if you want to exploit the full customizability of the Inseq API.</p>
+<p>The utility we provide in this section allows you to generate Python and Shell code calling the Inseq CLI with the parameters you set in the interface. This is useful to understand how to use the Inseq API and quickly get up to speed with running PECoRe on your own models and data.</p>
+<p>Once you are satisfied with the parameters you set (including context/query strings in the <code>🐑 Demo</code> tab), just press the button and get your code snippets ready for usage! 🤗</p>
 """
 
 citation = r"""

presets.py

@@ -1,3 +1,5 @@
+import json
+
 SYSTEM_PROMPT = "You are a helpful assistant that provide concise and accurate answers."
 
 def set_cora_preset():
@@ -77,4 +79,97 @@ def set_mistral_instruct_preset():
         "[INST]{context}\n{current}[/INST]" # input_template
         "[INST]{current}[/INST]" # input_current_text_template
         "\n" # decoder_input_output_separator
-    )
+    )
+
+def update_code_snippets_fn(
+    input_current_text: str,
+    input_context_text: str,
+    output_current_text: str,
+    output_context_text: str,
+    model_name_or_path: str,
+    attribution_method: str,
+    attributed_fn: str | None,
+    context_sensitivity_metric: str,
+    context_sensitivity_std_threshold: float,
+    context_sensitivity_topk: int,
+    attribution_std_threshold: float,
+    attribution_topk: int,
+    input_template: str,
+    output_template: str,
+    contextless_input_template: str,
+    contextless_output_template: str,
+    special_tokens_to_keep: str | list[str] | None,
+    decoder_input_output_separator: str,
+    model_kwargs: str,
+    tokenizer_kwargs: str,
+    generation_kwargs: str,
+    attribution_kwargs: str,
+) -> tuple[str, str]:
+    def get_kwargs_str(kwargs: str, name: str, pad: str = " " * 4) -> str:
+        kwargs_dict = json.loads(kwargs)
+        return nl + pad + name + '=' + str(kwargs_dict) + ',' if kwargs_dict else ''
+    nl = "\n"
+    tq = "\"\"\""
+    # Python
+    python = f"""#!pip install inseq
+import inseq
+from inseq.commands.attribute_context import attribute_context_with_model
+
+inseq_model = inseq.load_model(
+    "{model_name_or_path}",
+    "{attribution_method}",{get_kwargs_str(model_kwargs, "model_kwargs")}{get_kwargs_str(tokenizer_kwargs, "tokenizer_kwargs")}
+)
+
+pecore_args = AttributeContextArgs(
+    save_path="pecore_output.json",
+    viz_path="pecore_output.html",
+    model_name_or_path="{model_name_or_path}",
+    attribution_method="{attribution_method}",
+    attributed_fn="{attributed_fn}",
+    context_sensitivity_metric="{context_sensitivity_metric}",
+    special_tokens_to_keep={special_tokens_to_keep},
+    context_sensitivity_std_threshold={context_sensitivity_std_threshold},
+    attribution_std_threshold={attribution_std_threshold},
+    input_current_text=\"\"\"{input_current_text}\"\"\",
+    input_template=\"\"\"{input_template}\"\"\",
+    output_template="{output_template}",
+    contextless_input_current_text=\"\"\"{contextless_input_template}\"\"\",
+    contextless_output_current_text=\"\"\"{contextless_output_template}\"\"\",
+    context_sensitivity_topk={context_sensitivity_topk if context_sensitivity_topk > 0 else None},
+    attribution_topk={attribution_topk if attribution_topk > 0 else None},
+    input_context_text={tq + input_context_text + tq if input_context_text else None},
+    output_context_text={tq + output_context_text + tq if output_context_text else None},
+    output_current_text={tq + output_current_text + tq if output_current_text else None},
+    decoder_input_output_separator={tq + decoder_input_output_separator + tq if decoder_input_output_separator else None},{get_kwargs_str(model_kwargs, "model_kwargs")}{get_kwargs_str(tokenizer_kwargs, "tokenizer_kwargs")}{get_kwargs_str(generation_kwargs, "generation_kwargs")}{get_kwargs_str(attribution_kwargs, "attribution_kwargs")}
+)
+out = attribute_context_with_model(pecore_args, loaded_model)"""
+    # Bash
+    bash = f"""pip install inseq
+inseq attribute-context \\
+    --save-path pecore_output.json \\
+    --viz-path pecore_output.html \\
+    --model-name-or-path "{model_name_or_path}" \\
+    --attribution-method "{attribution_method}" \\
+    --attributed-fn "{attributed_fn}" \\
+    --context-sensitivity-metric "{context_sensitivity_metric}" \\
+    --special-tokens-to-keep {" ".join(special_tokens_to_keep)} \\
+    --context-sensitivity-std-threshold {context_sensitivity_std_threshold} \\
+    --attribution-std-threshold {attribution_std_threshold} \\
+    --input-current-text "{input_current_text}" \\
+    --input-template "{input_template}" \\
+    --output-template "{output_template}" \\
+    --contextless-input-current-text "{contextless_input_template}" \\
+    --contextless-output-current-text "{contextless_output_template}" \\
+    --context-sensitivity-topk {context_sensitivity_topk if context_sensitivity_topk > 0 else None} \\
+    --attribution-topk {attribution_topk if attribution_topk > 0 else None} \\
+    --input-context-text "{input_context_text}" \\
+    --output-context-text "{output_context_text}" \\
+    --output-current-text "{output_current_text}" \\
+    --decoder-input-output-separator "{decoder_input_output_separator}" \\
+    --model-kwargs "{str(model_kwargs).replace(nl, "")}" \\
+    --tokenizer-kwargs "{str(tokenizer_kwargs).replace(nl, "")} \\
+    --generation-kwargs "{str(generation_kwargs).replace(nl, "")}" \\
+    --attribution-kwargs "{str(attribution_kwargs).replace(nl, "")}"
+    """
+    return python, bash
+    

style.py

@@ -38,4 +38,8 @@ custom_css = """
 .footer-custom-block a {
     margin-right: 15px;
 }
+
+ol {
+    padding-left: 30px;
+}
 """