Fix bugs with CORA and Qwen

app.py

@@ -17,7 +17,6 @@ from contents import (
     subtitle,
     title,
     powered_by,
-    support,
 )
 from gradio_highlightedtextbox import HighlightedTextbox
 from gradio_modal import Modal
@@ -82,6 +81,8 @@ def pecore(
             model_kwargs=json.loads(model_kwargs),
             tokenizer_kwargs=json.loads(tokenizer_kwargs),
         )
+    if loaded_model.tokenizer.pad_token is None:
+        loaded_model.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
     kwargs = {}
     if context_sensitivity_topk > 0:
         kwargs["context_sensitivity_topk"] = context_sensitivity_topk
@@ -160,6 +161,8 @@ def preload_model(
             model_kwargs=json.loads(model_kwargs),
             tokenizer_kwargs=json.loads(tokenizer_kwargs),
         )
+    if loaded_model.tokenizer.pad_token is None:
+        loaded_model.tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
 
 with gr.Blocks(css=custom_css) as demo:
@@ -372,7 +375,7 @@ with gr.Blocks(css=custom_css) as demo:
         gr.Markdown("#### Text Format Parameters")
         with gr.Row(equal_height=True):
             input_template = gr.Textbox(
-                value="<Q>:{current} <P>:{context}",
+                value="<Q>: {current} <P>: {context}",
                 label="Contextual input template",
                 info="Template to format the input for the model. Use {current} and {context} placeholders for Input Query and Input Context, respectively.",
                 interactive=True,
@@ -384,7 +387,7 @@ with gr.Blocks(css=custom_css) as demo:
                 interactive=True,
             )
             contextless_input_template = gr.Textbox(
-                value="<Q>:{current}",
+                value="<Q>: {current}",
                 label="Contextless input template",
                 info="Template to format the input query in the non-contextual setting. Use {current} placeholder for Input Query.",
                 interactive=True,
@@ -488,25 +491,37 @@ with gr.Blocks(css=custom_css) as demo:
         gr.Markdown("If you use the Inseq implementation of PECoRe (<a href=\"https://inseq.org/en/latest/main_classes/cli.html#attribute-context\"><code>inseq attribute-context</code></a>, including this demo), please also cite:")
         gr.Code(inseq_citation, interactive=False, label="Inseq (Sarti et al., 2023)")
     with gr.Row(elem_classes="footer-container"):
-        gr.Markdown(powered_by)
-        gr.Markdown(support)
+        with gr.Column():
+            gr.Markdown(powered_by)
+        with gr.Column():
+            with gr.Row(elem_classes="footer-custom-block"):
+                with gr.Column(scale=0.25, 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.25, 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.25, 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>""")
+                with gr.Column(scale=0.25, min_width=120):
+                    gr.Markdown("""<a href='https://www.esciencecenter.nl/' target='_blank'><img src="file/img/escience_logo_white_contour.png" width=120px /></a>""")
     with Modal(visible=False) as code_modal:
         gr.Markdown(show_code_modal)
         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.Column(scale=0.5):
+                python_code_snippet = gr.Code(
+                    value="""Generate Python code snippet by pressing the button.""",
+                    language="python",
+                    label="Python",
+                    interactive=False,
+                    show_label=True,
+                )
+            with gr.Column(scale=0.5):
+                shell_code_snippet = gr.Code(
+                    value="""Generate Shell code snippet by pressing the button.""",
+                    language="shell",
+                    label="Shell",
+                    interactive=False,
+                    show_label=True,
+                )
 
     # Main logic
 
@@ -604,8 +619,8 @@ with gr.Blocks(css=custom_css) as demo:
             model_name_or_path,
             input_template,
             contextless_input_template,
-            decoder_input_output_separator,
             special_tokens_to_keep,
+            generation_kwargs,
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 
@@ -620,8 +635,8 @@ with gr.Blocks(css=custom_css) as demo:
             model_name_or_path,
             input_template,
             contextless_input_template,
-            decoder_input_output_separator,
             special_tokens_to_keep,
+            generation_kwargs,
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 
@@ -631,8 +646,8 @@ with gr.Blocks(css=custom_css) as demo:
             model_name_or_path,
             input_template,
             contextless_input_template,
-            decoder_input_output_separator,
             special_tokens_to_keep,
+            generation_kwargs,
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 
@@ -642,8 +657,8 @@ with gr.Blocks(css=custom_css) as demo:
             model_name_or_path,
             input_template,
             contextless_input_template,
-            decoder_input_output_separator,
             special_tokens_to_keep,
+            generation_kwargs,
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 
@@ -653,7 +668,7 @@ with gr.Blocks(css=custom_css) as demo:
             model_name_or_path,
             input_template,
             contextless_input_template,
-            decoder_input_output_separator,
+            generation_kwargs,
         ],
     ).success(preload_model, inputs=load_model_args, cancels=load_model_event)
 

contents.py

@@ -44,7 +44,7 @@ how_to_use = """
 example_explanation = """
 <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>
+<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 the matching tokens stating the number of inhabitants are identified as salient (<code>,</code> and <code>287</code> for the generated <code>287</code>, while <code>235</code> is also found salient for the generated <code>,</code>). In this case, the influential context found by PECoRe is lexically equal to the generated output, but in principle better LMs might not use their inputs verbatim, hence the interest for using model internals with PECoRe.</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>
@@ -95,8 +95,6 @@ inseq_citation = """@inproceedings{sarti-etal-2023-inseq,
 
 powered_by = """<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></div>"""
 
-support = """<div class="footer-custom-block"><b>Built by <a href="https://gsarti.com" target="_blank">Gabriele Sarti</a><br> with the support of</b> <a href='https://www.rug.nl/research/clcg/research/cl/' target='_blank'><img src="file/img/rug_logo_white_contour.png" width=170px /></a><a href='https://projects.illc.uva.nl/indeep/' target='_blank'><img src="file/img/indeep_logo_white_contour.png" width=100px /></a><a href='https://www.esciencecenter.nl/' target='_blank'><img src="file/img/escience_logo_white_contour.png" width=120px /></a></div>"""
-
 examples = [
     [
         "How many inhabitants does Groningen have?",

presets.py

@@ -5,8 +5,8 @@ SYSTEM_PROMPT = "You are a helpful assistant that provide concise and accurate a
 def set_cora_preset():
     return (
         "gsarti/cora_mgen",  # model_name_or_path
-        "<Q>:{current} <P>:{context}",  # input_template
-        "<Q>:{current}",  # input_current_text_template
+        "<Q>: {current} <P>: {context}",  # input_template
+        "<Q>: {current}",  # input_current_text_template
     )
 
 
@@ -29,20 +29,20 @@ def set_default_preset():
 def set_zephyr_preset():
     return (
         "stabilityai/stablelm-2-zephyr-1_6b",  # model_name_or_path
-        "<|system|>{system_prompt}<|endoftext|>\n<|user|>\n{context}\n\n{current}<|endoftext|>\n<|assistant|>".format(system_prompt=SYSTEM_PROMPT),  # input_template
-        "<|system|>{system_prompt}<|endoftext|>\n<|user|>\n{current}<|endoftext|>\n<|assistant|>".format(system_prompt=SYSTEM_PROMPT),  # input_current_text_template
-        "\n",  # decoder_input_output_separator
+        "<|system|>{system_prompt}<|endoftext|>\n<|user|>\n{context}\n\n{current}<|endoftext|>\n<|assistant|>\n".replace("{system_prompt}", SYSTEM_PROMPT),  # input_template
+        "<|system|>{system_prompt}<|endoftext|>\n<|user|>\n{current}<|endoftext|>\n<|assistant|>\n".replace("{system_prompt}", SYSTEM_PROMPT),  # input_current_text_template
         ["<|im_start|>", "<|im_end|>", "<|endoftext|>"],  # special_tokens_to_keep
+        '{\n\t"max_new_tokens": 50\n}',  # generation_kwargs
     )
 
 
 def set_chatml_preset():
     return (
         "Qwen/Qwen1.5-0.5B-Chat",  # model_name_or_path
-        "<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{context}\n\n{current}<|im_end|>\n<|im_start|>assistant".format(system_prompt=SYSTEM_PROMPT),  # input_template
-        "<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{current}<|im_end|>\n<|im_start|>assistant".format(system_prompt=SYSTEM_PROMPT),  # input_current_text_template
-        "\n",  # decoder_input_output_separator
+        "<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{context}\n\n{current}<|im_end|>\n<|im_start|>assistant\n".replace("{system_prompt}", SYSTEM_PROMPT),  # input_template
+        "<|im_start|>system\n{system_prompt}<|im_end|>\n<|im_start|>user\n{current}<|im_end|>\n<|im_start|>assistant\n".replace("{system_prompt}", SYSTEM_PROMPT),  # input_current_text_template
         ["<|im_start|>", "<|im_end|>"],  # special_tokens_to_keep
+        '{\n\t"max_new_tokens": 50\n}',  # generation_kwargs
     )
 
 
@@ -58,19 +58,19 @@ def set_mmt_preset():
 def set_towerinstruct_preset():
     return (
         "Unbabel/TowerInstruct-7B-v0.1",  # model_name_or_path
-        "<|im_start|>user\nSource: {current}\nContext: {context}\nTranslate the above text into French. Use the context to guide your answer.\nTarget:<|im_end|>\n<|im_start|>assistant",  # input_template
-        "<|im_start|>user\nSource: {current}\nTranslate the above text into French.\nTarget:<|im_end|>\n<|im_start|>assistant",  # input_current_text_template
-        "\n",  # decoder_input_output_separator
+        "<|im_start|>user\nSource: {current}\nContext: {context}\nTranslate the above text into French. Use the context to guide your answer.\nTarget:<|im_end|>\n<|im_start|>assistant\n",  # input_template
+        "<|im_start|>user\nSource: {current}\nTranslate the above text into French.\nTarget:<|im_end|>\n<|im_start|>assistant\n",  # input_current_text_template
         ["<|im_start|>", "<|im_end|>"],  # special_tokens_to_keep
+        '{\n\t"max_new_tokens": 50\n}',  # generation_kwargs
     )
 
 def set_gemma_preset():
     return (
         "google/gemma-2b-it", # model_name_or_path
-        "<start_of_turn>user\n{context}\n{current}<end_of_turn>\n<start_of_turn>model", # input_template
-        "<start_of_turn>user\n{current}<end_of_turn>\n<start_of_turn>model", # input_current_text_template
-        "\n", # decoder_input_output_separator
+        "<start_of_turn>user\n{context}\n{current}<end_of_turn>\n<start_of_turn>model\n", # input_template
+        "<start_of_turn>user\n{current}<end_of_turn>\n<start_of_turn>model\n", # input_current_text_template
         ["<start_of_turn>", "<end_of_turn>"], # special_tokens_to_keep
+        '{\n\t"max_new_tokens": 50\n}',  # generation_kwargs
     )
 
 def set_mistral_instruct_preset():
@@ -78,7 +78,7 @@ def set_mistral_instruct_preset():
         "mistralai/Mistral-7B-Instruct-v0.2" # model_name_or_path
         "[INST]{context}\n{current}[/INST]" # input_template
         "[INST]{current}[/INST]" # input_current_text_template
-        "\n" # decoder_input_output_separator
+        '{\n\t"max_new_tokens": 50\n}',  # generation_kwargs
     )
 
 def update_code_snippets_fn(
@@ -137,7 +137,7 @@ def update_code_snippets_fn(
     # Python
     python = f"""#!pip install inseq
 import inseq
-from inseq.commands.attribute_context import attribute_context_with_model
+from inseq.commands.attribute_contex.attribute_context import attribute_context_with_model, AttributeContextArgs
 
 inseq_model = inseq.load_model(
     "{model_name_or_path}",
@@ -160,7 +160,7 @@ pecore_args = AttributeContextArgs(
     viz_path="pecore_output.html",{py_get_kwargs_str(model_kwargs, "model_kwargs")}{py_get_kwargs_str(tokenizer_kwargs, "tokenizer_kwargs")}{py_get_kwargs_str(generation_kwargs, "generation_kwargs")}{py_get_kwargs_str(attribution_kwargs, "attribution_kwargs")}
 )
 
-out = attribute_context_with_model(pecore_args, loaded_model)"""
+out = attribute_context_with_model(pecore_args, inseq_model)"""
     # Bash
     bash = f"""# pip install inseq
 inseq attribute-context \\