add gradio app

app.py

@@ -0,0 +1,274 @@
+import gradio as gr
+from datasets import disable_caching, load_dataset
+from transformer_ranker import TransformerRanker, prepare_popular_models
+import traceback
+
+from utils import (
+    DISABLED_BUTTON_VARIANT, ENABLED_BUTTON_VARIANT, CSS, HEADLINE, FOOTER,
+    EmbeddingProgressTracker, check_dataset_exists, check_dataset_is_loaded,
+    compute_ratio, ensure_one_lm_selected, get_dataset_info
+)
+
+disable_caching()
+
+THEME = "pseudolab/huggingface-korea-theme"
+DEFAULT_SAMPLES = 1000
+MAX_SAMPLES = 5000
+LANGUAGE_MODELS = prepare_popular_models('base') + prepare_popular_models('large')
+
+# Add a tiny model for demonstration on CPU
+LANGUAGE_MODELS = ['prajjwal1/bert-tiny'] + list(dict.fromkeys(LANGUAGE_MODELS))
+LANGUAGE_MODELS.insert(LANGUAGE_MODELS.index("bert-base-cased") + 1, "bert-base-uncased")
+
+# Preselect some small models
+DEFAULT_MODELS = [
+    "prajjwal1/bert-tiny", "google/electra-small-discriminator", 
+    "distilbert-base-cased", "sentence-transformers/all-MiniLM-L12-v2"
+]
+
+
+with gr.Blocks(css=CSS, theme=THEME) as demo:
+
+    ########## STEP 1: Load the Dataset ##########
+
+    gr.Markdown(HEADLINE)
+
+    gr.Markdown("## Step 1: Load a Dataset")
+    with gr.Group():
+        dataset = gr.State(None)
+
+        dataset_name = gr.Textbox(
+            label="Enter the name of your dataset",
+            placeholder="Examples: trec, ag_news, sst2, conll2003, leondz/wnut_17",
+            max_lines=1,
+        )
+        select_dataset_button = gr.Button(
+            value="Load dataset", interactive=False, variant=DISABLED_BUTTON_VARIANT
+        )
+
+        # Activate the "Load dataset" button if dataset was found
+        dataset_name.change(
+            check_dataset_exists, inputs=dataset_name, outputs=select_dataset_button
+        )
+
+    gr.Markdown(
+        "*The number of samples that can be used in this demo is limited to save resources. "
+        "To run an estimate on the full dataset, check out the "
+        "[library](https://github.com/flairNLP/transformer-ranker).*"
+    )
+
+    ########## Step 1.1 Dataset preprocessing ##########
+
+    with gr.Accordion("Dataset settings", open=False) as dataset_config:
+        with gr.Row() as dataset_details:
+            dataset_name_label = gr.Label("", label="Dataset Name")
+            num_samples = gr.State(0)
+            num_samples_label = gr.Label("", label="Number of Samples")
+            num_samples.change(
+                lambda x: str(x), inputs=[num_samples], outputs=[num_samples_label]
+            )
+
+        with gr.Row():
+            text_column = gr.Dropdown("", label="Text Column")
+            text_pair_column = gr.Dropdown("", label="Text Pair Column")
+
+        with gr.Row():
+            label_column = gr.Dropdown("", label="Label Column")
+            task_category = gr.Dropdown("", label="Task Type")
+
+        with gr.Group():
+            downsample_ratio = gr.State(0.0)
+            num_samples_to_use = gr.Slider(
+                20, MAX_SAMPLES, label="Samples to use", value=DEFAULT_SAMPLES, step=1
+            )
+            downsample_ratio_label = gr.Label("", label="Ratio of dataset to use")
+            downsample_ratio.change(
+                lambda x: f"{x:.1%}",
+                inputs=[downsample_ratio],
+                outputs=[downsample_ratio_label],
+            )
+
+            num_samples_to_use.change(
+                compute_ratio,
+                inputs=[num_samples_to_use, num_samples],
+                outputs=downsample_ratio,
+            )
+            num_samples.change(
+                compute_ratio,
+                inputs=[num_samples_to_use, num_samples],
+                outputs=downsample_ratio,
+            )
+
+    # Download the dataset and show details
+    def select_dataset(dataset_name):
+        try:
+            dataset = load_dataset(dataset_name, trust_remote_code=True)
+            dataset_info = get_dataset_info(dataset)
+        except ValueError:
+            gr.Warning("Dataset collections are not supported. Please use a single dataset.")
+
+        return (
+            gr.update(value="Loaded", interactive=False, variant=DISABLED_BUTTON_VARIANT),
+            gr.Accordion(open=True),
+            dataset_name,
+            dataset,
+            *dataset_info
+        )
+
+    select_dataset_button.click(
+        select_dataset,
+        inputs=[dataset_name],
+        outputs=[
+            select_dataset_button,
+            dataset_config,
+            dataset_name_label,
+            dataset,
+            task_category,
+            text_column,
+            text_pair_column,
+            label_column,
+            num_samples,
+        ],
+        scroll_to_output=True,
+    )
+
+    ########## STEP 2 ##########
+
+    gr.Markdown("## Step 2: Select a List of Language Models")
+    with gr.Group():
+        model_options = [
+            (model_handle.split("/")[-1], model_handle)
+            for model_handle in LANGUAGE_MODELS
+        ]
+        models = gr.CheckboxGroup(
+            choices=model_options, label="Select Models", value=DEFAULT_MODELS
+        )
+
+    ########## STEP 3: Run Language Model Ranking ##########
+
+    gr.Markdown("## Step 3: Rank LMs")
+
+    with gr.Group():
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                estimator = gr.Dropdown(
+                    choices=["hscore", "logme", "knn"],
+                    label="Transferability metric",
+                    value="hscore",
+                )
+                layer_pooling_options = ["lastlayer", "layermean", "bestlayer"]
+                layer_pooling = gr.Dropdown(
+                    choices=["lastlayer", "layermean", "bestlayer"],
+                    label="Layer pooling",
+                    value="layermean",
+                )
+        submit_button = gr.Button("Run Ranking", interactive=False, variant=DISABLED_BUTTON_VARIANT)
+
+        # Make button active if the dataset is loaded
+        dataset.change(
+            check_dataset_is_loaded,
+            inputs=[dataset, text_column, label_column, task_category],
+            outputs=submit_button
+        )
+
+        label_column.change(
+            check_dataset_is_loaded,
+            inputs=[dataset, text_column, label_column, task_category],
+            outputs=submit_button
+        )
+
+        text_column.change(
+            check_dataset_is_loaded,
+            inputs=[dataset, text_column, label_column, task_category],
+            outputs=submit_button
+        )
+
+    def rank_models(
+        dataset,
+        downsample_ratio,
+        selected_models,
+        layer_pooling,
+        estimator,
+        text_column,
+        text_pair_column,
+        label_column,
+        task_category,
+        progress=gr.Progress(),
+    ):
+
+        if text_column == "-":
+            raise gr.Error("Text column is not set.")
+
+        if label_column == "-":
+            raise gr.Error("Label column is not set.")
+
+        if task_category == "-":
+            raise gr.Error(
+                "Task category is not set. The dataset must support classification or regression tasks."
+            )
+
+        if text_pair_column == "-":
+            text_pair_column = None
+
+        progress(0.0, "Starting")
+
+        with EmbeddingProgressTracker(progress=progress, model_names=selected_models) as tracker:
+            try:
+                ranker = TransformerRanker(
+                    dataset,
+                    dataset_downsample=downsample_ratio,
+                    text_column=text_column,
+                    text_pair_column=text_pair_column,
+                    label_column=label_column,
+                    task_category=task_category,
+                )
+
+                results = ranker.run(
+                    models=selected_models,
+                    layer_aggregator=layer_pooling,
+                    estimator=estimator,
+                    batch_size=64,
+                    tracker=tracker,
+                )
+
+                sorted_results = sorted(
+                    results._results.items(), key=lambda item: item[1], reverse=True
+                )
+                return [
+                    (i + 1, model, score) for i, (model, score) in enumerate(sorted_results)
+                ]
+            except Exception as e:
+                gr.Error("The dataset is not supported.")
+
+    gr.Markdown("## Results")
+    ranking_results = gr.Dataframe(
+        headers=["Rank", "Model", "Score"], datatype=["number", "str", "number"]
+    )
+
+    submit_button.click(
+        rank_models,
+        inputs=[
+            dataset,
+            downsample_ratio,
+            models,
+            layer_pooling,
+            estimator,
+            text_column,
+            text_pair_column,
+            label_column,
+            task_category,
+        ],
+        outputs=ranking_results,
+        scroll_to_output=True,
+    )
+
+    gr.Markdown(
+        "*The results are ranked by their transferability score, with the most suitable model listed first. "
+        "This ranking allows focusing on the higher-ranked models for further exploration and fine-tuning.*"
+    )
+
+    gr.Markdown(FOOTER)
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=3)
+    demo.launch(max_threads=6)

