Load model after preset

app.py

@@ -24,7 +24,7 @@ from presets import (
 from style import custom_css
 from utils import get_formatted_attribute_context_results
 
-from inseq import list_feature_attribution_methods, list_step_functions, load_model
+from inseq import list_feature_attribution_methods, list_step_functions
 from inseq.commands.attribute_context.attribute_context import (
     AttributeContextArgs,
     attribute_context_with_model,
@@ -65,7 +65,7 @@ def pecore(
         )
     if loaded_model is None or model_name_or_path != loaded_model.model_name:
         gr.Info("Loading model...")
-        loaded_model = load_model(
+        loaded_model = HuggingfaceModel.load(
             model_name_or_path,
             attribution_method,
             model_kwargs=json.loads(model_kwargs),
@@ -130,7 +130,7 @@ def preload_model(
     global loaded_model
     if loaded_model is None or model_name_or_path != loaded_model.model_name:
         gr.Info("Loading model...")
-        loaded_model = load_model(
+        loaded_model = HuggingfaceModel.load(
             model_name_or_path,
             attribution_method,
             model_kwargs=json.loads(model_kwargs),
@@ -192,7 +192,9 @@ with gr.Blocks(css=custom_css) as demo:
             outputs=pecore_output_highlights,
         )
     with gr.Tab("⚙️ Parameters") as params_tab:
-        gr.Markdown("## ✨ Presets")
+        gr.Markdown(
+            "## ✨ Presets\nSelect a preset to load default parameters into the fields below."
+        )
         with gr.Row(equal_height=True):
             with gr.Column():
                 default_preset = gr.Button("Default", variant="secondary")
@@ -218,7 +220,7 @@ with gr.Blocks(css=custom_css) as demo:
                     "Present for multilingual MT models such as <a href='https://huggingface.co/facebook/nllb-200-distilled-600M' target='_blank'>NLLB</a> and <a href='https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt' target='_blank'>mBART</a> using language tags."
                 )
             with gr.Column(scale=1):
-                chatml_template = gr.Button("ChatML Template", variant="secondary")
+                chatml_template = gr.Button("Qwen ChatML", variant="secondary")
                 gr.Markdown(
                     "Preset for models using the <a href='https://github.com/MicrosoftDocs/azure-docs/blob/main/articles/ai-services/openai/includes/chat-markup-language.md' target='_blank'>ChatML conversational template</a>.\nUses <code><|im_start|></code>, <code><|im_end|></code> special tokens."
                 )
@@ -401,6 +403,15 @@ with gr.Blocks(css=custom_css) as demo:
     gr.Markdown(how_to_use)
     gr.Markdown(citation)
 
+    # Main logic
+
+    load_model_args = [
+        model_name_or_path,
+        attribution_method,
+        model_kwargs,
+        tokenizer_kwargs,
+    ]
+
     attribute_input_button.click(
         pecore,
         inputs=[
@@ -435,7 +446,7 @@ with gr.Blocks(css=custom_css) as demo:
 
     load_model_button.click(
         preload_model,
-        inputs=[model_name_or_path, attribution_method, model_kwargs, tokenizer_kwargs],
+        inputs=load_model_args,
         outputs=[],
     )
 
@@ -461,11 +472,13 @@ with gr.Blocks(css=custom_css) as demo:
 
     # Presets
 
-    default_preset.click(**reset_kwargs)
+    default_preset.click(**reset_kwargs).success(preload_model, inputs=load_model_args)
+
     cora_preset.click(**reset_kwargs).then(
         set_cora_preset,
         outputs=[model_name_or_path, input_template, contextless_input_current_text],
-    )
+    ).success(preload_model, inputs=load_model_args)
+
     zephyr_preset.click(**reset_kwargs).then(
         set_zephyr_preset,
         outputs=[
@@ -474,11 +487,13 @@ with gr.Blocks(css=custom_css) as demo:
             contextless_input_current_text,
             decoder_input_output_separator,
         ],
-    )
+    ).success(preload_model, inputs=load_model_args)
+
     multilingual_mt_template.click(**reset_kwargs).then(
         set_mmt_preset,
         outputs=[model_name_or_path, input_template, output_template, tokenizer_kwargs],
-    )
+    ).success(preload_model, inputs=load_model_args)
+
     chatml_template.click(**reset_kwargs).then(
         set_chatml_preset,
         outputs=[
@@ -488,7 +503,8 @@ with gr.Blocks(css=custom_css) as demo:
             decoder_input_output_separator,
             special_tokens_to_keep,
         ],
-    )
+    ).success(preload_model, inputs=load_model_args)
+
     towerinstruct_template.click(**reset_kwargs).then(
         set_towerinstruct_preset,
         outputs=[
@@ -497,6 +513,6 @@ with gr.Blocks(css=custom_css) as demo:
             contextless_input_current_text,
             decoder_input_output_separator,
         ],
-    )
+    ).success(preload_model, inputs=load_model_args)
 
 demo.launch(allowed_paths=["outputs/"])