From fb4cbf082f318e8216dcc937f6c32ce8e49653fe Mon Sep 17 00:00:00 2001 From: xadupre Date: Mon, 23 Jun 2025 15:10:32 +0200 Subject: [PATCH] handle specific models --- onnx_diagnostic/torch_models/hghub/hub_api.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/onnx_diagnostic/torch_models/hghub/hub_api.py b/onnx_diagnostic/torch_models/hghub/hub_api.py index 4d82ea92..cb450cb7 100644 --- a/onnx_diagnostic/torch_models/hghub/hub_api.py +++ b/onnx_diagnostic/torch_models/hghub/hub_api.py @@ -206,6 +206,10 @@ def task_from_id( data = load_architecture_task() if model_id in data: return data[model_id] + if not config.architectures or not config.architectures: + # Some hardcoded values until a better solution is found. + if model_id.startswith("google/bert_"): + return "fill-mask" assert config.architectures is not None and len(config.architectures) == 1, ( f"Cannot return the task of {model_id!r}, pipeline_tag is not setup, " f"architectures={config.architectures} in config={config}. "