Skip to content

Commit 7553646

Browse files
authored
Partial revert #2851: allow task override on sentence-similarity (#2861)
1 parent 166f174 commit 7553646

File tree

3 files changed

+9
-3
lines changed

3 files changed

+9
-3
lines changed

src/huggingface_hub/inference/_providers/hf_inference.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,13 @@ def _prepare_url(self, api_key: str, mapped_model: str) -> str:
3838
# hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment)
3939
if mapped_model.startswith(("http://", "https://")):
4040
return mapped_model
41-
return f"{self.base_url}/models/{mapped_model}"
41+
return (
42+
# Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks.
43+
f"{self.base_url}/pipeline/{self.task}/{mapped_model}"
44+
if self.task in ("feature-extraction", "sentence-similarity")
45+
# Otherwise, we use the default endpoint
46+
else f"{self.base_url}/models/{mapped_model}"
47+
)
4248

4349
def _prepare_payload(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
4450
if isinstance(inputs, bytes):

tests/cassettes/TestInferenceClient.test_sentence_similarity[hf-inference,sentence-similarity].yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ interactions:
1717
X-Amzn-Trace-Id:
1818
- 0434ff33-56fe-49db-9380-17b81e41f756
1919
method: POST
20-
uri: https://router.huggingface.co/hf-inference/models/sentence-transformers/all-MiniLM-L6-v2
20+
uri: https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/all-MiniLM-L6-v2
2121
response:
2222
body:
2323
string: '[0.7785724997520447,0.4587624967098236,0.29062220454216003]'

tests/cassettes/test_async_sentence_similarity.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ interactions:
33
body: null
44
headers: {}
55
method: POST
6-
uri: https://router.huggingface.co/hf-inference/models/sentence-transformers/all-MiniLM-L6-v2
6+
uri: https://router.huggingface.co/hf-inference/pipeline/sentence-similarity/sentence-transformers/all-MiniLM-L6-v2
77
response:
88
body:
99
string: '[0.7785724997520447,0.4587624967098236,0.29062220454216003]'

0 commit comments

Comments
 (0)