Remove the task argument for ExecuTorchModelForXxx (#12)

* remove the need to provide the task argument

* update tests

Author: echarlaix

optimum/executorch/modeling.py

@@ -34,6 +34,7 @@
     _load_for_executorch,
 )
 
+from ..exporters import TasksManager
 from ..exporters.executorch import main_export
 from ..modeling_base import OptimizedModel
 
@@ -119,7 +120,6 @@ def from_pretrained(
         cls,
         model_name_or_path: Union[str, Path],
         export: bool = True,
-        task: str = "",
         recipe: str = "",
         config: "PretrainedConfig" = None,
         subfolder: str = "",
@@ -140,8 +140,6 @@ def from_pretrained(
             export (`bool`, *optional*, defaults to `True`):
                 If `True`, the model will be exported from eager to ExecuTorch after fetched from huggingface.co. `model_name_or_path` must be a valid model ID on huggingface.co.
                 If `False`, the previously exported ExecuTorch model will be loaded from a local path. `model_name_or_path` must be a valid local directory where a `model.pte` is stored.
-            task (`str`, defaults to `""`):
-                The task to export the model for, e.g. "text-generation". It is required to specify a task when `export` is `True`.
             recipe (`str`, defaults to `""`):
                 The recipe to use to do the export, e.g. "xnnpack". It is required to specify a task when `export` is `True`.
             config (`PretrainedConfig`, *optional*):
@@ -180,13 +178,10 @@ def from_pretrained(
 
         if export:
             # Fetch the model from huggingface.co and export it to ExecuTorch
-            if task == "":
-                raise ValueError("Please specify a task to export the model for.")
             if recipe == "":
                 raise ValueError("Please specify a recipe to export the model for.")
             return cls._export(
                 model_id=model_name_or_path,
-                task=task,
                 recipe=recipe,
                 config=config,
                 **kwargs,
@@ -261,7 +256,6 @@ def _save_pretrained(self, save_directory):
     def _export(
         cls,
         model_id: str,
-        task: str,
         recipe: str,
         config: PretrainedConfig,
         cache_dir: str = HUGGINGFACE_HUB_CACHE,
@@ -280,8 +274,6 @@ def _export(
         Args:
             model_id (`str`):
                 Model ID on huggingface.co, for example: `model_name_or_path="meta-llama/Llama-3.2-1B"`.
-            task (`str`):
-                The task to export the model for, e.g. "text-generation".
             recipe (`str`):
                 The recipe to use to do the export, e.g. "xnnpack".
             config (`PretrainedConfig`, *optional*):
@@ -314,6 +306,10 @@ def _export(
             `ExecuTorchModelForCausalLM`: The loaded and exported ExecuTorch model.
 
         """
+        task = kwargs.pop("task", None)
+        if task is not None:
+            logger.warning(f"task was provided and set to {task} but not used, will be ignored")
+
         if use_auth_token is not None:
             warnings.warn(
                 "The `use_auth_token` argument is deprecated and will be removed soon. Please use the `token` argument instead.",
@@ -325,12 +321,11 @@ def _export(
 
         save_dir = TemporaryDirectory()
         save_dir_path = Path(save_dir.name)
-
         # Export to ExecuTorch and save the pte file to the temporary directory
         main_export(
             model_name_or_path=model_id,
             output_dir=save_dir_path,
-            task=task,
+            task=TasksManager.infer_task_from_model(cls.auto_model_class),
             recipe=recipe,
             subfolder=subfolder,
             revision=revision,

tests/models/test_modeling.py

@@ -34,7 +34,6 @@ def test_load_model_from_hub(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path="NousResearch/Llama-3.2-1B",
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)

tests/models/test_modeling_gemma.py

@@ -36,7 +36,6 @@ def test_gemma_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)

tests/models/test_modeling_gemma2.py

@@ -36,7 +36,6 @@ def test_gemma2_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)

tests/models/test_modeling_llama.py

@@ -36,7 +36,6 @@ def test_llama3_2_1b_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)
@@ -61,7 +60,6 @@ def test_llama3_2_3b_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)

tests/models/test_modeling_olmo.py

@@ -34,7 +34,6 @@ def test_olmo_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)

tests/models/test_modeling_qwen2.py

@@ -34,7 +34,6 @@ def test_qwen2_5_text_generation_with_xnnpack(self):
         model = ExecuTorchModelForCausalLM.from_pretrained(
             model_name_or_path=model_id,
             export=True,
-            task="text-generation",
             recipe="xnnpack",
         )
         self.assertIsInstance(model, ExecuTorchModelForCausalLM)