Select pipeline_tag task by default.

app.py

@@ -63,12 +63,26 @@ tasks_mapping = {
     "semantic-segmentation": "Semantic Segmentation",
     "seq2seq-lm": "Text to Text Generation",
     "sequence-classification": "Text Classification",
-    "speech-seq2seq": "Speech to Speech Generation",
+    "speech-seq2seq": "Audio to Audio",
     "token-classification": "Token Classification",
 }
 reverse_tasks_mapping = {v: k for k, v in tasks_mapping.items()}
 tasks_labels = list(tasks_mapping.keys())
 
+# Map pipeline_tag to internal exporters features/tasks
+tags_to_tasks_mapping = {
+    "feature-extraction": "default",
+    "text-generation": "causal-lm",
+    "image-classification": "image-classification",
+    "image-segmentation": "image-segmentation",
+    "fill-mask": "masked-lm",
+    "object-detection": "object-detection",
+    "question-answering": "question-answering",
+    "text2text-generation": "seq2seq-lm",
+    "text-classification": "sequence-classification",
+    "token-classification": "token-classification",
+}
+
 def error_str(error, title="Error", model=None, task=None, framework=None, compute_units=None, precision=None, tolerance=None, destination=None, open_discussion=True):
     if not error: return ""
 
@@ -112,8 +126,18 @@ def get_pr_url(api, repo_id, title):
             and discussion.title == title
         ):
             return f"https://huggingface.co/{repo_id}/discussions/{discussion.num}"
-        
-def supported_frameworks(model_id):
+
+def retrieve_model_info(model_id):
+    api = HfApi()
+    model_info = api.model_info(model_id)
+    tags = model_info.tags
+    frameworks = [tag for tag in tags if tag in ["pytorch", "tf"]]
+    return {
+        "pipeline_tag": model_info.pipeline_tag,
+        "frameworks": sorted(["PyTorch" if f == "pytorch" else "TensorFlow" for f in frameworks]),
+    }
+
+def supported_frameworks(model_info):
     """
     Return a list of supported frameworks (`PyTorch` or `TensorFlow`) for a given model_id.
     Only PyTorch and Tensorflow are supported.
@@ -130,6 +154,7 @@ def on_model_change(model):
     error = None
     frameworks = []
     selected_framework = None
+    selected_task = None
 
     try:
         config_file = hf_hub_download(model, filename="config.json")
@@ -144,17 +169,30 @@ def on_model_change(model):
 
         features = FeaturesManager.get_supported_features_for_model_type(model_type)
         tasks = list(features.keys())
-        tasks = [tasks_mapping[task] for task in tasks]
 
-        frameworks = supported_frameworks(model)
+        model_info = retrieve_model_info(model)
+        frameworks = model_info["frameworks"]
         selected_framework = frameworks[0] if len(frameworks) > 0 else None
+
+        pipeline_tag = model_info["pipeline_tag"]
+        # Select the task corresponding to the pipeline tag
+        if tasks:
+            if pipeline_tag in tags_to_tasks_mapping:
+                selected_task = tags_to_tasks_mapping[pipeline_tag]
+            else:
+                selected_task = tasks[0]
+
+            # Convert to UI labels
+            tasks = [tasks_mapping[task] for task in tasks]
+            selected_task = tasks_mapping[selected_task]
+
     except Exception as e:
         error = e
         model_type = None
 
     return (
         gr.update(visible=bool(model_type)),                                                    # Settings column
-        gr.update(choices=tasks, value=tasks[0] if tasks else None),                            # Tasks
+        gr.update(choices=tasks, value=selected_task),                                          # Tasks
         gr.update(visible=len(frameworks)>1, choices=frameworks, value=selected_framework),     # Frameworks
         gr.update(value=error_str(error, model=model)),                                         # Error
     )