Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
+205
@@ -0,0 +1,205 @@
|
||||
from typing import Dict, Literal, Optional, Union
|
||||
|
||||
from huggingface_hub.inference._providers.featherless_ai import (
|
||||
FeatherlessConversationalTask,
|
||||
FeatherlessTextGenerationTask,
|
||||
)
|
||||
from huggingface_hub.utils import logging
|
||||
|
||||
from ._common import TaskProviderHelper, _fetch_inference_provider_mapping
|
||||
from .black_forest_labs import BlackForestLabsTextToImageTask
|
||||
from .cerebras import CerebrasConversationalTask
|
||||
from .cohere import CohereConversationalTask
|
||||
from .fal_ai import (
|
||||
FalAIAutomaticSpeechRecognitionTask,
|
||||
FalAITextToImageTask,
|
||||
FalAITextToSpeechTask,
|
||||
FalAITextToVideoTask,
|
||||
)
|
||||
from .fireworks_ai import FireworksAIConversationalTask
|
||||
from .groq import GroqConversationalTask
|
||||
from .hf_inference import (
|
||||
HFInferenceBinaryInputTask,
|
||||
HFInferenceConversational,
|
||||
HFInferenceFeatureExtractionTask,
|
||||
HFInferenceTask,
|
||||
)
|
||||
from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
|
||||
from .nebius import (
|
||||
NebiusConversationalTask,
|
||||
NebiusFeatureExtractionTask,
|
||||
NebiusTextGenerationTask,
|
||||
NebiusTextToImageTask,
|
||||
)
|
||||
from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask
|
||||
from .nscale import NscaleConversationalTask, NscaleTextToImageTask
|
||||
from .openai import OpenAIConversationalTask
|
||||
from .replicate import ReplicateTask, ReplicateTextToImageTask, ReplicateTextToSpeechTask
|
||||
from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
|
||||
from .together import TogetherConversationalTask, TogetherTextGenerationTask, TogetherTextToImageTask
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
PROVIDER_T = Literal[
|
||||
"black-forest-labs",
|
||||
"cerebras",
|
||||
"cohere",
|
||||
"fal-ai",
|
||||
"featherless-ai",
|
||||
"fireworks-ai",
|
||||
"groq",
|
||||
"hf-inference",
|
||||
"hyperbolic",
|
||||
"nebius",
|
||||
"novita",
|
||||
"nscale",
|
||||
"openai",
|
||||
"replicate",
|
||||
"sambanova",
|
||||
"together",
|
||||
]
|
||||
|
||||
PROVIDER_OR_POLICY_T = Union[PROVIDER_T, Literal["auto"]]
|
||||
|
||||
PROVIDERS: Dict[PROVIDER_T, Dict[str, TaskProviderHelper]] = {
|
||||
"black-forest-labs": {
|
||||
"text-to-image": BlackForestLabsTextToImageTask(),
|
||||
},
|
||||
"cerebras": {
|
||||
"conversational": CerebrasConversationalTask(),
|
||||
},
|
||||
"cohere": {
|
||||
"conversational": CohereConversationalTask(),
|
||||
},
|
||||
"fal-ai": {
|
||||
"automatic-speech-recognition": FalAIAutomaticSpeechRecognitionTask(),
|
||||
"text-to-image": FalAITextToImageTask(),
|
||||
"text-to-speech": FalAITextToSpeechTask(),
|
||||
"text-to-video": FalAITextToVideoTask(),
|
||||
},
|
||||
"featherless-ai": {
|
||||
"conversational": FeatherlessConversationalTask(),
|
||||
"text-generation": FeatherlessTextGenerationTask(),
|
||||
},
|
||||
"fireworks-ai": {
|
||||
"conversational": FireworksAIConversationalTask(),
|
||||
},
|
||||
"groq": {
|
||||
"conversational": GroqConversationalTask(),
|
||||
},
|
||||
"hf-inference": {
|
||||
"text-to-image": HFInferenceTask("text-to-image"),
|
||||
"conversational": HFInferenceConversational(),
|
||||
"text-generation": HFInferenceTask("text-generation"),
|
||||
"text-classification": HFInferenceTask("text-classification"),
|
||||
"question-answering": HFInferenceTask("question-answering"),
|
||||
"audio-classification": HFInferenceBinaryInputTask("audio-classification"),
|
||||
"automatic-speech-recognition": HFInferenceBinaryInputTask("automatic-speech-recognition"),
|
||||
"fill-mask": HFInferenceTask("fill-mask"),
|
||||
"feature-extraction": HFInferenceFeatureExtractionTask(),
|
||||
"image-classification": HFInferenceBinaryInputTask("image-classification"),
|
||||
"image-segmentation": HFInferenceBinaryInputTask("image-segmentation"),
|
||||
"document-question-answering": HFInferenceTask("document-question-answering"),
|
||||
"image-to-text": HFInferenceBinaryInputTask("image-to-text"),
|
||||
"object-detection": HFInferenceBinaryInputTask("object-detection"),
|
||||
"audio-to-audio": HFInferenceBinaryInputTask("audio-to-audio"),
|
||||
"zero-shot-image-classification": HFInferenceBinaryInputTask("zero-shot-image-classification"),
|
||||
"zero-shot-classification": HFInferenceTask("zero-shot-classification"),
|
||||
"image-to-image": HFInferenceBinaryInputTask("image-to-image"),
|
||||
"sentence-similarity": HFInferenceTask("sentence-similarity"),
|
||||
"table-question-answering": HFInferenceTask("table-question-answering"),
|
||||
"tabular-classification": HFInferenceTask("tabular-classification"),
|
||||
"text-to-speech": HFInferenceTask("text-to-speech"),
|
||||
"token-classification": HFInferenceTask("token-classification"),
|
||||
"translation": HFInferenceTask("translation"),
|
||||
"summarization": HFInferenceTask("summarization"),
|
||||
"visual-question-answering": HFInferenceBinaryInputTask("visual-question-answering"),
|
||||
},
|
||||
"hyperbolic": {
|
||||
"text-to-image": HyperbolicTextToImageTask(),
|
||||
"conversational": HyperbolicTextGenerationTask("conversational"),
|
||||
"text-generation": HyperbolicTextGenerationTask("text-generation"),
|
||||
},
|
||||
"nebius": {
|
||||
"text-to-image": NebiusTextToImageTask(),
|
||||
"conversational": NebiusConversationalTask(),
|
||||
"text-generation": NebiusTextGenerationTask(),
|
||||
"feature-extraction": NebiusFeatureExtractionTask(),
|
||||
},
|
||||
"novita": {
|
||||
"text-generation": NovitaTextGenerationTask(),
|
||||
"conversational": NovitaConversationalTask(),
|
||||
"text-to-video": NovitaTextToVideoTask(),
|
||||
},
|
||||
"nscale": {
|
||||
"conversational": NscaleConversationalTask(),
|
||||
"text-to-image": NscaleTextToImageTask(),
|
||||
},
|
||||
"openai": {
|
||||
"conversational": OpenAIConversationalTask(),
|
||||
},
|
||||
"replicate": {
|
||||
"text-to-image": ReplicateTextToImageTask(),
|
||||
"text-to-speech": ReplicateTextToSpeechTask(),
|
||||
"text-to-video": ReplicateTask("text-to-video"),
|
||||
},
|
||||
"sambanova": {
|
||||
"conversational": SambanovaConversationalTask(),
|
||||
"feature-extraction": SambanovaFeatureExtractionTask(),
|
||||
},
|
||||
"together": {
|
||||
"text-to-image": TogetherTextToImageTask(),
|
||||
"conversational": TogetherConversationalTask(),
|
||||
"text-generation": TogetherTextGenerationTask(),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def get_provider_helper(
|
||||
provider: Optional[PROVIDER_OR_POLICY_T], task: str, model: Optional[str]
|
||||
) -> TaskProviderHelper:
|
||||
"""Get provider helper instance by name and task.
|
||||
|
||||
Args:
|
||||
provider (`str`, *optional*): name of the provider, or "auto" to automatically select the provider for the model.
|
||||
task (`str`): Name of the task
|
||||
model (`str`, *optional*): Name of the model
|
||||
Returns:
|
||||
TaskProviderHelper: Helper instance for the specified provider and task
|
||||
|
||||
Raises:
|
||||
ValueError: If provider or task is not supported
|
||||
"""
|
||||
|
||||
if (model is None and provider in (None, "auto")) or (
|
||||
model is not None and model.startswith(("http://", "https://"))
|
||||
):
|
||||
provider = "hf-inference"
|
||||
|
||||
if provider is None:
|
||||
logger.info(
|
||||
"Defaulting to 'auto' which will select the first provider available for the model, sorted by the user's order in https://hf.co/settings/inference-providers."
|
||||
)
|
||||
provider = "auto"
|
||||
|
||||
if provider == "auto":
|
||||
if model is None:
|
||||
raise ValueError("Specifying a model is required when provider is 'auto'")
|
||||
provider_mapping = _fetch_inference_provider_mapping(model)
|
||||
provider = next(iter(provider_mapping)).provider
|
||||
|
||||
provider_tasks = PROVIDERS.get(provider) # type: ignore
|
||||
if provider_tasks is None:
|
||||
raise ValueError(
|
||||
f"Provider '{provider}' not supported. Available values: 'auto' or any provider from {list(PROVIDERS.keys())}."
|
||||
"Passing 'auto' (default value) will automatically select the first provider available for the model, sorted "
|
||||
"by the user's order in https://hf.co/settings/inference-providers."
|
||||
)
|
||||
|
||||
if task not in provider_tasks:
|
||||
raise ValueError(
|
||||
f"Task '{task}' not supported for provider '{provider}'. Available tasks: {list(provider_tasks.keys())}"
|
||||
)
|
||||
return provider_tasks[task]
|
||||
+299
@@ -0,0 +1,299 @@
|
||||
from functools import lru_cache
|
||||
from typing import Any, Dict, List, Optional, Union, overload
|
||||
|
||||
from huggingface_hub import constants
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters
|
||||
from huggingface_hub.inference._generated.types.chat_completion import ChatCompletionInputMessage
|
||||
from huggingface_hub.utils import build_hf_headers, get_token, logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
# Dev purposes only.
|
||||
# If you want to try to run inference for a new model locally before it's registered on huggingface.co
|
||||
# for a given Inference Provider, you can add it to the following dictionary.
|
||||
HARDCODED_MODEL_INFERENCE_MAPPING: Dict[str, Dict[str, InferenceProviderMapping]] = {
|
||||
# "HF model ID" => InferenceProviderMapping object initialized with "Model ID on Inference Provider's side"
|
||||
#
|
||||
# Example:
|
||||
# "Qwen/Qwen2.5-Coder-32B-Instruct": InferenceProviderMapping(hf_model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
|
||||
# provider_id="Qwen2.5-Coder-32B-Instruct",
|
||||
# task="conversational",
|
||||
# status="live")
|
||||
"cerebras": {},
|
||||
"cohere": {},
|
||||
"fal-ai": {},
|
||||
"fireworks-ai": {},
|
||||
"groq": {},
|
||||
"hf-inference": {},
|
||||
"hyperbolic": {},
|
||||
"nebius": {},
|
||||
"nscale": {},
|
||||
"replicate": {},
|
||||
"sambanova": {},
|
||||
"together": {},
|
||||
}
|
||||
|
||||
|
||||
@overload
|
||||
def filter_none(obj: Dict[str, Any]) -> Dict[str, Any]: ...
|
||||
@overload
|
||||
def filter_none(obj: List[Any]) -> List[Any]: ...
|
||||
|
||||
|
||||
def filter_none(obj: Union[Dict[str, Any], List[Any]]) -> Union[Dict[str, Any], List[Any]]:
|
||||
if isinstance(obj, dict):
|
||||
cleaned: Dict[str, Any] = {}
|
||||
for k, v in obj.items():
|
||||
if v is None:
|
||||
continue
|
||||
if isinstance(v, (dict, list)):
|
||||
v = filter_none(v)
|
||||
# remove empty nested dicts
|
||||
if isinstance(v, dict) and not v:
|
||||
continue
|
||||
cleaned[k] = v
|
||||
return cleaned
|
||||
|
||||
if isinstance(obj, list):
|
||||
return [filter_none(v) if isinstance(v, (dict, list)) else v for v in obj]
|
||||
|
||||
raise ValueError(f"Expected dict or list, got {type(obj)}")
|
||||
|
||||
|
||||
class TaskProviderHelper:
|
||||
"""Base class for task-specific provider helpers."""
|
||||
|
||||
def __init__(self, provider: str, base_url: str, task: str) -> None:
|
||||
self.provider = provider
|
||||
self.task = task
|
||||
self.base_url = base_url
|
||||
|
||||
def prepare_request(
|
||||
self,
|
||||
*,
|
||||
inputs: Any,
|
||||
parameters: Dict[str, Any],
|
||||
headers: Dict,
|
||||
model: Optional[str],
|
||||
api_key: Optional[str],
|
||||
extra_payload: Optional[Dict[str, Any]] = None,
|
||||
) -> RequestParameters:
|
||||
"""
|
||||
Prepare the request to be sent to the provider.
|
||||
|
||||
Each step (api_key, model, headers, url, payload) can be customized in subclasses.
|
||||
"""
|
||||
# api_key from user, or local token, or raise error
|
||||
api_key = self._prepare_api_key(api_key)
|
||||
|
||||
# mapped model from HF model ID
|
||||
provider_mapping_info = self._prepare_mapping_info(model)
|
||||
|
||||
# default HF headers + user headers (to customize in subclasses)
|
||||
headers = self._prepare_headers(headers, api_key)
|
||||
|
||||
# routed URL if HF token, or direct URL (to customize in '_prepare_route' in subclasses)
|
||||
url = self._prepare_url(api_key, provider_mapping_info.provider_id)
|
||||
|
||||
# prepare payload (to customize in subclasses)
|
||||
payload = self._prepare_payload_as_dict(inputs, parameters, provider_mapping_info=provider_mapping_info)
|
||||
if payload is not None:
|
||||
payload = recursive_merge(payload, extra_payload or {})
|
||||
|
||||
# body data (to customize in subclasses)
|
||||
data = self._prepare_payload_as_bytes(inputs, parameters, provider_mapping_info, extra_payload)
|
||||
|
||||
# check if both payload and data are set and return
|
||||
if payload is not None and data is not None:
|
||||
raise ValueError("Both payload and data cannot be set in the same request.")
|
||||
if payload is None and data is None:
|
||||
raise ValueError("Either payload or data must be set in the request.")
|
||||
return RequestParameters(
|
||||
url=url, task=self.task, model=provider_mapping_info.provider_id, json=payload, data=data, headers=headers
|
||||
)
|
||||
|
||||
def get_response(
|
||||
self,
|
||||
response: Union[bytes, Dict],
|
||||
request_params: Optional[RequestParameters] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Return the response in the expected format.
|
||||
|
||||
Override this method in subclasses for customized response handling."""
|
||||
return response
|
||||
|
||||
def _prepare_api_key(self, api_key: Optional[str]) -> str:
|
||||
"""Return the API key to use for the request.
|
||||
|
||||
Usually not overwritten in subclasses."""
|
||||
if api_key is None:
|
||||
api_key = get_token()
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
f"You must provide an api_key to work with {self.provider} API or log in with `huggingface-cli login`."
|
||||
)
|
||||
return api_key
|
||||
|
||||
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
|
||||
"""Return the mapped model ID to use for the request.
|
||||
|
||||
Usually not overwritten in subclasses."""
|
||||
if model is None:
|
||||
raise ValueError(f"Please provide an HF model ID supported by {self.provider}.")
|
||||
|
||||
# hardcoded mapping for local testing
|
||||
if HARDCODED_MODEL_INFERENCE_MAPPING.get(self.provider, {}).get(model):
|
||||
return HARDCODED_MODEL_INFERENCE_MAPPING[self.provider][model]
|
||||
|
||||
provider_mapping = None
|
||||
for mapping in _fetch_inference_provider_mapping(model):
|
||||
if mapping.provider == self.provider:
|
||||
provider_mapping = mapping
|
||||
break
|
||||
|
||||
if provider_mapping is None:
|
||||
raise ValueError(f"Model {model} is not supported by provider {self.provider}.")
|
||||
|
||||
if provider_mapping.task != self.task:
|
||||
raise ValueError(
|
||||
f"Model {model} is not supported for task {self.task} and provider {self.provider}. "
|
||||
f"Supported task: {provider_mapping.task}."
|
||||
)
|
||||
|
||||
if provider_mapping.status == "staging":
|
||||
logger.warning(
|
||||
f"Model {model} is in staging mode for provider {self.provider}. Meant for test purposes only."
|
||||
)
|
||||
if provider_mapping.status == "error":
|
||||
logger.warning(
|
||||
f"Our latest automated health check on model '{model}' for provider '{self.provider}' did not complete successfully. "
|
||||
"Inference call might fail."
|
||||
)
|
||||
return provider_mapping
|
||||
|
||||
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
|
||||
"""Return the headers to use for the request.
|
||||
|
||||
Override this method in subclasses for customized headers.
|
||||
"""
|
||||
return {**build_hf_headers(token=api_key), **headers}
|
||||
|
||||
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
|
||||
"""Return the URL to use for the request.
|
||||
|
||||
Usually not overwritten in subclasses."""
|
||||
base_url = self._prepare_base_url(api_key)
|
||||
route = self._prepare_route(mapped_model, api_key)
|
||||
return f"{base_url.rstrip('/')}/{route.lstrip('/')}"
|
||||
|
||||
def _prepare_base_url(self, api_key: str) -> str:
|
||||
"""Return the base URL to use for the request.
|
||||
|
||||
Usually not overwritten in subclasses."""
|
||||
# Route to the proxy if the api_key is a HF TOKEN
|
||||
if api_key.startswith("hf_"):
|
||||
logger.info(f"Calling '{self.provider}' provider through Hugging Face router.")
|
||||
return constants.INFERENCE_PROXY_TEMPLATE.format(provider=self.provider)
|
||||
else:
|
||||
logger.info(f"Calling '{self.provider}' provider directly.")
|
||||
return self.base_url
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
"""Return the route to use for the request.
|
||||
|
||||
Override this method in subclasses for customized routes.
|
||||
"""
|
||||
return ""
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
"""Return the payload to use for the request, as a dict.
|
||||
|
||||
Override this method in subclasses for customized payloads.
|
||||
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
|
||||
"""
|
||||
return None
|
||||
|
||||
def _prepare_payload_as_bytes(
|
||||
self,
|
||||
inputs: Any,
|
||||
parameters: Dict,
|
||||
provider_mapping_info: InferenceProviderMapping,
|
||||
extra_payload: Optional[Dict],
|
||||
) -> Optional[bytes]:
|
||||
"""Return the body to use for the request, as bytes.
|
||||
|
||||
Override this method in subclasses for customized body data.
|
||||
Only one of `_prepare_payload_as_dict` and `_prepare_payload_as_bytes` should return a value.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
class BaseConversationalTask(TaskProviderHelper):
|
||||
"""
|
||||
Base class for conversational (chat completion) tasks.
|
||||
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/chat
|
||||
"""
|
||||
|
||||
def __init__(self, provider: str, base_url: str):
|
||||
super().__init__(provider=provider, base_url=base_url, task="conversational")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/chat/completions"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self,
|
||||
inputs: List[Union[Dict, ChatCompletionInputMessage]],
|
||||
parameters: Dict,
|
||||
provider_mapping_info: InferenceProviderMapping,
|
||||
) -> Optional[Dict]:
|
||||
return filter_none({"messages": inputs, **parameters, "model": provider_mapping_info.provider_id})
|
||||
|
||||
|
||||
class BaseTextGenerationTask(TaskProviderHelper):
|
||||
"""
|
||||
Base class for text-generation (completion) tasks.
|
||||
The schema follows the OpenAI API format defined here: https://platform.openai.com/docs/api-reference/completions
|
||||
"""
|
||||
|
||||
def __init__(self, provider: str, base_url: str):
|
||||
super().__init__(provider=provider, base_url=base_url, task="text-generation")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/completions"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return {"prompt": inputs, **filter_none(parameters), "model": provider_mapping_info.provider_id}
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _fetch_inference_provider_mapping(model: str) -> List["InferenceProviderMapping"]:
|
||||
"""
|
||||
Fetch provider mappings for a model from the Hub.
|
||||
"""
|
||||
from huggingface_hub.hf_api import HfApi
|
||||
|
||||
info = HfApi().model_info(model, expand=["inferenceProviderMapping"])
|
||||
provider_mapping = info.inference_provider_mapping
|
||||
if provider_mapping is None:
|
||||
raise ValueError(f"No provider mapping found for model {model}")
|
||||
return provider_mapping
|
||||
|
||||
|
||||
def recursive_merge(dict1: Dict, dict2: Dict) -> Dict:
|
||||
return {
|
||||
**dict1,
|
||||
**{
|
||||
key: recursive_merge(dict1[key], value)
|
||||
if (key in dict1 and isinstance(dict1[key], dict) and isinstance(value, dict))
|
||||
else value
|
||||
for key, value in dict2.items()
|
||||
},
|
||||
}
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
import time
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
|
||||
from huggingface_hub.utils import logging
|
||||
from huggingface_hub.utils._http import get_session
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
MAX_POLLING_ATTEMPTS = 6
|
||||
POLLING_INTERVAL = 1.0
|
||||
|
||||
|
||||
class BlackForestLabsTextToImageTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(provider="black-forest-labs", base_url="https://api.us1.bfl.ai", task="text-to-image")
|
||||
|
||||
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
|
||||
headers = super()._prepare_headers(headers, api_key)
|
||||
if not api_key.startswith("hf_"):
|
||||
_ = headers.pop("authorization")
|
||||
headers["X-Key"] = api_key
|
||||
return headers
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return f"/v1/{mapped_model}"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
parameters = filter_none(parameters)
|
||||
if "num_inference_steps" in parameters:
|
||||
parameters["steps"] = parameters.pop("num_inference_steps")
|
||||
if "guidance_scale" in parameters:
|
||||
parameters["guidance"] = parameters.pop("guidance_scale")
|
||||
|
||||
return {"prompt": inputs, **parameters}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
"""
|
||||
Polling mechanism for Black Forest Labs since the API is asynchronous.
|
||||
"""
|
||||
url = _as_dict(response).get("polling_url")
|
||||
session = get_session()
|
||||
for _ in range(MAX_POLLING_ATTEMPTS):
|
||||
time.sleep(POLLING_INTERVAL)
|
||||
|
||||
response = session.get(url, headers={"Content-Type": "application/json"}) # type: ignore
|
||||
response.raise_for_status() # type: ignore
|
||||
response_json: Dict = response.json() # type: ignore
|
||||
status = response_json.get("status")
|
||||
logger.info(
|
||||
f"Polling generation result from {url}. Current status: {status}. "
|
||||
f"Will retry after {POLLING_INTERVAL} seconds if not ready."
|
||||
)
|
||||
|
||||
if (
|
||||
status == "Ready"
|
||||
and isinstance(response_json.get("result"), dict)
|
||||
and (sample_url := response_json["result"].get("sample"))
|
||||
):
|
||||
image_resp = session.get(sample_url)
|
||||
image_resp.raise_for_status()
|
||||
return image_resp.content
|
||||
|
||||
raise TimeoutError(f"Failed to get the image URL after {MAX_POLLING_ATTEMPTS} attempts.")
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
class CerebrasConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="cerebras", base_url="https://api.cerebras.ai")
|
||||
+32
@@ -0,0 +1,32 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
|
||||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
_PROVIDER = "cohere"
|
||||
_BASE_URL = "https://api.cohere.com"
|
||||
|
||||
|
||||
class CohereConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/compatibility/v1/chat/completions"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
|
||||
response_format = parameters.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
json_schema_details = response_format.get("json_schema")
|
||||
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
|
||||
payload["response_format"] = { # type: ignore [index]
|
||||
"type": "json_object",
|
||||
"schema": json_schema_details["schema"],
|
||||
}
|
||||
|
||||
return payload
|
||||
+172
@@ -0,0 +1,172 @@
|
||||
import base64
|
||||
import time
|
||||
from abc import ABC
|
||||
from typing import Any, Dict, Optional, Union
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from huggingface_hub import constants
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
|
||||
from huggingface_hub.utils import get_session, hf_raise_for_status
|
||||
from huggingface_hub.utils.logging import get_logger
|
||||
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Arbitrary polling interval
|
||||
_POLLING_INTERVAL = 0.5
|
||||
|
||||
|
||||
class FalAITask(TaskProviderHelper, ABC):
|
||||
def __init__(self, task: str):
|
||||
super().__init__(provider="fal-ai", base_url="https://fal.run", task=task)
|
||||
|
||||
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
|
||||
headers = super()._prepare_headers(headers, api_key)
|
||||
if not api_key.startswith("hf_"):
|
||||
headers["authorization"] = f"Key {api_key}"
|
||||
return headers
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return f"/{mapped_model}"
|
||||
|
||||
|
||||
class FalAIAutomaticSpeechRecognitionTask(FalAITask):
|
||||
def __init__(self):
|
||||
super().__init__("automatic-speech-recognition")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
if isinstance(inputs, str) and inputs.startswith(("http://", "https://")):
|
||||
# If input is a URL, pass it directly
|
||||
audio_url = inputs
|
||||
else:
|
||||
# If input is a file path, read it first
|
||||
if isinstance(inputs, str):
|
||||
with open(inputs, "rb") as f:
|
||||
inputs = f.read()
|
||||
|
||||
audio_b64 = base64.b64encode(inputs).decode()
|
||||
content_type = "audio/mpeg"
|
||||
audio_url = f"data:{content_type};base64,{audio_b64}"
|
||||
|
||||
return {"audio_url": audio_url, **filter_none(parameters)}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
text = _as_dict(response)["text"]
|
||||
if not isinstance(text, str):
|
||||
raise ValueError(f"Unexpected output format from FalAI API. Expected string, got {type(text)}.")
|
||||
return text
|
||||
|
||||
|
||||
class FalAITextToImageTask(FalAITask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-image")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload: Dict[str, Any] = {
|
||||
"prompt": inputs,
|
||||
**filter_none(parameters),
|
||||
}
|
||||
if "width" in payload and "height" in payload:
|
||||
payload["image_size"] = {
|
||||
"width": payload.pop("width"),
|
||||
"height": payload.pop("height"),
|
||||
}
|
||||
if provider_mapping_info.adapter_weights_path is not None:
|
||||
lora_path = constants.HUGGINGFACE_CO_URL_TEMPLATE.format(
|
||||
repo_id=provider_mapping_info.hf_model_id,
|
||||
revision="main",
|
||||
filename=provider_mapping_info.adapter_weights_path,
|
||||
)
|
||||
payload["loras"] = [{"path": lora_path, "scale": 1}]
|
||||
if provider_mapping_info.provider_id == "fal-ai/lora":
|
||||
# little hack: fal requires the base model for stable-diffusion-based loras but not for flux-based
|
||||
# See payloads in https://fal.ai/models/fal-ai/lora/api vs https://fal.ai/models/fal-ai/flux-lora/api
|
||||
payload["model_name"] = "stabilityai/stable-diffusion-xl-base-1.0"
|
||||
|
||||
return payload
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
url = _as_dict(response)["images"][0]["url"]
|
||||
return get_session().get(url).content
|
||||
|
||||
|
||||
class FalAITextToSpeechTask(FalAITask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-speech")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return {"text": inputs, **filter_none(parameters)}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
url = _as_dict(response)["audio"]["url"]
|
||||
return get_session().get(url).content
|
||||
|
||||
|
||||
class FalAITextToVideoTask(FalAITask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-video")
|
||||
|
||||
def _prepare_base_url(self, api_key: str) -> str:
|
||||
if api_key.startswith("hf_"):
|
||||
return super()._prepare_base_url(api_key)
|
||||
else:
|
||||
logger.info(f"Calling '{self.provider}' provider directly.")
|
||||
return "https://queue.fal.run"
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
if api_key.startswith("hf_"):
|
||||
# Use the queue subdomain for HF routing
|
||||
return f"/{mapped_model}?_subdomain=queue"
|
||||
return f"/{mapped_model}"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return {"prompt": inputs, **filter_none(parameters)}
|
||||
|
||||
def get_response(
|
||||
self,
|
||||
response: Union[bytes, Dict],
|
||||
request_params: Optional[RequestParameters] = None,
|
||||
) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
|
||||
request_id = response_dict.get("request_id")
|
||||
if not request_id:
|
||||
raise ValueError("No request ID found in the response")
|
||||
if request_params is None:
|
||||
raise ValueError(
|
||||
"A `RequestParameters` object should be provided to get text-to-video responses with Fal AI."
|
||||
)
|
||||
|
||||
# extract the base url and query params
|
||||
parsed_url = urlparse(request_params.url)
|
||||
# a bit hacky way to concatenate the provider name without parsing `parsed_url.path`
|
||||
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}{'/fal-ai' if parsed_url.netloc == 'router.huggingface.co' else ''}"
|
||||
query_param = f"?{parsed_url.query}" if parsed_url.query else ""
|
||||
|
||||
# extracting the provider model id for status and result urls
|
||||
# from the response as it might be different from the mapped model in `request_params.url`
|
||||
model_id = urlparse(response_dict.get("response_url")).path
|
||||
status_url = f"{base_url}{str(model_id)}/status{query_param}"
|
||||
result_url = f"{base_url}{str(model_id)}{query_param}"
|
||||
|
||||
status = response_dict.get("status")
|
||||
logger.info("Generating the video.. this can take several minutes.")
|
||||
while status != "COMPLETED":
|
||||
time.sleep(_POLLING_INTERVAL)
|
||||
status_response = get_session().get(status_url, headers=request_params.headers)
|
||||
hf_raise_for_status(status_response)
|
||||
status = status_response.json().get("status")
|
||||
|
||||
response = get_session().get(result_url, headers=request_params.headers).json()
|
||||
url = _as_dict(response)["video"]["url"]
|
||||
return get_session().get(url).content
|
||||
+38
@@ -0,0 +1,38 @@
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
|
||||
from ._common import BaseConversationalTask, BaseTextGenerationTask, filter_none
|
||||
|
||||
|
||||
_PROVIDER = "featherless-ai"
|
||||
_BASE_URL = "https://api.featherless.ai"
|
||||
|
||||
|
||||
class FeatherlessTextGenerationTask(BaseTextGenerationTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
params = filter_none(parameters.copy())
|
||||
params["max_tokens"] = params.pop("max_new_tokens", None)
|
||||
|
||||
return {"prompt": inputs, **params, "model": provider_mapping_info.provider_id}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
output = _as_dict(response)["choices"][0]
|
||||
return {
|
||||
"generated_text": output["text"],
|
||||
"details": {
|
||||
"finish_reason": output.get("finish_reason"),
|
||||
"seed": output.get("seed"),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class FeatherlessConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
|
||||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
class FireworksAIConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="fireworks-ai", base_url="https://api.fireworks.ai")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/inference/v1/chat/completions"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
|
||||
response_format = parameters.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
json_schema_details = response_format.get("json_schema")
|
||||
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
|
||||
payload["response_format"] = { # type: ignore [index]
|
||||
"type": "json_object",
|
||||
"schema": json_schema_details["schema"],
|
||||
}
|
||||
return payload
|
||||
@@ -0,0 +1,9 @@
|
||||
from ._common import BaseConversationalTask
|
||||
|
||||
|
||||
class GroqConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="groq", base_url="https://api.groq.com")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/openai/v1/chat/completions"
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
import json
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub import constants
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _b64_encode, _bytes_to_dict, _open_as_binary
|
||||
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
|
||||
from huggingface_hub.utils import build_hf_headers, get_session, get_token, hf_raise_for_status
|
||||
|
||||
|
||||
class HFInferenceTask(TaskProviderHelper):
|
||||
"""Base class for HF Inference API tasks."""
|
||||
|
||||
def __init__(self, task: str):
|
||||
super().__init__(
|
||||
provider="hf-inference",
|
||||
base_url=constants.INFERENCE_PROXY_TEMPLATE.format(provider="hf-inference"),
|
||||
task=task,
|
||||
)
|
||||
|
||||
def _prepare_api_key(self, api_key: Optional[str]) -> str:
|
||||
# special case: for HF Inference we allow not providing an API key
|
||||
return api_key or get_token() # type: ignore[return-value]
|
||||
|
||||
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
|
||||
if model is not None and model.startswith(("http://", "https://")):
|
||||
return InferenceProviderMapping(
|
||||
provider="hf-inference", providerId=model, hf_model_id=model, task=self.task, status="live"
|
||||
)
|
||||
model_id = model if model is not None else _fetch_recommended_models().get(self.task)
|
||||
if model_id is None:
|
||||
raise ValueError(
|
||||
f"Task {self.task} has no recommended model for HF Inference. Please specify a model"
|
||||
" explicitly. Visit https://huggingface.co/tasks for more info."
|
||||
)
|
||||
_check_supported_task(model_id, self.task)
|
||||
return InferenceProviderMapping(
|
||||
provider="hf-inference", providerId=model_id, hf_model_id=model_id, task=self.task, status="live"
|
||||
)
|
||||
|
||||
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
|
||||
# hf-inference provider can handle URLs (e.g. Inference Endpoints or TGI deployment)
|
||||
if mapped_model.startswith(("http://", "https://")):
|
||||
return mapped_model
|
||||
return (
|
||||
# Feature-extraction and sentence-similarity are the only cases where we handle models with several tasks.
|
||||
f"{self.base_url}/models/{mapped_model}/pipeline/{self.task}"
|
||||
if self.task in ("feature-extraction", "sentence-similarity")
|
||||
# Otherwise, we use the default endpoint
|
||||
else f"{self.base_url}/models/{mapped_model}"
|
||||
)
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
if isinstance(inputs, bytes):
|
||||
raise ValueError(f"Unexpected binary input for task {self.task}.")
|
||||
if isinstance(inputs, Path):
|
||||
raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
|
||||
return {"inputs": inputs, "parameters": filter_none(parameters)}
|
||||
|
||||
|
||||
class HFInferenceBinaryInputTask(HFInferenceTask):
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return None
|
||||
|
||||
def _prepare_payload_as_bytes(
|
||||
self,
|
||||
inputs: Any,
|
||||
parameters: Dict,
|
||||
provider_mapping_info: InferenceProviderMapping,
|
||||
extra_payload: Optional[Dict],
|
||||
) -> Optional[bytes]:
|
||||
parameters = filter_none(parameters)
|
||||
extra_payload = extra_payload or {}
|
||||
has_parameters = len(parameters) > 0 or len(extra_payload) > 0
|
||||
|
||||
# Raise if not a binary object or a local path or a URL.
|
||||
if not isinstance(inputs, (bytes, Path)) and not isinstance(inputs, str):
|
||||
raise ValueError(f"Expected binary inputs or a local path or a URL. Got {inputs}")
|
||||
|
||||
# Send inputs as raw content when no parameters are provided
|
||||
if not has_parameters:
|
||||
with _open_as_binary(inputs) as data:
|
||||
data_as_bytes = data if isinstance(data, bytes) else data.read()
|
||||
return data_as_bytes
|
||||
|
||||
# Otherwise encode as b64
|
||||
return json.dumps({"inputs": _b64_encode(inputs), "parameters": parameters, **extra_payload}).encode("utf-8")
|
||||
|
||||
|
||||
class HFInferenceConversational(HFInferenceTask):
|
||||
def __init__(self):
|
||||
super().__init__("conversational")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload = filter_none(parameters)
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
payload_model = parameters.get("model") or mapped_model
|
||||
|
||||
if payload_model is None or payload_model.startswith(("http://", "https://")):
|
||||
payload_model = "dummy"
|
||||
|
||||
response_format = parameters.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
payload["response_format"] = {
|
||||
"type": "json_object",
|
||||
"value": response_format["json_schema"]["schema"],
|
||||
}
|
||||
return {**payload, "model": payload_model, "messages": inputs}
|
||||
|
||||
def _prepare_url(self, api_key: str, mapped_model: str) -> str:
|
||||
base_url = (
|
||||
mapped_model
|
||||
if mapped_model.startswith(("http://", "https://"))
|
||||
else f"{constants.INFERENCE_PROXY_TEMPLATE.format(provider='hf-inference')}/models/{mapped_model}"
|
||||
)
|
||||
return _build_chat_completion_url(base_url)
|
||||
|
||||
|
||||
def _build_chat_completion_url(model_url: str) -> str:
|
||||
# Strip trailing /
|
||||
model_url = model_url.rstrip("/")
|
||||
|
||||
# Append /chat/completions if not already present
|
||||
if model_url.endswith("/v1"):
|
||||
model_url += "/chat/completions"
|
||||
|
||||
# Append /v1/chat/completions if not already present
|
||||
if not model_url.endswith("/chat/completions"):
|
||||
model_url += "/v1/chat/completions"
|
||||
|
||||
return model_url
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def _fetch_recommended_models() -> Dict[str, Optional[str]]:
|
||||
response = get_session().get(f"{constants.ENDPOINT}/api/tasks", headers=build_hf_headers())
|
||||
hf_raise_for_status(response)
|
||||
return {task: next(iter(details["widgetModels"]), None) for task, details in response.json().items()}
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _check_supported_task(model: str, task: str) -> None:
|
||||
from huggingface_hub.hf_api import HfApi
|
||||
|
||||
model_info = HfApi().model_info(model)
|
||||
pipeline_tag = model_info.pipeline_tag
|
||||
tags = model_info.tags or []
|
||||
is_conversational = "conversational" in tags
|
||||
if task in ("text-generation", "conversational"):
|
||||
if pipeline_tag == "text-generation":
|
||||
# text-generation + conversational tag -> both tasks allowed
|
||||
if is_conversational:
|
||||
return
|
||||
# text-generation without conversational tag -> only text-generation allowed
|
||||
if task == "text-generation":
|
||||
return
|
||||
raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
|
||||
|
||||
if pipeline_tag == "text2text-generation":
|
||||
if task == "text-generation":
|
||||
return
|
||||
raise ValueError(f"Model '{model}' doesn't support task '{task}'.")
|
||||
|
||||
if pipeline_tag == "image-text-to-text":
|
||||
if is_conversational and task == "conversational":
|
||||
return # Only conversational allowed if tagged as conversational
|
||||
raise ValueError("Non-conversational image-text-to-text task is not supported.")
|
||||
|
||||
if (
|
||||
task in ("feature-extraction", "sentence-similarity")
|
||||
and pipeline_tag in ("feature-extraction", "sentence-similarity")
|
||||
and task in tags
|
||||
):
|
||||
# feature-extraction and sentence-similarity are interchangeable for HF Inference
|
||||
return
|
||||
|
||||
# For all other tasks, just check pipeline tag
|
||||
if pipeline_tag != task:
|
||||
raise ValueError(
|
||||
f"Model '{model}' doesn't support task '{task}'. Supported tasks: '{pipeline_tag}', got: '{task}'"
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
class HFInferenceFeatureExtractionTask(HFInferenceTask):
|
||||
def __init__(self):
|
||||
super().__init__("feature-extraction")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
if isinstance(inputs, bytes):
|
||||
raise ValueError(f"Unexpected binary input for task {self.task}.")
|
||||
if isinstance(inputs, Path):
|
||||
raise ValueError(f"Unexpected path input for task {self.task} (got {inputs})")
|
||||
|
||||
# Parameters are sent at root-level for feature-extraction task
|
||||
# See specs: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/feature-extraction/spec/input.json
|
||||
return {"inputs": inputs, **filter_none(parameters)}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
if isinstance(response, bytes):
|
||||
return _bytes_to_dict(response)
|
||||
return response
|
||||
+47
@@ -0,0 +1,47 @@
|
||||
import base64
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
|
||||
|
||||
|
||||
class HyperbolicTextToImageTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(provider="hyperbolic", base_url="https://api.hyperbolic.xyz", task="text-to-image")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/images/generations"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
parameters = filter_none(parameters)
|
||||
if "num_inference_steps" in parameters:
|
||||
parameters["steps"] = parameters.pop("num_inference_steps")
|
||||
if "guidance_scale" in parameters:
|
||||
parameters["cfg_scale"] = parameters.pop("guidance_scale")
|
||||
# For Hyperbolic, the width and height are required parameters
|
||||
if "width" not in parameters:
|
||||
parameters["width"] = 512
|
||||
if "height" not in parameters:
|
||||
parameters["height"] = 512
|
||||
return {"prompt": inputs, "model_name": mapped_model, **parameters}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
return base64.b64decode(response_dict["images"][0]["image"])
|
||||
|
||||
|
||||
class HyperbolicTextGenerationTask(BaseConversationalTask):
|
||||
"""
|
||||
Special case for Hyperbolic, where text-generation task is handled as a conversational task.
|
||||
"""
|
||||
|
||||
def __init__(self, task: str):
|
||||
super().__init__(
|
||||
provider="hyperbolic",
|
||||
base_url="https://api.hyperbolic.xyz",
|
||||
)
|
||||
self.task = task
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
import base64
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import (
|
||||
BaseConversationalTask,
|
||||
BaseTextGenerationTask,
|
||||
TaskProviderHelper,
|
||||
filter_none,
|
||||
)
|
||||
|
||||
|
||||
class NebiusTextGenerationTask(BaseTextGenerationTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
output = _as_dict(response)["choices"][0]
|
||||
return {
|
||||
"generated_text": output["text"],
|
||||
"details": {
|
||||
"finish_reason": output.get("finish_reason"),
|
||||
"seed": output.get("seed"),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class NebiusConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="nebius", base_url="https://api.studio.nebius.ai")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
|
||||
response_format = parameters.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
json_schema_details = response_format.get("json_schema")
|
||||
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
|
||||
payload["guided_json"] = json_schema_details["schema"] # type: ignore [index]
|
||||
return payload
|
||||
|
||||
|
||||
class NebiusTextToImageTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(task="text-to-image", provider="nebius", base_url="https://api.studio.nebius.ai")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/images/generations"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
parameters = filter_none(parameters)
|
||||
if "guidance_scale" in parameters:
|
||||
parameters.pop("guidance_scale")
|
||||
if parameters.get("response_format") not in ("b64_json", "url"):
|
||||
parameters["response_format"] = "b64_json"
|
||||
|
||||
return {"prompt": inputs, **parameters, "model": mapped_model}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
return base64.b64decode(response_dict["data"][0]["b64_json"])
|
||||
|
||||
|
||||
class NebiusFeatureExtractionTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(task="feature-extraction", provider="nebius", base_url="https://api.studio.nebius.ai")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/embeddings"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return {"input": inputs, "model": provider_mapping_info.provider_id}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
embeddings = _as_dict(response)["data"]
|
||||
return [embedding["embedding"] for embedding in embeddings]
|
||||
+69
@@ -0,0 +1,69 @@
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import (
|
||||
BaseConversationalTask,
|
||||
BaseTextGenerationTask,
|
||||
TaskProviderHelper,
|
||||
filter_none,
|
||||
)
|
||||
from huggingface_hub.utils import get_session
|
||||
|
||||
|
||||
_PROVIDER = "novita"
|
||||
_BASE_URL = "https://api.novita.ai"
|
||||
|
||||
|
||||
class NovitaTextGenerationTask(BaseTextGenerationTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
# there is no v1/ route for novita
|
||||
return "/v3/openai/completions"
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
output = _as_dict(response)["choices"][0]
|
||||
return {
|
||||
"generated_text": output["text"],
|
||||
"details": {
|
||||
"finish_reason": output.get("finish_reason"),
|
||||
"seed": output.get("seed"),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class NovitaConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
# there is no v1/ route for novita
|
||||
return "/v3/openai/chat/completions"
|
||||
|
||||
|
||||
class NovitaTextToVideoTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task="text-to-video")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return f"/v3/hf/{mapped_model}"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
return {"prompt": inputs, **filter_none(parameters)}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
if not (
|
||||
isinstance(response_dict, dict)
|
||||
and "video" in response_dict
|
||||
and isinstance(response_dict["video"], dict)
|
||||
and "video_url" in response_dict["video"]
|
||||
):
|
||||
raise ValueError("Expected response format: { 'video': { 'video_url': string } }")
|
||||
|
||||
video_url = response_dict["video"]["video_url"]
|
||||
return get_session().get(video_url).content
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
import base64
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
|
||||
from ._common import BaseConversationalTask, TaskProviderHelper, filter_none
|
||||
|
||||
|
||||
class NscaleConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com")
|
||||
|
||||
|
||||
class NscaleTextToImageTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(provider="nscale", base_url="https://inference.api.nscale.com", task="text-to-image")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/images/generations"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
# Combine all parameters except inputs and parameters
|
||||
parameters = filter_none(parameters)
|
||||
if "width" in parameters and "height" in parameters:
|
||||
parameters["size"] = f"{parameters.pop('width')}x{parameters.pop('height')}"
|
||||
if "num_inference_steps" in parameters:
|
||||
parameters.pop("num_inference_steps")
|
||||
if "cfg_scale" in parameters:
|
||||
parameters.pop("cfg_scale")
|
||||
payload = {
|
||||
"response_format": "b64_json",
|
||||
"prompt": inputs,
|
||||
"model": mapped_model,
|
||||
**parameters,
|
||||
}
|
||||
return payload
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
return base64.b64decode(response_dict["data"][0]["b64_json"])
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
from typing import Optional
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._providers._common import BaseConversationalTask
|
||||
|
||||
|
||||
class OpenAIConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="openai", base_url="https://api.openai.com")
|
||||
|
||||
def _prepare_api_key(self, api_key: Optional[str]) -> str:
|
||||
if api_key is None:
|
||||
raise ValueError("You must provide an api_key to work with OpenAI API.")
|
||||
if api_key.startswith("hf_"):
|
||||
raise ValueError(
|
||||
"OpenAI provider is not available through Hugging Face routing, please use your own OpenAI API key."
|
||||
)
|
||||
return api_key
|
||||
|
||||
def _prepare_mapping_info(self, model: Optional[str]) -> InferenceProviderMapping:
|
||||
if model is None:
|
||||
raise ValueError("Please provide an OpenAI model ID, e.g. `gpt-4o` or `o1`.")
|
||||
return InferenceProviderMapping(
|
||||
provider="openai", providerId=model, task="conversational", status="live", hf_model_id=model
|
||||
)
|
||||
+72
@@ -0,0 +1,72 @@
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import TaskProviderHelper, filter_none
|
||||
from huggingface_hub.utils import get_session
|
||||
|
||||
|
||||
_PROVIDER = "replicate"
|
||||
_BASE_URL = "https://api.replicate.com"
|
||||
|
||||
|
||||
class ReplicateTask(TaskProviderHelper):
|
||||
def __init__(self, task: str):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
|
||||
|
||||
def _prepare_headers(self, headers: Dict, api_key: str) -> Dict:
|
||||
headers = super()._prepare_headers(headers, api_key)
|
||||
headers["Prefer"] = "wait"
|
||||
return headers
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
if ":" in mapped_model:
|
||||
return "/v1/predictions"
|
||||
return f"/v1/models/{mapped_model}/predictions"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
payload: Dict[str, Any] = {"input": {"prompt": inputs, **filter_none(parameters)}}
|
||||
if ":" in mapped_model:
|
||||
version = mapped_model.split(":", 1)[1]
|
||||
payload["version"] = version
|
||||
return payload
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
if response_dict.get("output") is None:
|
||||
raise TimeoutError(
|
||||
f"Inference request timed out after 60 seconds. No output generated for model {response_dict.get('model')}"
|
||||
"The model might be in cold state or starting up. Please try again later."
|
||||
)
|
||||
output_url = (
|
||||
response_dict["output"] if isinstance(response_dict["output"], str) else response_dict["output"][0]
|
||||
)
|
||||
return get_session().get(output_url).content
|
||||
|
||||
|
||||
class ReplicateTextToImageTask(ReplicateTask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-image")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
|
||||
if provider_mapping_info.adapter_weights_path is not None:
|
||||
payload["input"]["lora_weights"] = f"https://huggingface.co/{provider_mapping_info.hf_model_id}"
|
||||
return payload
|
||||
|
||||
|
||||
class ReplicateTextToSpeechTask(ReplicateTask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-speech")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload: Dict = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info) # type: ignore[assignment]
|
||||
payload["input"]["text"] = payload["input"].pop("prompt") # rename "prompt" to "text" for TTS
|
||||
return payload
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import BaseConversationalTask, TaskProviderHelper, filter_none
|
||||
|
||||
|
||||
class SambanovaConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
response_format_config = parameters.get("response_format")
|
||||
if isinstance(response_format_config, dict):
|
||||
if response_format_config.get("type") == "json_schema":
|
||||
json_schema_config = response_format_config.get("json_schema", {})
|
||||
strict = json_schema_config.get("strict")
|
||||
if isinstance(json_schema_config, dict) and (strict is True or strict is None):
|
||||
json_schema_config["strict"] = False
|
||||
|
||||
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
|
||||
return payload
|
||||
|
||||
|
||||
class SambanovaFeatureExtractionTask(TaskProviderHelper):
|
||||
def __init__(self):
|
||||
super().__init__(provider="sambanova", base_url="https://api.sambanova.ai", task="feature-extraction")
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
return "/v1/embeddings"
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
parameters = filter_none(parameters)
|
||||
return {"input": inputs, "model": provider_mapping_info.provider_id, **parameters}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
embeddings = _as_dict(response)["data"]
|
||||
return [embedding["embedding"] for embedding in embeddings]
|
||||
+88
@@ -0,0 +1,88 @@
|
||||
import base64
|
||||
from abc import ABC
|
||||
from typing import Any, Dict, Optional, Union
|
||||
|
||||
from huggingface_hub.hf_api import InferenceProviderMapping
|
||||
from huggingface_hub.inference._common import RequestParameters, _as_dict
|
||||
from huggingface_hub.inference._providers._common import (
|
||||
BaseConversationalTask,
|
||||
BaseTextGenerationTask,
|
||||
TaskProviderHelper,
|
||||
filter_none,
|
||||
)
|
||||
|
||||
|
||||
_PROVIDER = "together"
|
||||
_BASE_URL = "https://api.together.xyz"
|
||||
|
||||
|
||||
class TogetherTask(TaskProviderHelper, ABC):
|
||||
"""Base class for Together API tasks."""
|
||||
|
||||
def __init__(self, task: str):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL, task=task)
|
||||
|
||||
def _prepare_route(self, mapped_model: str, api_key: str) -> str:
|
||||
if self.task == "text-to-image":
|
||||
return "/v1/images/generations"
|
||||
elif self.task == "conversational":
|
||||
return "/v1/chat/completions"
|
||||
elif self.task == "text-generation":
|
||||
return "/v1/completions"
|
||||
raise ValueError(f"Unsupported task '{self.task}' for Together API.")
|
||||
|
||||
|
||||
class TogetherTextGenerationTask(BaseTextGenerationTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
output = _as_dict(response)["choices"][0]
|
||||
return {
|
||||
"generated_text": output["text"],
|
||||
"details": {
|
||||
"finish_reason": output.get("finish_reason"),
|
||||
"seed": output.get("seed"),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TogetherConversationalTask(BaseConversationalTask):
|
||||
def __init__(self):
|
||||
super().__init__(provider=_PROVIDER, base_url=_BASE_URL)
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
payload = super()._prepare_payload_as_dict(inputs, parameters, provider_mapping_info)
|
||||
response_format = parameters.get("response_format")
|
||||
if isinstance(response_format, dict) and response_format.get("type") == "json_schema":
|
||||
json_schema_details = response_format.get("json_schema")
|
||||
if isinstance(json_schema_details, dict) and "schema" in json_schema_details:
|
||||
payload["response_format"] = { # type: ignore [index]
|
||||
"type": "json_object",
|
||||
"schema": json_schema_details["schema"],
|
||||
}
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
class TogetherTextToImageTask(TogetherTask):
|
||||
def __init__(self):
|
||||
super().__init__("text-to-image")
|
||||
|
||||
def _prepare_payload_as_dict(
|
||||
self, inputs: Any, parameters: Dict, provider_mapping_info: InferenceProviderMapping
|
||||
) -> Optional[Dict]:
|
||||
mapped_model = provider_mapping_info.provider_id
|
||||
parameters = filter_none(parameters)
|
||||
if "num_inference_steps" in parameters:
|
||||
parameters["steps"] = parameters.pop("num_inference_steps")
|
||||
if "guidance_scale" in parameters:
|
||||
parameters["guidance"] = parameters.pop("guidance_scale")
|
||||
|
||||
return {"prompt": inputs, "response_format": "base64", **parameters, "model": mapped_model}
|
||||
|
||||
def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
|
||||
response_dict = _as_dict(response)
|
||||
return base64.b64decode(response_dict["data"][0]["b64_json"])
|
||||
Reference in New Issue
Block a user