Skip to content

Commit 448fbb5

Browse files
committed
Fix 'sentence-transformers/all-MiniLM-L6-v2' doesn't support task 'feature-extraction' (#2968)
1 parent fefa7cc commit 448fbb5

File tree

2 files changed

+23
-0
lines changed

2 files changed

+23
-0
lines changed

src/huggingface_hub/inference/_providers/hf_inference.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,14 @@ def _check_supported_task(model: str, task: str) -> None:
151151
return # Only conversational allowed if tagged as conversational
152152
raise ValueError("Non-conversational image-text-to-text task is not supported.")
153153

154+
if (
155+
task in ("feature-extraction", "sentence-similarity")
156+
and pipeline_tag in ("feature-extraction", "sentence-similarity")
157+
and task in tags
158+
):
159+
# feature-extraction and sentence-similarity are interchangeable for HF Inference
160+
return
161+
154162
# For all other tasks, just check pipeline tag
155163
if pipeline_tag != task:
156164
raise ValueError(

tests/test_inference_providers.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -585,6 +585,21 @@ def test_prepare_payload_as_dict_conversational(self, mapped_model, parameters,
585585
"conversational",
586586
True,
587587
),
588+
# Feature-extraction / sentence-similarity are interchangeable for HF Inference
589+
(
590+
"sentence-similarity",
591+
["tag1", "feature-extraction", "sentence-similarity"],
592+
"feature-extraction",
593+
False,
594+
),
595+
(
596+
"feature-extraction",
597+
["tag1", "feature-extraction", "sentence-similarity"],
598+
"sentence-similarity",
599+
False,
600+
),
601+
# if pipeline_tag is not feature-extraction or sentence-similarity, raise
602+
("text-generation", ["tag1", "feature-extraction", "sentence-similarity"], "sentence-similarity", True),
588603
# Other tasks
589604
(
590605
"audio-classification",

0 commit comments

Comments
 (0)