utils.py

@@ -0,0 +1,205 @@
+import gradio as gr
+from datasets import concatenate_datasets
+from huggingface_hub import HfApi
+from huggingface_hub.errors import HFValidationError
+from requests.exceptions import HTTPError
+from transformer_ranker import Result
+from transformer_ranker.datacleaner import DatasetCleaner, TaskCategory
+from transformer_ranker.embedder import Embedder
+import math
+
+DISABLED_BUTTON_VARIANT = "huggingface"
+ENABLED_BUTTON_VARIANT = "primary"
+
+HEADLINE = """
+<h1 align="center">TransformerRanker</h1>
+<p align="center" style="max-width: 560px; margin: auto;">
+    A very simple library that helps you find the best-suited language model for your NLP task.
+    All you need to do is to select a dataset and a list of pre-trained language models (LMs) from the 🤗 HuggingFace Hub.
+    TransformerRanker will quickly estimate which of these LMs will perform best on the given dataset!
+</p>
+<p align="center" style="font-weight: bold; margin-top: 20px; display: flex; justify-content: center; gap: 10px;">
+    <a href="https://github.com/flairNLP/transformer-ranker">
+        <img src="https://img.shields.io/github/stars/flairNLP/transformer-ranker?style=social&label=Repository" alt="GitHub Badge">
+    </a>
+    <a href="https://pypi.org/project/transformer-ranker/">
+        <img src="https://img.shields.io/badge/Package-orange?style=flat&logo=python" alt="Package Badge">
+    </a>
+    <a href="https://github.com/flairNLP/transformer-ranker/blob/main/examples/01-walkthrough.md">
+        <img src="https://img.shields.io/badge/Tutorials-blue?style=flat&logo=readthedocs&logoColor=white" alt="Tutorials Badge">
+    </a>
+    <img src="https://img.shields.io/badge/license-MIT-green?style=flat" alt="License: MIT">
+</p>
+<p align="center">Developed at <a href="https://www.informatik.hu-berlin.de/en/forschung-en/gebiete/ml-en/">Humboldt University of Berlin</a>.</p>
+"""
+
+FOOTER = """
+**Note:** This demonstration currently runs on a CPU and is suited for smaller models only.  
+**Developers:** [@plonerma](https://huggingface.co/plonerma) and [@lukasgarbas](https://huggingface.co/lukasgarbas). 
+For feedback, suggestions, or contributions, reach out via GitHub or leave a message in the [discussions](https://huggingface.co/spaces/lukasgarbas/transformer-ranker/discussions).
+"""
+
+CSS = """
+.gradio-container{max-width: 800px !important}
+a {color: #ff9d00;}
+@media (prefers-color-scheme: dark) { a {color: #be185d;} }
+"""
+
+
+hf_api = HfApi()
+
+
+def check_dataset_exists(dataset_name):
+    """Update loading button if dataset can be found"""
+    try:
+        hf_api.dataset_info(dataset_name)
+        return gr.update(interactive=True, variant=ENABLED_BUTTON_VARIANT)
+
+    except (HTTPError, HFValidationError):
+        return gr.update(value="Load dataset", interactive=False, variant=DISABLED_BUTTON_VARIANT)
+
+def check_dataset_is_loaded(dataset, text_column, label_column, task_category):
+    if dataset and text_column != "-" and label_column != "-" and task_category != "-":
+        return gr.update(interactive=True, variant=ENABLED_BUTTON_VARIANT)
+    else:
+        return gr.update(interactive=False, variant=DISABLED_BUTTON_VARIANT)
+
+
+def get_dataset_info(dataset):
+    """Show information for dataset settings"""
+    joined_dataset = concatenate_datasets(list(dataset.values()))
+    datacleaner = DatasetCleaner()
+
+    try:
+        text_column = datacleaner._find_column(joined_dataset, "text column")
+    except ValueError:
+        gr.Warning("Text column can not be found. Select it in the dataset settings.")
+        text_column = "-"
+
+    try:
+        label_column = datacleaner._find_column(joined_dataset, "label column")
+    except ValueError:
+        gr.Warning("Label column can not be found. Select it in the dataset settings.")
+        label_column = "-"
+
+    task_category = "-"
+    if label_column != "-":
+        try:
+            # Find or set the task_category
+            task_category = datacleaner._find_task_category(joined_dataset, label_column)
+        except ValueError:
+            gr.Warning(
+                "Task category could not be determined. The dataset must support classification or regression tasks.",
+            )
+            pass
+
+    num_samples = len(joined_dataset)
+
+    return (
+        gr.update(
+            value=task_category,
+            choices=[str(t) for t in TaskCategory],
+            interactive=True,
+        ),
+        gr.update(
+            value=text_column, choices=joined_dataset.column_names, interactive=True
+        ),
+        gr.update(
+            value="-", choices=["-", *joined_dataset.column_names], interactive=True
+        ),
+        gr.update(
+            value=label_column, choices=joined_dataset.column_names, interactive=True
+        ),
+        num_samples,
+    )
+
+
+def compute_ratio(num_samples_to_use, num_samples):
+    if num_samples > 0:
+        return num_samples_to_use / num_samples
+    else:
+        return 0.0
+
+
+def ensure_one_lm_selected(checkbox_values, previous_values):
+    if not any(checkbox_values):
+        return previous_values
+    return checkbox_values
+
+
+# Apply monkey patch to enable callbacks
+_old_embed = Embedder.embed
+
+def _new_embed(embedder, sentences, batch_size: int = 32, **kw):
+    if embedder.tracker is not None:
+        embedder.tracker.update_num_batches(math.ceil(len(sentences) / batch_size))
+
+    return _old_embed(embedder, sentences, batch_size=batch_size, **kw)
+
+Embedder.embed = _new_embed
+
+_old_embed_batch = Embedder.embed_batch
+
+def _new_embed_batch(embedder, *args, **kw):
+    r = _old_embed_batch(embedder, *args, **kw)
+    if embedder.tracker is not None:
+        embedder.tracker.update_batch_complete()
+    return r
+
+Embedder.embed_batch = _new_embed_batch
+
+_old_init = Embedder.__init__
+
+def _new_init(embedder, *args, tracker=None, **kw):
+    _old_init(embedder, *args, **kw)
+    embedder.tracker = tracker
+
+Embedder.__init__ = _new_init
+
+
+class EmbeddingProgressTracker:
+    def __init__(self, *, progress, model_names):
+        self.model_names = model_names
+        self.progress_bar = progress
+
+    @property
+    def total(self):
+        return len(self.model_names)
+
+    def __enter__(self):
+        self.progress_bar = gr.Progress(track_tqdm=False)
+        self.current_model = -1
+        self.batches_complete = 0
+        self.batches_total = None
+        return self
+
+    def __exit__(self, typ, value, tb):
+        if typ is None:
+            self.progress_bar(1.0, desc="Done")
+        else:
+            self.progress_bar(1.0, desc="Error")
+
+        # Do not suppress any errors
+        return False
+
+    def update_num_batches(self, total):
+        self.current_model += 1
+        self.batches_complete = 0
+        self.batches_total = total
+        self.update_bar()
+
+    def update_batch_complete(self):
+        self.batches_complete += 1
+        self.update_bar()
+
+    def update_bar(self):
+        i = self.current_model
+
+        description = f"Running {self.model_names[i]} ({i + 1} / {self.total})"
+
+        progress = i / self.total
+        if self.batches_total is not None:
+            progress += (self.batches_complete / self.batches_total) / self.total
+
+        self.progress_bar(progress=progress, desc=description)
+